Issues with SceneBatch.to_agent

[bmacadam-sfu]

Hi there,

I caught the following errors in the to_agent method for SceneBatch while trying to visualize some scenes:

1. The index_neighbors function does not preserve StateArrays/StateTensors. This can be fixed by checking if a statetensor is passed and using its formatting.

        def index_neighbors(x: Tensor | StateTensor) -> Tensor | StateTensor:
            index_neighbors = x[others_mask].reshape([batch_size, num_agents-1]+list(x.shape[2:]))
            if isinstance(x, StateTensor):
                index_neighbors = StateTensor.from_array(index_neighbors, x._format)
            return index_neighbors

2. The index_agent function doesn't play well with map names, this can be fixed by wrapping each map_name into a singleton list:

            map_names=index_agent_list([[m] for m in self.map_names]),

[bmacadam-sfu]

Of course, for the index_neighbors problem, you would presumably want to fix the underlying issues with reshape and StateArrays/StateTensors.