`info` API compatibility with Gym

## 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](https://gist.github.com/vwxyzjn/02e92fb724fb790d7132c8833be1200f)

## 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