info API compatibility with Gym

[vwxyzjn]

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

[erikfrey]

The Gym wrapper is explicitly meant for correct interop with Gym users - Brax is unopinionated and happy to support whatever is the expected output here. Just let us know what change you would like and we'll make it.

[erikfrey]

Actually taking a second read through this issue, I am concerned this would tank performance when stepping through vectorized environments. Turning a single dict with 2048-length arrays into 2048 dicts with 1-length arrays could be expensive, particularly if it's done every step.

You could try it, but I suspect it will slow things down significantly.

[vwxyzjn]

Turning a single dict with 2048-length arrays into 2048 dicts with 1-length arrays could be expensive, particularly if it's done every step.

That's my thought on this, too. Hence the reason maybe gym should reconsider how the info API is designed. The current design is not friendly to these types of 2048-length arrays.