jax-imprl

A JAX accelerated version of IMPRL (Inspection and Maintenance Planning with Reinforcement Learning), a library for applying reinforcement learning to inspection and maintenance planning of deteriorating engineering systems.

Installation

1. Create a virtual environment

conda create --name jax_imprl_env -y python==3.9
conda activate jax_imprl_env

2. Install the dependencies

pip install poetry==1.8 # or conda install -c conda-forge poetry==1.8
poetry install

Installing additional packages

You can them add via poetry add (official docs) in the command line.

For example, to install Jupyter notebook,

# Allow >=7.1.2, <8.0.0 versions
poetry add notebook@^7.1.2

This will resolve the package dependencies (and adjust versions of transitive dependencies if necessary) and install the package. If the package dependency cannot be resolved, try to relax the package version and try again.

3. Setup wandb

For logging, the library relies on wandb. You can log into wandb using your private API key,

wandb login
# <enter wandb API key>

Docker (Optional)

You can use the following Docker image to run the code in a containerized environment.

docker pull nvidia/cuda

In case you don't have accesss to NVIDIA GPUs, you can rent a cloud instance here, and load the above Docker image. For example, vast.ai at ~$0.35/hour (pricing)

https://cloud.vast.ai/?ref_id=113803&creator_id=113803&name=JAX%2BRL

Related Work

• IMPRL

• IMP-MARL: a platform for benchmarking the scalability of cooperative MARL methods in real-world engineering applications.

  • Environments: (Correlated and uncorrelated) k-out-of-n systems and offshore wind structural systems.
  • RL solvers: Provides wrappers for interfacing with several (MA)RL libraries such as EPyMARL, Rllib, MARLlib etc.

Acknowledgements

This repository is inspired by the following projects:

PureJAXRL by luchris429

CleanRl started by vwxyzjn