Support Jax 0.6 outside of envpool

[adzcai]

What?

Right now running with Jax 0.6.0 causes any script to fail (see #157). This change makes it so that it only fails for the relevant envpool environments and other environments can run properly.

Why?

Preparing for a more significant fix later on to allow updated Jax versions.

How?

In make_env.py, we just put the EnvFactory import inside if TYPE_CHECKING since it's only used to annotate the return type of EnvFactory. We move the JaxEnvFactory inside the relevant branch as well.

[EdanToledo]

This is great! Can we just add to the readme this requirement so people are aware that envpool only works with jax less than 0.6.0. This unfortunately means all the sebulba architecture algorithm cannot be updated beyond this version of jax. I think unfortunately we are going to start getting into these issues where keeping up with jax versions becomes tough due to environment maintainers stopping to ensure their envs work with up to date jax versions. It might be worth starting to think about how to maintain this. Potentially we can just do it with release tags of stoix.

[adzcai]

Maybe the best way to ensure this is to move the envpool dependency into an extra in pyproject.toml, alongside the Jax constraint, so that users have to manually install it via uv sync --extra envpool:

[project.optional-dependencies]
envpool = [
    "envpool>=0.8.4 ; sys_platform == 'linux'",
    "jax>=0.4.25,<0.6.0",
]

Then inside the main project.dependencies we can leave it unconstrained. I've also added an explicit runtime check inside env_factory.py.