Maskable Recurrent PPO

sb3_contrib/__init__.py

@@ -4,6 +4,7 @@
 from sb3_contrib.crossq import CrossQ
 from sb3_contrib.ppo_mask import MaskablePPO
 from sb3_contrib.ppo_recurrent import RecurrentPPO
+from sb3_contrib.ppo_mask_recurrent import MaskableRecurrentPPO
 from sb3_contrib.qrdqn import QRDQN
 from sb3_contrib.tqc import TQC
 from sb3_contrib.trpo import TRPO
@@ -21,4 +22,5 @@
     "CrossQ",
     "MaskablePPO",
     "RecurrentPPO",
+    "MaskableRecurrentPPO"
 ]

sb3_contrib/common/maskable_recurrent/type_aliases.py

@@ -0,0 +1,32 @@
+from typing import NamedTuple
+
+import torch as th
+from stable_baselines3.common.type_aliases import TensorDict
+from sb3_contrib.common.recurrent.type_aliases import (
+    RNNStates,
+)
+
+class MaskableRecurrentRolloutBufferSamples(NamedTuple):
+    observations: th.Tensor
+    actions: th.Tensor
+    old_values: th.Tensor
+    old_log_prob: th.Tensor
+    advantages: th.Tensor
+    returns: th.Tensor
+    lstm_states: RNNStates
+    episode_starts: th.Tensor
+    mask: th.Tensor
+    action_masks: th.Tensor
+
+
+class MaskableRecurrentDictRolloutBufferSamples(NamedTuple):
+    observations: TensorDict
+    actions: th.Tensor
+    old_values: th.Tensor
+    old_log_prob: th.Tensor
+    advantages: th.Tensor
+    returns: th.Tensor
+    lstm_states: RNNStates
+    episode_starts: th.Tensor
+    mask: th.Tensor
+    action_masks: th.Tensor

sb3_contrib/ppo_mask_recurrent/__init__.py

@@ -0,0 +1,4 @@
+from sb3_contrib.ppo_mask_recurrent.policies import CnnLstmPolicy, MlpLstmPolicy, MultiInputLstmPolicy
+from sb3_contrib.ppo_mask_recurrent.ppo_mask_recurrent import MaskableRecurrentPPO
+
+__all__ = ["CnnLstmPolicy", "MlpLstmPolicy", "MultiInputLstmPolicy", "MaskableRecurrentPPO"]

sb3_contrib/ppo_mask_recurrent/policies.py

@@ -0,0 +1,9 @@
+from sb3_contrib.common.maskable_recurrent.policies import (
+    MaskableRecurrentActorCriticPolicy,
+    MaskableRecurrentActorCriticCnnPolicy,
+    MaskableRecurrentMultiInputActorCriticPolicy,
+)
+
+MlpLstmPolicy = MaskableRecurrentActorCriticPolicy
+CnnLstmPolicy = MaskableRecurrentActorCriticCnnPolicy
+MultiInputLstmPolicy = MaskableRecurrentMultiInputActorCriticPolicy