Multi-Agent Data Canonicalization

In cooperative multi agent reinforcement learning with homogeneous agents, each agent collects its own experience. However, because the agents are interchangeable, their experiences can be shared to improve sample efficiency but only if the observations are first transformed into a consistent canonical form. Each agent's observation is a matrix where the rows correspond to the feature vectors of all agents in the environment. The row order depends on the observing agent's perspective: its own features might appear at any position depending on its index. Implement a function canonicalize observations that standardizes a batch of multi agent observations so that: 1. The observing agent's own feature vector is always placed in the first row. 2. The remaining agents' feature vectors are sorted in lexicographic order (sorted primarily by the first feature, then by the second feature for ties, and so on). This ensures that regardless of which agent produced the observation, the resulting canonical form is identical for equivalent environmental configurations. Inputs: observations: A numpy array of shape (B, N, F) where B is the batch size, N is the number of agents, and F is the feature dimension per agent. agent ids: A numpy array of shape (B,) of integers indicating which row index in each observation corresponds to the observing agent (the "self" agent). Returns: A numpy array of shape (B, N, F) containing the canonicalized observations, with values rounded to 4 decimal places.