prob_actor_masked.py

prob_actor_masked.py

from tensordict import TensorDict
from tensordict.nn import TensorDictModule as Mod
from torchrl.modules import MaskedCategorical, ProbabilisticActor, MLP
from torchrl.envs import set_exploration_type
import torch

td = TensorDict(
    observation=torch.randn(3, 4),
    mask=torch.zeros(3, 10, dtype=torch.bool).bernoulli_(0.9),
    batch_size=(3,)
)

policy = ProbabilisticActor(
    module=Mod(MLP(in_features=4, out_features=10, num_cells=(32,)), in_keys=["observation"], out_keys=["logits"]),
    in_keys=["logits", "mask"],
    distribution_class=MaskedCategorical,
    return_log_prob=True,
)
with set_exploration_type("RANDOM"):
    print(policy(td))