Fix import errors for matplotlib
and seaborn
visualization.
-
Fixed most of the typing issues under pylint --strict. 1.1 This required unifying the interface of
step
,step_env
,get_obs
, 'is_terminal` for proper inheritance. -
spaces
now depends ongymnasium
instead ofgym
, adding an additional dependency. Maybe later ongym
dependency can be completely removed since it's not maintained anymore? haven't looked into that. -
Removed
_DeviceArray
from the tests as it's deprecated. Uses jax version0.4.24
-
Ran everything in
python 3.10
so technically you could support that too now. -
In
wrapper/brax.py
,GymnaxtoBraxWrapper
has two new methodsbackend
andobservation_size
, as its required to be defined under newest brax version. Used brax version0.10.0
.
Thanks to @Aidandos !
- Gym, EvoJAX, Brax, DM env wrappers
- Reacher environment inspired by Lenton et al. (2021)
- Swimmer environment inspired by Lenton et al. (2021)
- Basic Pong environment inspired by Kirsch (2018)
- Fixed Minatar tests for jax arrays.
- Fixed reward setting in
DiscountingChain-v0
- Fixed
Tuple
space check via enumerate.
- Refactored wrappers into separate sub-directory.
- Fix deprecated
tree_multimap
. - Fix device grabbing when using an
jnp.array
to set default in Acrobot env.
- Fix import structure.
- Fully functional API using flax-style
EnvState
andEnvParams
for callingenv.step
. - MinAtar environments are not operational yet.
- Release to ensure that
evosax
0.0.9 release can work withGymFitness
backend ingymnax
.
- Basically everything :)
- Adds main
gym
-like infrastructure that isjit
,vmap
&pmap
compatible:- Jittable environment base class:
environment.py
- Spaces:
Discrete
,Continuous
,Box
- Translation of numpy state to state dictionary for JAX
step
.
- Jittable environment base class:
- Adds base set of environments:
- OpenAI's
classic_control
:Pendulum-v0
,CartPole-v0
,MountainCar-v0
,ContinuousMountainCar-v0
,Acrobot-v0
- DeepMind's
bsuite
:Catch-bsuite
,DeepSea-bsuite
,DiscountingChain-bsuite
,MemoryChain-bsuite
,UmbrellaChain-bsuite
,MNISTBandit-bsuite
,SimpleBandit-bsuite
MinAtar
:Asterix-MinAtar
,Breakout-MinAtar
,Freeway-MinAtar
,Seaquest-MinAtar
,SpaceInvaders-MinAtar
- OpenAI's
- Adds
tests
for comparinggym
/numpy
reset
+step
with JAX version.tests/classic_control/test_gym_env.py
tests/bsuite/test_bsuite_env.py
- Adds set of
notebooks
walking through the individual environments. - Adds set of
examples
incorporatinggymnax
into JAX-based RL experiments.- Anakin agent - port DM Colab for gymnax
- CMA-ES policy evolution - port blogpost with experimental minimal agent
- Adds benchmark infrastructure and numbers on different devices (CPU/GPU/TPU):
- Transitions/Second & Specific rollout types vs Torch setup
- CPU: Intel Xeon 2.4 GHz
- GPU: V100, A100, RTX 2080Ti
- TPU: V2, V3 -
vmap
+pmap
- Transitions/Second & Specific rollout types vs Torch setup
- Adds set of
experimental
utilities:dojos
: Multi-transition rollout wrapper vialax.scan
+jit
for sequentialstep
.agents
: Minimal and evaluation agent wrappers.
- Adds miscellaneous environments:
Bandit-misc
,Rooms-misc