Fix actor critic separate weights PackedSequence

sample_factory/model/actor_critic.py

@@ -4,6 +4,7 @@
 
 import torch
 from torch import Tensor, nn
+from torch.nn.utils.rnn import PackedSequence, pack_padded_sequence, pad_packed_sequence
 
 from sample_factory.algo.utils.action_distributions import is_continuous_action_space, sample_actions_log_probs
 from sample_factory.algo.utils.running_mean_std import RunningMeanStdInPlace, running_mean_std_summaries
@@ -218,19 +219,46 @@ def _core_rnn(self, head_output, rnn_states):
         This is actually pretty slow due to all these split and cat operations.
         Consider using shared weights when training RNN policies.
         """
-
         num_cores = len(self.cores)
-        head_outputs_split = head_output.chunk(num_cores, dim=1)
+
         rnn_states_split = rnn_states.chunk(num_cores, dim=1)
 
-        outputs, new_rnn_states = [], []
-        for i, c in enumerate(self.cores):
-            output, new_rnn_state = c(head_outputs_split[i], rnn_states_split[i])
-            outputs.append(output)
-            new_rnn_states.append(new_rnn_state)
+        if isinstance(head_output, PackedSequence):
+            # We cannot chunk PackedSequence directly, we first have to to unpack it,
+            # chunk, then pack chunks again to be able to process then through the cores.
+            # Finally we have to return concatenated outputs so we repeat the proces,
+            # but this time using concatenation - unpack, cat and pack.
+
+            unpacked_head_output, lengths = pad_packed_sequence(head_output)
+            unpacked_head_output_split = unpacked_head_output.chunk(num_cores, dim=2)
+            head_outputs_split = [
+                pack_padded_sequence(unpacked_head_output_split[i], lengths, enforce_sorted=False)
+                for i in range(num_cores)
+            ]
+
+            unpacked_outputs, new_rnn_states = [], []
+            for i, c in enumerate(self.cores):
+                output, new_rnn_state = c(head_outputs_split[i], rnn_states_split[i])
+                unpacked_output, lengths = pad_packed_sequence(output)
+                unpacked_outputs.append(unpacked_output)
+                new_rnn_states.append(new_rnn_state)
+
+            unpacked_outputs = torch.cat(unpacked_outputs, dim=2)
+            outputs = pack_padded_sequence(unpacked_outputs, lengths, enforce_sorted=False)
+        else:
+            head_outputs_split = head_output.chunk(num_cores, dim=1)
+            rnn_states_split = rnn_states.chunk(num_cores, dim=1)
+
+            outputs, new_rnn_states = [], []
+            for i, c in enumerate(self.cores):
+                output, new_rnn_state = c(head_outputs_split[i], rnn_states_split[i])
+                outputs.append(output)
+                new_rnn_states.append(new_rnn_state)
+
+            outputs = torch.cat(outputs, dim=1)
 
-        outputs = torch.cat(outputs, dim=1)
         new_rnn_states = torch.cat(new_rnn_states, dim=1)
+
         return outputs, new_rnn_states
 
     @staticmethod