Jaxplorer is a Jax reinforcement learning (RL) framework for exploring new ideas.
Warning
This project is still experimental, APIs could change without notice.
Pure Jax RL tasks are not supported at this stage.
Note
For PyTorch version, please check Explorer.
- Add more descriptions about slurm and experimental result analysis.
- Add more algorithms, such as DQN for Atari games.
- Python: 3.11
- Jax: >=0.4.20
- MuJoCo:
pip install 'mujoco>=2.3.6,<3.0' - Gymnasium:
pip install 'gymnasium[box2d,mujoco]>=0.29.1,<1.0' - Gym Games: >=2.0.0.
- Others:
pip install -r requirements.txt.
- Deep Q-Learning (DQN)
- Double Deep Q-learning (DDQN)
- Maxmin Deep Q-learning (MaxminDQN)
- Proximal Policy Optimisation (PPO)
- Soft Actor-Critic (SAC)
- Deep Deterministic Policy Gradients (DDPG)
- Twin Delayed Deep Deterministic Policy Gradients (TD3)
- Continuous Deep Q-Learning with Model-based Acceleration (NAF): model-free version; a different exploration strategy is applied for simplicity.
All hyperparameters including parameters for grid search are stored in a configuration file in directory configs. To run an experiment, a configuration index is first used to generate a configuration dict corresponding to this specific configuration index. Then we run an experiment defined by this configuration dict. All results including log files are saved in directory logs. Please refer to the code for details.
For example, run the experiment with configuration file classic_dqn.json and configuration index 1:
python main.py --config_file ./configs/classic_dqn.json --config_idx 1
The models are tested for one episode after every test_per_episodes training episodes which can be set in the configuration file.
First, we calculate the number of total combinations in a configuration file (e.g. classic_dqn.json):
python utils/sweeper.py
The output will be:
Number of total combinations in classic_dqn.json: 18
Then we run through all configuration indexes from 1 to 18. The simplest way is using a bash script:
for index in {1..18}
do
python main.py --config_file ./configs/classic_dqn.json --config_idx $index
doneParallel is usually a better choice to schedule a large number of jobs:
parallel --eta --ungroup python main.py --config_file ./configs/classic_dqn.json --config_idx {1} ::: $(seq 1 18)Any configuration index that has the same remainder (divided by the number of total combinations) should have the same configuration dict. So for multiple runs, we just need to add the number of total combinations to the configuration index. For example, 5 runs for configuration index 1:
for index in 1 19 37 55 73
do
python main.py --config_file ./configs/classic_dqn.json --config_idx $index
done
Or a simpler way:
parallel --eta --ungroup python main.py --config_file ./configs/classic_dqn.json --config_idx {1} ::: $(seq 1 18 90)
Slurm is supported as well. Please check submit.py.
To analyze the experimental results, just run:
python analysis.py
Inside analysis.py, unfinished_index will print out the configuration indexes of unfinished jobs based on the existence of the result file. memory_info will print out the memory usage information and generate a histogram to show the distribution of memory usages in directory logs/classic_dqn/0. Similarly, time_info will print out the time information and generate a histogram to show the distribution of time in directory logs/classic_dqn/0. Finally, analyze will generate csv files that store training and test results. Please check analysis.py for more details. More functions are available in utils/plotter.py.
If you find this repo useful to your research, please cite this repo:
@misc{Jaxplorer,
author = {Lan, Qingfeng},
title = {A Jax Reinforcement Learning Framework for Exploring New Ideas},
year = {2024},
publisher = {GitHub},
journal = {GitHub Repository},
howpublished = {\url{https://github.com/qlan3/Jaxplorer}}
}