GENTRL (Generative Tensorial Reinforcement Learning)

GENTRL is a variational autoencoder (VAE) that uses a rich prior distribution in its latent space. It employs tensor decompositions to encode relationships between molecular structures and their properties, which also allows it to learn from data with missing values.

The training process consists of two distinct stages:

1. Autoencoder Pretraining: The model learns a mapping of the chemical space onto a latent manifold by maximizing the evidence lower bound (ELBO).
2. Reinforcement Learning (RL) Optimization: After pretraining, all model parameters are frozen except for the learnable prior. The model then explores the chemical space to identify molecules that maximize a specific reward function.

To use GENTRL for molecular generation, follow these steps to set up the environment and run the training pipelines:

1. Install RDKit: Required for molecular processing. Follow the official RDKit installation guide.
2. Install GENTRL: Install the package from the local repository using python setup.py install.
3. Install MOSES: The example uses the MOSES dataset. Install it from the MOSES repository.
4. Pretrain the Autoencoder: Run the pretrain.ipynb notebook to learn the mapping of chemical space on the latent manifold.
5. Optimize via Reinforcement Learning: Run the train_rl.ipynb notebook to explore the chemical space and find molecules with high rewards.

# Install GENTRL
python setup.py install

To create a GENTRL model, you need to instantiate an encoder and a decoder, then pass them to the gentrl.GENTRL class.

Common components include:

• gentrl.RNNEncoder(latent_size=...): An encoder that maps molecules to a latent space.
• gentrl.DilConvDecoder(latent_input_size=...): A decoder that reconstructs molecules from the latent space.

The GENTRL constructor requires the encoder, decoder, and configuration for layers (e.g., 50 * [('c', 20)]) and latent inputs.

Note: Ensure you move the model to a GPU using .cuda() if using CUDA-enabled hardware.

import gentrl
import torch

enc = gentrl.RNNEncoder(latent_size=50)
dec = gentrl.DilConvDecoder(latent_input_size=50)
model = gentrl.GENTRL(enc, dec, 50 * [('c', 20)], [('c', 20)], beta=0.001)
model.cuda()

Use the model.sample(n) method to generate a batch of n molecular SMILES strings from the trained model.

Because generative models may produce invalid SMILES, it is a best practice to filter the output using a molecular parser (like RDKit's get_mol) to ensure only valid molecules are used for downstream tasks or visualization.

# Sample 100 molecules
sampled = model.sample(100)

# Filter for valid molecules
sampled_valid = [s for s in sampled if get_mol(s)]

You can restore a previously trained GENTRL model from a saved directory using the .load() method.

model.load('saved_gentrl_after_rl/')
model.cuda()

To create a GENTRL model, you must provide an encoder and a decoder. The gentrl.GENTRL constructor takes the encoder, the decoder, a list of layer configurations for the encoder, a list of layer configurations for the decoder, and a beta parameter.

Example configuration:

• enc: An instance of gentrl.RNNEncoder.
• dec: An instance of gentrl.DilConvDecoder.
• latent_size: The size of the latent space.
• beta: A hyperparameter for the model training.

You can move the model to a GPU using .cuda().

import gentrl
import torch

torch.cuda.set_device(0)

enc = gentrl.RNNEncoder(latent_size=50)
dec = gentrl.DilConvDecoder(latent_input_size=50)
# The lists define layer configurations (e.g., [('c', 20)])
model = gentrl.GENTRL(enc, dec, 50 * [('c', 20)], [('c', 20)], beta=0.001)
model.cuda()

The gentrl.MolecularDataset class is used to load molecular data for training. You provide a list of sources, where each source is a dictionary defining:

• path: The path to the CSV file.
• smiles: The column name containing SMILES strings.
• prob: The probability/weight of this source (e.g., 1).
• Additional property columns (e.g., 'plogP': 'plogP').

You must also specify the props argument, which is a list of the property column names you want to include in the dataset.

import gentrl
from torch.utils.data import DataLoader

md = gentrl.MolecularDataset(sources=[
    {
        'path': 'train_plogp_plogpm.csv',
        'smiles': 'SMILES',
        'prob': 1,
        'plogP': 'plogP',
    }
], props=['plogP'])

train_loader = DataLoader(md, batch_size=50, shuffle=True, num_workers=1, drop_last=True)

You can persist GENTRL model states to disk using .save(path) and restore them using .load(path). Ensure the directory exists before saving.

# Loading a model
model.load('saved_gentrl/')
model.cuda()

# Saving a model
model.save('./saved_gentrl_after_rl/')

You can persist the state of your trained GENTRL model to a directory using the .save() method.

import os

os.makedirs('./saved_gentrl/', exist_ok=True)
model.save('./saved_gentrl/')

Once the model and data loader are initialized, you can start the pretraining process using the train_as_vaelp method. This method typically requires a DataLoader and a learning rate (lr).

# Assuming 'model' and 'train_loader' are already defined
model.train_as_vaelp(train_loader, lr=1e-4)