[Feature Request] ProbabilisticActor in a Composite Action Space Environment

[sakakibara-yuuki]

Motivation

Using PPO tutorial and #1473 issue as a guide, I created an environment with a composite action space.

I also created a module using ProbabilisticActor based on the #1473.
However, the values sampled by distribution_class were not written to action entries.
(In the PPO tutorial and no composite action space version of #1473, which does not use a composite action space, the values are written to action.)

when action space is not composite.

class Module(nn.Module):
    def forward(self, x):
        return x[..., :3], x[..., 3:6]
module = TensorDictModule(Module(),
                          in_keys=["x"],
                          out_keys=["loc", "scale"])
actor = ProbabilisticActor(module,
                           in_keys=["loc", "scale"],
                           distribution_class=d.Normal,
                           # distribution_kwargs={"distribution_map": {"normal": d.Normal, "categ": d.Categorical}}
                          )
data = TensorDict({"x": torch.rand(10)}, [])
module(data)
print(actor(data))

then, I get action entries.

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        loc: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        scale: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        x: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

when action space is composite, in #1473.

class Module(nn.Module):
    def forward(self, x):
        return x[..., :3], x[..., 3:6], x[..., 6:]
module = TensorDictModule(Module(),
                          in_keys=["x"],
                          out_keys=[("params", "normal", "loc"), ("params", "normal", "scale"), ("params", "categ", "logits")])
actor = ProbabilisticActor(module,
                           in_keys=["params"],
                           distribution_class=CompositeDistribution,
                           distribution_kwargs={"distribution_map": {"normal": d.Normal, "categ": d.Categorical}}
                          )
data = TensorDict({"x": torch.rand(10)}, [])
module(data)
print(actor(data))

then, I don't get action entries !

TensorDict(
    fields={
        categ: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        normal: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        params: TensorDict(
            fields={
                categ: TensorDict(
                    fields={
                        logits: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([]),
                    device=None,
                    is_shared=False),
                normal: TensorDict(
                    fields={
                        loc: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
                        scale: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([]),
                    device=None,
                    is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False),
        x: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

Solution

I would like to obtain the following results (The output tensordict contains an action entry)

TensorDict(
    fields={
        action: TensorDict(
            fields={
                categ: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
                normal: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False),
        params: TensorDict(
            fields={
                 categ: TensorDict(
                      fields={
                          logits: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False)},
                      batch_size=torch.Size([]),
                      device=None,
                      is_shared=False),
                  normal: TensorDict(
                      fields={
                          loc: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
                          scale: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
                      batch_size=torch.Size([]),
                      device=None,
                      is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False),
        x: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

Could you provide guidance on how to correctly write the sampled values into the action key when using ProbabilisticActor in a composite action space?

Alternatives

A possible symptomatic remedy would be to create a TensorDictModule that moves the normal and categ keys under the action key after using the ProbabilisticActor.

Or, Including the action key in module output_keys and the distribution_map of ProbabilisticActor solves the problem, but causes inconvenience when using PPOLoss.

class Module(nn.Module):
    def forward(self, x):
        return x[..., :3], x[..., 3:6], x[..., 6:]
module = TensorDictModule(Module(),
                          in_keys=["x"],
                          out_keys=[("params", "action", "normal", "loc"), ("params", "action", "normal", "scale"), ("params", "action", "categ", "logits")])
actor = ProbabilisticActor(module,
                           out_keys=[("action", "normal"), ("action", "categ")],
                           in_keys=["params"],
                           distribution_class=CompositeDistribution,
                           distribution_kwargs={"distribution_map": {("action", "normal"): d.Normal, ("action", "categ"): d.Categorical}}
                          )
data = TensorDict({"x": torch.rand(10)}, [])
module(data)
print(actor(data))

This may be due to the fact that out_keys is not taken into account when out_tensors in the forward method of ProbabilisticTensorModule is an instance of TensorDictBase.

(tensordict/nn/probabilistic.py line 379)

    @dispatch(auto_batch_size=False)
    @set_skip_existing(None)
    def forward(
        self,
        tensordict: TensorDictBase,
        tensordict_out: TensorDictBase | None = None,
        _requires_sample: bool = True,
    ) -> TensorDictBase:
        if tensordict_out is None:
            tensordict_out = tensordict

        dist = self.get_dist(tensordict)
        if _requires_sample:
            out_tensors = self._dist_sample(dist, interaction_type=interaction_type())
            if isinstance(out_tensors, TensorDictBase):
                tensordict_out.update(out_tensors)                                        #   <= no using out_keys
                if self.return_log_prob:
                    tensordict_out = dist.log_prob(tensordict_out)
            else:
                if isinstance(out_tensors, Tensor):
                    out_tensors = (out_tensors,)
                tensordict_out.update(
                    {key: value for key, value in zip(self.out_keys, out_tensors)}    #  <= using out_keys
                )
                if self.return_log_prob:
                    log_prob = dist.log_prob(*out_tensors)
                    tensordict_out.set(self.log_prob_key, log_prob)
 

Additional context

I could not determine if this is an issue or a specification, but I decided to post it.
I apologize for my poor English.

Checklist

• I have checked that there is no similar issue in the repo (required)

[vmoens]

Have a look at the example
https://github.com/pytorch/rl/blob/main/examples/agents/composite_actor.py