Observation norm for RL problems with early stopping via Welford's parallel algorithm

[keraJLi]

This addresses #97, which was prematurely closed by @maxencefaldor.

Switched to Welford's parallel algorithm for combining statistics collected during policy rollouts with global observation normalization statistics. This way, we don't have to go up to the maximum number of episode steps for every rollout, which can be super expensive.

This does not increase implementation complexity:

• gymnax.py: 63 insertions(+), 61 deletions(-)
• brax.py: 68 insertions(+), 64 deletions(-)

So +6 lines total, even though my code is very verbose.

I have tested both GymnaxProblem and BraxProblem by running ES training and calculating the statistics of the normalized rewards.

Caveat: we don't return all env states of the complete rollout. Doing this might be nice for visualization, but then we would also have to return state validity and mask out invalid states, which is not currently implemented and might easily be forgotten by users of the library. I suggest rolling out policies again for visualization, which is orders of magnitude cheaper than running a single generation.

[keraJLi]

Before this is merged, we should update the RL example notebook to handle the missing rollout states for brax.

[keraJLi]

I have updated the notebook, which now explicitly rolls out the search distribution mean. This did not create any notable extra complexity.