Added fixed-size sequence length for GRU

benchmarl/algorithms/common.py

@@ -157,10 +157,10 @@ def get_replay_buffer(
         memory_size = self.experiment_config.replay_buffer_memory_size(self.on_policy)
         sampling_size = self.experiment_config.train_minibatch_size(self.on_policy)
         if self.has_rnn:
-            sequence_length = -(
-                -self.experiment_config.collected_frames_per_batch(self.on_policy)
-                // self.experiment_config.n_envs_per_worker(self.on_policy)
-            )
+            sequence_length = self.model_config.rnn_sequence_length
+            assert (
+                sequence_length > sampling_size
+            ), "Sequence length must be greater than the training minibatch size"
             memory_size = -(-memory_size // sequence_length)
             sampling_size = -(-sampling_size // sequence_length)
 

benchmarl/conf/model/layers/gru.yaml

@@ -6,6 +6,7 @@ n_layers: 1
 bias: True
 dropout: 0
 compile: False
+rnn_sequence_length: 20
 
 mlp_num_cells: [256, 256]
 mlp_layer_class: torch.nn.Linear

benchmarl/experiment/experiment.py

@@ -536,6 +536,12 @@ def _setup_algorithm(self):
             for group in self.group_map.keys()
         }
 
+        if self.algorithm.has_rnn:
+            total_seq = self.config.collected_frames_per_batch(
+                self.on_policy
+            ) // self.config.n_envs_per_worker(self.on_policy)
+            self.n_chunks = total_seq // self.model_config.rnn_sequence_length
+
     def _setup_collector(self):
         self.policy = self.algorithm.get_policy_for_collection()
 
@@ -737,6 +743,18 @@ def _collection_loop(self):
                 group_batch = self.algorithm.process_batch(group, group_batch)
                 if not self.algorithm.has_rnn:
                     group_batch = group_batch.reshape(-1)
+                else:
+                    n_chunks = self.n_chunks
+                    shape = group_batch.shape
+                    group_batch = group_batch[
+                        :, : n_chunks * self.config.rnn_sequence_length
+                    ]
+                    group_batch = group_batch.reshape(
+                        shape[0], n_chunks, self.config.rnn_sequence_length, *shape[2:]
+                    )
+                    group_batch = group_batch.reshape(
+                        shape[0] * n_chunks, self.config.rnn_sequence_length, *shape[2:]
+                    )
 
                 group_buffer = self.replay_buffers[group]
                 group_buffer.extend(group_batch.to(group_buffer.storage.device))

benchmarl/models/gru.py

@@ -165,7 +165,7 @@ def forward(
         # is_init never has it
 
         assert is_init is not None, "We need to pass is_init"
-        training = h_0 is None
+        training = h_0 is None or input.dim() == 4
 
         missing_batch = False
         if (
@@ -198,24 +198,27 @@ def forward(
         is_init = is_init.unsqueeze(-2).expand(batch, seq, self.n_agents, 1)
 
         if training:
-            if self.centralised and self.share_params:
-                shape = (
-                    batch,
-                    self.n_layers,
-                    self.hidden_size,
+            if h_0 is None:
+                if self.centralised and self.share_params:
+                    shape = (
+                        batch,
+                        self.n_layers,
+                        self.hidden_size,
+                    )
+                else:
+                    shape = (
+                        batch,
+                        self.n_agents,
+                        self.n_layers,
+                        self.hidden_size,
+                    )
+                h_0 = torch.zeros(
+                    shape,
+                    device=self.device,
+                    dtype=torch.float,
                 )
             else:
-                shape = (
-                    batch,
-                    self.n_agents,
-                    self.n_layers,
-                    self.hidden_size,
-                )
-            h_0 = torch.zeros(
-                shape,
-                device=self.device,
-                dtype=torch.float,
-            )
+                h_0 = h_0[:, 0]
         if self.centralised:
             input = input.view(batch, seq, self.n_agents * self.input_size)
             is_init = is_init[..., 0, :]
@@ -321,7 +324,7 @@ def __init__(
             is_critic=kwargs.pop("is_critic"),
         )
 
-        self.hidden_state_name = (self.agent_group, f"_hidden_gru_{self.model_index}")
+        self.hidden_state_name = (self.agent_group, f"hidden_gru_{self.model_index}")
         self.rnn_keys = unravel_key_list(["is_init", self.hidden_state_name])
         self.in_keys += self.rnn_keys
 
@@ -439,7 +442,7 @@ def _forward(self, tensordict: TensorDictBase) -> TensorDictBase:
         )
         h_0 = tensordict.get(self.hidden_state_name, None)
         is_init = tensordict.get("is_init")
-        training = h_0 is None
+        training = h_0 is None or input.dim() == (4 if self.input_has_agent_dim else 3)
 
         # Has multi-agent input dimension
         if self.input_has_agent_dim:
@@ -494,6 +497,7 @@ class GruConfig(ModelConfig):
     bias: bool = MISSING
     dropout: float = MISSING
     compile: bool = MISSING
+    rnn_sequence_length: int = MISSING
 
     mlp_num_cells: Sequence[int] = MISSING
     mlp_layer_class: Type[nn.Module] = MISSING
@@ -514,7 +518,7 @@ def is_rnn(self) -> bool:
     def get_model_state_spec(self, model_index: int = 0) -> Composite:
         spec = Composite(
             {
-                f"_hidden_gru_{model_index}": Unbounded(
+                f"hidden_gru_{model_index}": Unbounded(
                     shape=(self.n_layers, self.hidden_size)
                 )
             }

test/conftest.py

@@ -136,6 +136,7 @@ def gru_mlp_sequence_config() -> ModelConfig:
                 bias=True,
                 dropout=0,
                 compile=False,
+                rnn_sequence_length=5,
             ),
             MlpConfig(num_cells=[4], activation_class=nn.Tanh, layer_class=nn.Linear),
         ],