Use 4 heads

ml-agents/mlagents/trainers/torch/networks.py

         self.entity_encoder = EntityEmbeddings(
             sum(_input_size), [sum(_input_size)], self.h_size
         )
-        self.self_attn = ResidualSelfAttention(self.h_size, num_heads=1)
+        self.self_attn = ResidualSelfAttention(self.h_size)
 
         encoder_input_size = self.h_size