Unity-Technologies/ml-agents · error · UnityTrainerException

Trying to export an attention mechanism that doesn't have a

Error message

Trying to export an attention mechanism that doesn't have a set                     number of elements.

What it means

A guard inside MultiHeadAttention's export path (forward/onnx export logic) that fires when entity_num_max_element was never set, so the attention module does not know the fixed number of entity elements required to build the ONNX graph. It is a sentinel error for an uninitialized setup step: the attention resolver was exported before the entity observations' sizes were recorded (self.entity_num_max_element is None), which makes static-shape export impossible.

Source

Thrown at ml-agents/mlagents/trainers/torch_entities/attention.py:158

            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(
            self.self_size + self.entity_size,
            1,
            self.embedding_size,
            kernel_init=Initialization.Normal,
            kernel_gain=(0.125 / self.embedding_size) ** 0.5,
        )

    def forward(self, x_self: torch.Tensor, entities: torch.Tensor) -> torch.Tensor:
        num_entities = self.entity_num_max_elements
        if num_entities < 0:
            if exporting_to_onnx.is_exporting():
                raise UnityTrainerException(
                    "Trying to export an attention mechanism that doesn't have a set max \
                    number of elements."
                )
            num_entities = entities.shape[1]

        if self.self_size > 0:
            expanded_self = x_self.reshape(-1, 1, self.self_size)
            expanded_self = torch.cat([expanded_self] * num_entities, dim=1)
            # Concatenate all observations with self
            entities = torch.cat([expanded_self, entities], dim=2)
        # Encode entities
        encoded_entities = self.self_ent_encoder(entities)
        return encoded_entities


class ResidualSelfAttention(torch.nn.Module):
    """
    Residual self attentioninspired from https://arxiv.org/pdf/1909.07528.pdf. Can be used

View on GitHub (pinned to 3ecb446f75)

Solutions

1. Ensure the model export (.onnx) is run after the network has processed at least one batch of real observations so entity sizes are known
2. Check that entity sensors are configured on the agent so the attention module receives entity observations
3. Upgrade ml-agents; export-shape handling for entity/attention networks has fixes in later releases

When it happens

Trigger: Thrown at ml-agents/mlagents/trainers/torch_entities/attention.py:158 when the library encounters an invalid state.

Common situations: See trigger scenarios.

AI-assisted analysis of Unity-Technologies/ml-agents@3ecb446f75 (2026-09-02). Data as JSON: /api/errors/e9c18efa47a653f9. Report an issue: GitHub.