Attempting to use EntityEmbedding directly as processor

ml-agents/mlagents/trainers/tests/torch/test_attention.py

     torch.manual_seed(1336)
     size, n_k, = 3, 5
     embedding_size = 64
-    entity_embeddings = EntityEmbedding(size, size, n_k, embedding_size, True)
+    entity_embeddings = EntityEmbedding(size, n_k, embedding_size)
+    entity_embeddings.add_self_embedding(size)
     transformer = ResidualSelfAttention(embedding_size, n_k)
     l_layer = linear_layer(embedding_size, size)
     optimizer = torch.optim.Adam(
     n_k = 5
     size = n_k + 1
     embedding_size = 64
-    entity_embedding = EntityEmbedding(
-        size, size, n_k, embedding_size, concat_self=False
-    )
+    entity_embedding = EntityEmbedding(size, n_k, embedding_size)  # no self
     transformer = ResidualSelfAttention(embedding_size)
     l_layer = LinearEncoder(embedding_size, 2, n_k)
     loss = torch.nn.CrossEntropyLoss()

ml-agents/mlagents/trainers/tests/torch/test_utils.py

         obs_shapes.append(vis_obs_shape)
     h_size = 128
     obs_spec = create_observation_specs_with_shapes(obs_shapes)
-    encoders, _, embedding_sizes = ModelUtils.create_input_processors(
+    encoders, embedding_sizes = ModelUtils.create_input_processors(
         obs_spec, h_size, encoder_type, normalize
     )
     total_output = sum(embedding_sizes)

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

 
     def __init__(
         self,
-        x_self_size: int,
-        concat_self: bool,
     ):
         """
         Constructs an EntityEmbedding module.
             self-attention.
         """
         super().__init__()
-        self.self_size: int = x_self_size
+        self.self_size: int = 0
-        # If not concatenating self, input to encoder is just entity size
-        if not concat_self:
-            self.self_size = 0
+        self.embedding_size = embedding_size
+            self.entity_size,
+            1,
+            self.embedding_size,
+            kernel_init=Initialization.Normal,
+            kernel_gain=(0.125 / self.embedding_size) ** 0.5,
+        )
+
+    def add_self_embedding(self, size: int) -> None:
+        self.self_size = size
+        self.self_ent_encoder = LinearEncoder(
-            embedding_size,
+            self.embedding_size,
-            kernel_gain=(0.125 / embedding_size) ** 0.5,
+            kernel_gain=(0.125 / self.embedding_size) ** 0.5,
         )
 
     def forward(self, x_self: torch.Tensor, entities: torch.Tensor) -> torch.Tensor:

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

 from mlagents.trainers.torch.model_serialization import exporting_to_onnx
 
 
-class Identity(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
-        return inputs
-
-
 class Normalizer(nn.Module):
     def __init__(self, vec_obs_size: int):
         super().__init__()

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

 from mlagents.trainers.torch.utils import ModelUtils
 from mlagents.trainers.torch.decoders import ValueHeads
 from mlagents.trainers.torch.layers import LSTM, LinearEncoder, Initialization
-from mlagents.trainers.torch.encoders import VectorInput, Identity
+from mlagents.trainers.torch.encoders import VectorInput
+    EntityEmbedding,
     ResidualSelfAttention,
     get_zero_entities_mask,
 )
             else 0
         )
 
-        self.processors, self.var_processors, self.embedding_sizes = ModelUtils.create_input_processors(
+        self.processors, self.embedding_sizes = ModelUtils.create_input_processors(
             observation_specs,
             self.h_size,
             network_settings.vis_encode_type,
         entity_num_max: int = 0
-        for var_processor in self.var_processors:
-            entity_max: int = var_processor.entity_num_max_elements
+        var_processors = [p for p in self.processors if isinstance(p, EntityEmbedding)]
+        for processor in var_processors:
+            entity_max: int = processor.entity_num_max_elements
-        if len(self.var_processors) > 0:
+        if len(var_processors) > 0:
             if sum(self.embedding_sizes):
                 self.x_self_encoder = LinearEncoder(
                     sum(self.embedding_sizes),
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         encodes = []
         var_len_inputs = []  # The list of variable length inputs
+        var_len_processors = [
+            p for p in self.processors if isinstance(p, EntityEmbedding)
+        ]
-            if not isinstance(processor, Identity):
+            if not isinstance(processor, EntityEmbedding):
                 # The input can be encoded without having to process other inputs
                 obs_input = inputs[idx]
                 processed_obs = processor(obs_input)
             masks = get_zero_entities_mask(var_len_inputs)
             embeddings: List[torch.Tensor] = []
             processed_self = self.x_self_encoder(encoded_self) if input_exist else None
-            for var_len_input, var_len_processor in zip(
-                var_len_inputs, self.var_processors
-            ):
-                embeddings.append(var_len_processor(processed_self, var_len_input))
+            for var_len_input, processor in zip(var_len_inputs, var_len_processors):
+                embeddings.append(processor(processed_self, var_len_input))
             qkv = torch.cat(embeddings, dim=1)
             attention_embedding = self.rsa(qkv, masks)
             if not input_exist:
                 end = start + vec_size
                 inputs.append(concatenated_vec_obs[:, start:end])
                 start = end
-            elif enc is not None:
+            elif isinstance(enc, EntityEmbedding):
+                inputs.append(var_len_inputs[var_len_index])
+                var_len_index += 1
+            else:  # visual input
-            else:
-                inputs.append(var_len_inputs[var_len_index])
-                var_len_index += 1
+
         # End of code to convert the vec and vis obs into a list of inputs for the network
         encoding, memories_out = self.network_body(
             inputs, memories=memories, sequence_length=1

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

     NatureVisualEncoder,
     SmallVisualEncoder,
     VectorInput,
-    Identity,
 )
 from mlagents.trainers.settings import EncoderType, ScheduleType
 from mlagents.trainers.torch.attention import EntityEmbedding
             return (VectorInput(shape[0], normalize), shape[0])
         # VARIABLE LENGTH
         if dim_prop in ModelUtils.VALID_VAR_LEN_PROP:
-            return (Identity, 0)
+            return (
+                EntityEmbedding(
+                    entity_size=shape[1],
+                    entity_num_max_elements=shape[0],
+                    embedding_size=h_size,
+                ),
+                0,
+            )
         # OTHER
         raise UnityTrainerException(f"Unsupported Sensor with specs {obs_spec}")
 
         h_size: int,
         vis_encode_type: EncoderType,
         normalize: bool = False,
-    ) -> Tuple[nn.ModuleList, nn.ModuleList, List[int]]:
+    ) -> Tuple[nn.ModuleList, List[int]]:
         """
         Creates visual and vector encoders, along with their normalizers.
         :param observation_specs: List of ObservationSpec that represent the observation dimensions.
             obs.
         :param normalize: Normalize all vector inputs.
         :return: Tuple of :
-         - ModuleList of the encoders (Identity if the input requires to be processed with a variable length
-         observation encoder)
+         - ModuleList of the encoders
-         - A list of the inputs that need to be processed by a variable length observation encoder.
-        var_encoders: List[nn.Module] = []
-        var_len_indices: List[int] = []
-        for idx, obs_spec in enumerate(observation_specs):
+        for obs_spec in observation_specs:
-            if encoder is Identity:
-                var_len_indices.append(idx)
-        for idx in var_len_indices:
-            var_encoders.append(
-                EntityEmbedding(
-                    x_self_size=0 if x_self_size == 0 else h_size,
-                    entity_size=observation_specs[idx].shape[1],
-                    entity_num_max_elements=observation_specs[idx].shape[0],
-                    embedding_size=h_size,
-                    concat_self=True,
-                )
-            )
-        return (nn.ModuleList(encoders), nn.ModuleList(var_encoders), embedding_sizes)
+        if x_self_size > 0:
+            for enc in encoders:
+                if isinstance(enc, EntityEmbedding):
+                    enc.add_self_embedding(h_size)
+        return (nn.ModuleList(encoders), embedding_sizes)
 
     @staticmethod
     def list_to_tensor(