Add padding to decoded DCT coefficients in decode method so that it always tries

processing_action_tokenizer.py

@@ -78,8 +78,13 @@ class UniversalActionProcessor(ProcessorMixin):
         decoded_actions = []
         for token in tokens:
             decoded_tokens = self.bpe_tokenizer.decode(token)
-            decoded_dct_coeff = np.array(list(map(ord, decoded_tokens))) + self.min_token
-            decoded_dct_coeff = decoded_dct_coeff.reshape(-1, self.action_dim)
+            decoded_flat_dct_coeff = np.array(list(map(ord, decoded_tokens))) + self.min_token
+            unpadded_size = decoded_flat_dct_coeff.size
+            padded_size = self.time_horizon * self.action_dim
+            padded_flat_dct_coeff = np.zeros(shape=padded_size, dtype=decoded_flat_dct_coeff.dtype)
+            copy_size = min(unpadded_size, padded_size)
+            padded_flat_dct_coeff[:copy_size] = decoded_flat_dct_coeff[:copy_size]
+            decoded_dct_coeff = padded_flat_dct_coeff.reshape(-1, self.action_dim)
             assert (
                 decoded_dct_coeff.shape
                 == (