├── .flake8 ├── .gitignore ├── .pre-commit-config.yaml ├── LICENSE ├── Makefile ├── README.md ├── docs └── _static │ └── cartpole_graph.png ├── jax_rl ├── MPO.py ├── SAC.py ├── TD3.py ├── __init__.py ├── buffers.py ├── dm_control_utils.py ├── gym_utils.py ├── main_dm_control.py ├── main_gym.py ├── models.py ├── saving.py ├── train_loops.py └── utils.py ├── poetry.lock ├── pyproject.toml └── tests ├── __init__.py ├── test_mpo.py ├── test_sac.py ├── test_td3.py └── utils.py /.flake8: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/.flake8 -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/.gitignore -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/.pre-commit-config.yaml -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/LICENSE -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/Makefile -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/README.md -------------------------------------------------------------------------------- /docs/_static/cartpole_graph.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/docs/_static/cartpole_graph.png -------------------------------------------------------------------------------- /jax_rl/MPO.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/jax_rl/MPO.py -------------------------------------------------------------------------------- /jax_rl/SAC.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/jax_rl/SAC.py -------------------------------------------------------------------------------- /jax_rl/TD3.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/jax_rl/TD3.py -------------------------------------------------------------------------------- /jax_rl/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /jax_rl/buffers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/jax_rl/buffers.py -------------------------------------------------------------------------------- /jax_rl/dm_control_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/jax_rl/dm_control_utils.py -------------------------------------------------------------------------------- /jax_rl/gym_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/jax_rl/gym_utils.py -------------------------------------------------------------------------------- /jax_rl/main_dm_control.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/jax_rl/main_dm_control.py -------------------------------------------------------------------------------- /jax_rl/main_gym.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/jax_rl/main_gym.py -------------------------------------------------------------------------------- /jax_rl/models.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/jax_rl/models.py -------------------------------------------------------------------------------- /jax_rl/saving.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/jax_rl/saving.py -------------------------------------------------------------------------------- /jax_rl/train_loops.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/jax_rl/train_loops.py -------------------------------------------------------------------------------- /jax_rl/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/jax_rl/utils.py -------------------------------------------------------------------------------- /poetry.lock: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/poetry.lock -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/pyproject.toml -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/test_mpo.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/tests/test_mpo.py -------------------------------------------------------------------------------- /tests/test_sac.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/tests/test_sac.py -------------------------------------------------------------------------------- /tests/test_td3.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/tests/test_td3.py -------------------------------------------------------------------------------- /tests/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/henry-prior/jax-rl/HEAD/tests/utils.py --------------------------------------------------------------------------------