What is the problem

Right now Brax provides a slightly different Gym vector API in the info field.

Gym's info dict is returned per environment, like if I have 5 envs, then the dict could look like
[{}, {}, {}, {}, {["truncated": True}

In Brax, it looks like

{"truncated":[False, False, False, False, True]}

This makes it difficult to develop wrappers compatible with both gym and brax. See #89.

As a result, gym wrapper + jumpy with brax does not work out of the box: see code

Pros and Cons

Brax's approach feels more ergonomic for high-throughput envs (note that envpool also adopts brax approach in the info field), whereas Gym's approach is maybe more efficient for "sparse" info keys that don't appear often (I could be wrong with this).

Not sure what's the best way to do this. CC @RedTachyon @jkterry1 @lebrice