{"id":19736014,"url":"https://github.com/ikostrikov/jaxrl","last_synced_at":"2025-04-04T22:04:31.960Z","repository":{"id":39159677,"uuid":"330059679","full_name":"ikostrikov/jaxrl","owner":"ikostrikov","description":"JAX (Flax) implementation of algorithms for Deep Reinforcement Learning with continuous action spaces. ","archived":false,"fork":false,"pushed_at":"2022-10-26T04:30:25.000Z","size":4377,"stargazers_count":597,"open_issues_count":3,"forks_count":62,"subscribers_count":12,"default_branch":"main","last_synced_at":"2024-06-24T04:55:02.074Z","etag":null,"topics":["batch-reinforcement-learning","behavioral-cloning","continuous-control","deep-deterministic-policy-gradient","deep-learning","deep-reinforcement-learning","flax","gym","jax","offline-reinforcement-learning","reinforcement-learning","sac","soft-actor-critic"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/ikostrikov.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2021-01-16T01:26:44.000Z","updated_at":"2024-06-19T15:59:25.000Z","dependencies_parsed_at":"2023-01-19T16:31:34.910Z","dependency_job_id":null,"html_url":"https://github.com/ikostrikov/jaxrl","commit_stats":null,"previous_names":["ikostrikov/jax-rl"],"tags_count":1,"template":true,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ikostrikov%2Fjaxrl","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ikostrikov%2Fjaxrl/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ikostrikov%2Fjaxrl/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ikostrikov%2Fjaxrl/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/ikostrikov","download_url":"https://codeload.github.com/ikostrikov/jaxrl/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247256110,"owners_count":20909240,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["batch-reinforcement-learning","behavioral-cloning","continuous-control","deep-deterministic-policy-gradient","deep-learning","deep-reinforcement-learning","flax","gym","jax","offline-reinforcement-learning","reinforcement-learning","sac","soft-actor-critic"],"created_at":"2024-11-12T01:04:45.718Z","updated_at":"2025-04-04T22:04:31.930Z","avatar_url":"https://github.com/ikostrikov.png","language":"Jupyter Notebook","funding_links":[],"categories":["Reinforcement Learning (RL) and Deep Reinforcement Learning (DRL)","Libraries"],"sub_categories":["RL/DRL Algorithm Implementations and Software Frameworks","Inactive Libraries"],"readme":"[![DOI](https://zenodo.org/badge/330059679.svg)](https://zenodo.org/badge/latestdoi/330059679)\n\n# JAX (Flax) RL\n\nThis repository contains JAX (Flax) implementations of Reinforcement Learning algorithms:\n\n* [Soft Actor Critic with learnable temperature](https://arxiv.org/abs/1812.05905)\n* [Advantage Weighted Actor Critic](https://arxiv.org/abs/2006.09359)\n* [Image Augmentation Is All You Need](https://arxiv.org/abs/2004.13649)(only [K=1, M=1])\n* [Deep Deterministic Policy Gradient](https://arxiv.org/abs/1509.02971) with [Clipped Double Q-Learning](https://arxiv.org/abs/1802.09477)\n* [Randomized Ensembled Double Q-Learning: Learning Fast Without a Model](https://arxiv.org/abs/2101.05982)\n* Behavioral Cloning\n\nThe goal of this repository is to provide simple and clean implementations to build research on top of. **Please do not use this repository for baseline results and use the original implementations instead ([SAC](https://github.com/rail-berkeley/softlearning/), [AWAC](https://github.com/vitchyr/rlkit/tree/master/examples/awac), [DrQ](https://github.com/denisyarats/drq)).**\n\nIf you use JAXRL in your work, please cite this repository in publications:\n```\n@misc{jaxrl,\n  author = {Kostrikov, Ilya},\n  doi = {10.5281/zenodo.5535154},\n  month = {10},\n  title = {{JAXRL: Implementations of Reinforcement Learning algorithms in JAX}},\n  url = {https://github.com/ikostrikov/jaxrl},\n  year = {2021}\n}\n```\n\nYou can find an updated version of this repository [here](https://github.com/ikostrikov/jaxrl2).\n\n# Changelog\n\n## January 10th, 2022\n- Added an implementation of [Randomized Ensembled Double Q-Learning: Learning Fast Without a Model](https://arxiv.org/abs/2101.05982)\n\n## July 20th, 2021\n- Added an implementation of [Deep Deterministic Policy Gradient](https://arxiv.org/abs/1509.02971) with [Clipped Double Q-Learning](https://arxiv.org/abs/1802.09477)\n\n## May 19th, 2021\n- Added an implementation of [Soft Actor Critic v1](https://arxiv.org/abs/1801.01290)\n\n## April 29th, 2021\n- Added an implementation of data augmentation from [Image Augmentation Is All You Need](https://arxiv.org/abs/2004.13649)\n\n# Installation\n\nPrerequisites:\n* Python 3.8-3.9 (not yet 3.10)\n* [Poetry](https://python-poetry.org)\n* patchelf\n\nSuggested build environment:\n```bash\n# general build dependencies\nsudo apt-get update; sudo apt-get install make build-essential libssl-dev zlib1g-dev \\\nlibbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \\\nlibncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev\n# mujoco dependencies\napt-get -y install wget unzip software-properties-common \\\n    libgl1-mesa-dev \\\n    libgl1-mesa-glx \\\n    libglew-dev \\\n    libosmesa6-dev patchelf\n# mujoco installation\ncurl -OL https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz\nmkdir ~/.mujoco\ntar -zxf mujoco210-linux-x86_64.tar.gz -C ~/.mujoco\nrm mujoco210-linux-x86_64.tar.gz\n```\n\nTo install, run\n\n```bash\npoetry install\n# For GPU support run\npip install \"jax[cuda]==0.3.10\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n```\n\nFor further instructions on running this code on GPU, please follow instructions from [the official repository](https://github.com/google/jax).\n\nFor MuJoCo inslattion, you may need to add the following lines in the `.bashrc`:\n\n```bash\nexport LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/costa/.mujoco/mujoco210/bin\nexport LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/nvidia\n```\n\n# Development \n\nIf you want to modify the code, install following the instructions above.\n\n\n# [Examples](examples/)\n\n# Troubleshooting\n\nIf you experience out-of-memory errors, especially with enabled video saving, please consider reading [docs](https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html#gpu-memory-allocation) on JAX GPU memory allocation. Also, you can try running with the following environment variable:\n\n```bash\nXLA_PYTHON_CLIENT_MEM_FRACTION=0.80 python ...\n```\n\nIf you run your code on a remote machine and want to save videos for DeepMind Control Suite, please use EGL for rendering:\n```bash\nMUJOCO_GL=egl python train.py --env_name=cheetah-run --save_dir=./tmp/ --save_video\n```\n\n# Tensorboard\n\nLaunch tensorboard to see training and evaluation logs\n\n```bash\ntensorboard --logdir=./tmp/\n```\n\n# Results\n\n## Continous control from states\n\n![gym](./learning_curves/images/results.png)\n\n## Continous control from pixels\n\n![gym](./learning_curves/images/results_drq.png)\n\n\n# Docker\n\n## Build\n\nCopy your MuJoCo key to ./vendor\n\n```bash\ncd remote\ndocker build -t ikostrikov/jaxrl . -f Dockerfile \n```\n\n## Test\n```bash\n sudo docker run -v \u003cexamples-dir\u003e:/jaxrl/ ikostrikov/jaxrl:latest python /jaxrl/train.py --env_name=HalfCheetah-v2 --save_dir=/jaxrl/tmp/\n\n# On GPU\n sudo docker run --rm --gpus all -v \u003cexamples-dir\u003e:/jaxrl/ --gpus=all ikostrikov/jaxrl:latest python /jaxrl/train.py --env_name=HalfCheetah-v2 --save_dir=/jaxrl/tmp/\n```\n\n# Contributing\n\nWhen contributing to this repository, please first discuss the change you wish to make via issue. If you are not familiar with pull requests, please read [this documentation](https://opensource.com/article/19/7/create-pull-request-github).\n\n# Acknowledgements \n\nThanks to [@evgenii-nikishin](https://github.com/evgenii-nikishin) for helping with JAX. And [@dibyaghosh](https://github.com/dibyaghosh) for helping with vmapped ensembles.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fikostrikov%2Fjaxrl","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fikostrikov%2Fjaxrl","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fikostrikov%2Fjaxrl/lists"}