These are some of the tools used on Towards Reinforcement Learning Controllers for Soft Robots Using Learned Environments, IEEE Robosoft 2024.
SoRoLEX is a framework for training and testing soft robot controllers using reinforcement learning with learned environments. Our implementation is fully built in JAX.
Requirements can be found in the lstm/setup/ directory. You can view them here.
Important Dependencies
Base Requirements:
- distrax==0.1.3
- flax==0.7.2
- gymnax==0.0.6
- pre-commit==3.3.3
- wandb==0.15.8
CPU-specific:
- jax==0.4.13
- jaxlib==0.4.13
GPU-specific:
- jax[cuda12_pip]==0.4.13
Please ensure you have the correct versions installed for your system (CPU or GPU).
- Build the Docker container with the provided script:
cd setup/docker && ./build.sh
- Add your WandB key to the
lstm/setup/dockerfolder:
echo <wandb_key> > setup/docker/wandb_key
👼 just add a wandb_key file without any extensions containing the key from the link above. the .gitignore is set up to ignore it and ensure the privacy of your key and your data.
Place the running script in the relevant directory
./run_docker.sh <gpu_id> python3 train.py <arguments>
For example, to train the agent on the learned environment using GPU 3, run:
cd envs
./run_docker.sh 3 python3 train.py
To train the lstm on GPU 5
cd lstm
./run_docker.sh 5 python3 train.py
This repo follows the jax-rl template. You can refer to that for more details
This work would have been possible without the following:
🚀 Jax Ecosystem ⚡
💪 Gymnax 🏋️♂️
🌟 PureJaxRL 🌟
If you use any of these tools, it would be really nice if you could please cite 😍 :
@INPROCEEDINGS{10522003,
author={Berdica, Uljad and Jackson, Matthew and Veronese, Niccolò Enrico and Foerster, Jakob and Maiolino Perla},
booktitle={2024 IEEE 7th International Conference on Soft Robotics (RoboSoft)},
title={Reinforcement Learning Controllers for Soft Robots Using Learned Environments},
year={2024},
pages={933-939},
doi={10.1109/RoboSoft60065.2024.10522003},
}