As we need a categorical sample for the next node position, we use the straight-through Gumbel softmax estimator [1]}, [2]}. In the practical implementation, we use dot-product attention [3]} to determine the logits \(z_{a_i \rightarrow v_j}\)  used for sampling the transition. To make the final graph-level prediction, pooling is applied on the agent embeddings, followed by a readout function. Edge features can also be trivially incorporated into this model. For details on this and the practical implementation of the model see Appendix . There we also discuss other simple extensions to the model, such as providing node count to the agents or global communication between them.
