Using MaskedCategorical with ProbabilisticActor

[MaHaArt]

I’m working on a reinforcement learning scenario with discrete action types and continuous parameters. I’m using a ProbabilisticActor with a CompositeDistribution. Initially, I used Categorical for the discrete action type and masked invalid actions directly in the logits. As a result, the KL divergence started to explode during training.

I’m now considering switching to torchrl.modules.MaskedCategorical instead of Categorical. However, it seems that the mask is not being passed correctly.

Question: Has anyone successfully used MaskedCategorical with a ProbabilisticActor and could share some hints?