[Proposal & Design Goals] Add `MJX` based version of the existing `MuJoCo` environments

[Kallinteris-Andreas]

Proposal

Create a reimplementation of the environments MuJuCo based on MJX

Unlike it was previously suggested, I suggest that the environments do not replace the existing MuJoCo environments, and therefore should be named something like "AntJ-v0" instead of "Ant-v6" because MJX is missing plenty of features

Alternative names are "AntJ-v5" or "AntJ-v6" (released along "Ant-v6")

The structure of code should be similar to the current MuJoCo environments, i.e a base class MJXEnv and then all the environments inherit it e.g. HalfCheetahMJXEnv(MJXEnv)

The ETA is 2024H2 as we are waiting for the MJX to implement some features such as tendons and some sensors

Along with this proposal, I will be posting an RFC of the prototype of the code for review. #834

Motivation

MJX is based on JAX and therefore support: XLA parallelization & JIT, allowing the users to collect many more environment steps per second compared to standard MuJoCo

Pitch

No response

Alternatives

The environments could instead be place in gymnasium_robotics, and I am also open to whatever naming scheme suggestions

Additional context

Old proposal for JAX based robotic environments: openai/gym#2456 (back then it was suggested to use brax, MJX did not exist yet)

Checklist

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

[Finebouche]

It would be incredible to have this

[pseudo-rnd-thoughts]

Would MJX/Ant-v0 work as a name? We might want a custom vector environment using vmap

[alik-git]

are there any updates on this? I would also like to use the gym api with an mjx backend for some work on humanoid robots.