fix: make FlattenObservationWrapper also flatten next_obs

[JesseSilverberg]

What?

This PR modifies FlattenObservationWrapper to also perform the same flattening transform on next_obs (right now, the wrapper only modifies the current observation).

Why?

Currently, running an algorithm that uses a replay buffer on an environment with non-flat observations will fail (e.g., DQN + gymnax/freeway). The error is one of shapes, caused by FlattenObservationWrapper not flattening extras['next_obs'].

How?

Since this change involved duplicating some logic, I also took the liberty of moving a repetitive set of two lines to its own function.

Extra

Flagging that I'm not familiar enough with Stoix to know whether if 'next_obs' in timestep.extras: is the appropriate check/assumption. It looks like .extras will always be there, so I don't think this will cause issues, but wanted to bring it to your attention.

[EdanToledo]

Hello, thanks so much for this. The code looks great but could you run the pre-commit to format the code. You simply need to install the dev requirements using pip install -e .[dev] in the directory where you cloned the repo. or you can simply install the requirements-dev.txt. Then you need to run pre-commit run --all. That should format the code and point out if theres any things you are missing.

[JesseSilverberg]

Oops, sorry about that. Checks should be passing now.