PPO attention net (GTrXLNet)

sb3_contrib/__init__.py

@@ -3,6 +3,7 @@
 from sb3_contrib.ars import ARS
 from sb3_contrib.ppo_mask import MaskablePPO
 from sb3_contrib.ppo_recurrent import RecurrentPPO
+from sb3_contrib.ppo_attention import AttentionPPO
 from sb3_contrib.qrdqn import QRDQN
 from sb3_contrib.tqc import TQC
 from sb3_contrib.trpo import TRPO
@@ -16,6 +17,7 @@
     "ARS",
     "MaskablePPO",
     "RecurrentPPO",
+    "AttentionPPO",
     "QRDQN",
     "TQC",
     "TRPO",

sb3_contrib/common/attention/type_aliases.py

@@ -0,0 +1,33 @@
+from typing import NamedTuple, Tuple
+
+import torch as th
+from stable_baselines3.common.type_aliases import TensorDict
+
+
+class AttnMemory(NamedTuple):
+    pi: th.Tensor #Tuple[th.Tensor, ...]
+    # vf: Tuple[th.Tensor, ...]
+
+
+class AttentionRolloutBufferSamples(NamedTuple):
+    observations: th.Tensor
+    actions: th.Tensor
+    old_values: th.Tensor
+    old_log_prob: th.Tensor
+    advantages: th.Tensor
+    returns: th.Tensor
+    attn_memory: AttnMemory
+    episode_starts: th.Tensor
+    mask: th.Tensor
+
+
+class AttentionDictRolloutBufferSamples(NamedTuple):
+    observations: TensorDict
+    actions: th.Tensor
+    old_values: th.Tensor
+    old_log_prob: th.Tensor
+    advantages: th.Tensor
+    returns: th.Tensor
+    attn_memory: AttnMemory
+    episode_starts: th.Tensor
+    mask: th.Tensor

sb3_contrib/ppo_attention/__init__.py

@@ -0,0 +1,4 @@
+from sb3_contrib.ppo_attention.policies import CnnAttnPolicy, MlpAttnPolicy, MultiInputAttnPolicy
+from sb3_contrib.ppo_attention.ppo_attention import AttentionPPO
+
+__all__ = ["CnnAttnPolicy", "MlpAttnPolicy", "MultiInputAttnPolicy", "AttentionPPO"]