[Bug Report] Jax, pytorch and dependency imports when importing gymnasium

[vmoens]

Describe the bug

Importing Jax and PyTorch when doing

import gymnasium

(which occurs via imports in gymnasium.experimental.wrappers is the cause of multiple issues:

• It does not catch errors that are unrelated to an ImportError. If pytorch or jax have errors such as OSError or other we can't import gymnasium for reasons that are unrelated to gymnasium itself. PyTorch and Jax should not condition gymnasium usage.
• Besides, when running scripts in parallel with multiprocessing in spawn mode, importing gymnasium is an operation one would expect to be super fast and hence not to cause major slowdowns in code execution. However, import torch is a very slow operation, and import gymnasium will suffer from it.

## Solution

Any import to optional dependencies should be hidden from the general library loading.

Code example

No response

System info

No response

Additional context

No response

Checklist

• I have checked that there is no similar issue in the repo

[pseudo-rnd-thoughts]

Jax and PyTorch are optional installs but I agree that we import them at the beginning of the file and will cause slow downs if the modules are not actually used.
Do you know what the best solution to this is?

[RedTachyon]

We should probably rearrange what is imported in which __init__, currently around experimental, but when those things become stable, then throughout the repo.

For now the only "leaking" part do seem to be the wrappers (gymnasium/__init__.py -> gymnasium/experimental/__init__.py -> gymnasium/experimental/wrappers/__init__.py -> all individual wrappers), so if we remove that from the exports, it should be enough.

More generally, I'd consider removing experimental from the root __init__ in general - this way, it only gets imported if someone explicitly does import gymnasium.experimental<...>. Even then, we should keep the different dependencies modular, so that e.g. accessing gymnasium.experimental.VectorEnv doesn't touch any jax or pytorch stuff

FWIW this is what I meant when I was talking about making sure that none of the new experimental features things leak out to the main repo, but I didn't actually check the current state, so that's my bad. I can try to unfuck this before the next release. (though it might be as simple as that one import I mentioned earlier)

[RedTachyon]

Closed via #323, please reopen if this is (still/again) an issue