search

aksub99 / MolDQN-pytorch

A PyTorch Implementation of "Optimization of Molecules via Deep Reinforcement Learning".
MIT License
74 stars 28 forks source link
chemistry deep-learning dqn-pytorch drug-discovery inverse-design machine-learning materials-informatics materials-science molecule python pytorch pytorch-implementation pytorch-rl reinforcement-learning

readme

MolDQN-pytorch

MIT license made-with-python

PyTorch implementation of MolDQN as described in Optimization of Molecules via Deep Reinforcement Learning by Zhenpeng Zhou, Steven Kearnes, Li Li, Richard N. Zare and Patrick Riley.

Installation

From source:

1) Install rdkit.
conda create -c rdkit -n my-rdkit-env rdkit
conda activate my-rdkit-env
conda install -c conda-forge rdkit

2) Clone this repository.
git clone https://github.com/aksub99/MolDQN-pytorch.git
cd MolDQN-pytorch

3) Install the requirements given in requirements.txt.
pip install -r requirements.txt

4) Install baselines.
pip install "git+https://github.com/openai/baselines.git"

From Docker:

Using a docker image requires an NVIDIA GPU. If you do not have a GPU please follow the directions for installing from source In order to get GPU support you will have to use the nvidia-docker plugin.

# Build the Dockerfile in Dockerfiles/Dockerfile to create a Docker image.
cd Dockerfiles
docker build -t moldqn_pytorch:latest .

# This will create a container from the image we just created.
nvidia-docker run -[Options] moldqn_pytorch:latest python path/to/main.py

Please remember to modify the TB_LOG_PATH variable in main.py depending on where you wish to store your tensorboard runs file.

Training the MolDQN:

python main.py

A simple example to train the model on a single property optimization task can be seen in examples/MolDQN-pytorch.ipynb.

Results:

The following was the reward curve obtained when the model was trained for 5000 episodes on a single property optimization task (QED in this case).

References:

The original tensorflow implementation can be found at https://github.com/google-research/google-research/tree/master/mol_dqn This repository re-uses some code from the original implementation.