Getting Started - Training a Conformer Generation Agent
This section walks through how to write a short Python script for training a deep reinforcement learning model for generating low-energy conformers of a molecule. The main components needed for training a model will be covered in this section. Additional options such as hyperparamters and custom environments will be covered in the next sections of the tutorial.
Creating a Training Script
The full code for this example can be found in examples/example1.py.
Configuring the Environment
In reinforcement learning, the environment simulates the task that the agent will be trained on - in this case the task of conformer generation. conformer_rl contains implementations of several general environments (see environments in the API reference for more info) that can be initialized to simulate the task of conformer generation for a specific molecule. This initialization requires a MolConfig object to be passed in, which contains information specific to the molecule of interest, such as the molecule structure and normalizing constants specific to that molecule.
For convenience, conformer_rl has functions for generating a MolConfig object (as well as calculating normalizing constants, if necessary), given a molecule input. The library supports several formats, including rdkit molecule (rdkit.rdchem.mol) objects, SMILES strings, and MOL files (see Molecule Config Generators in the API reference). Also included are functions for generating rdkit mol objects for certain classes of molecules and polymers, such as alkanes and lignins (see Molecule Generation in the API reference).
A Note on Normalizing Constants
When using an environment that utilizes Gibbs Score reward 1 (or equivalently, the Boltzmann Factor reward), the normalizing constants \(E_0\) and \(Z_0\) need to be provided for the molecule. The functions in Molecule Config Generators include options for generating these normalizing constants automatically using the function calculate_normalizers(). Note that an agent can be still trained if these constants are set to arbitrary values, but the rewards may not be numerically stable which might affect agent performance. Additionally, when training multiple models on the same molecule to assess model performance, the same set of normalizing constants should be used in each run and normalizing constants should not be re-generated in each run (since the constants generated by calculate_normalizers() are not deterministic).
Example Environment Configuration
Suppose we want to generate conformers for a branched alkane molecule with 14 carbon atoms. We will first generate an rdkit.rdchem.mol object for the molecule. Using the function generate_branched_alkane(), we generate our rdkit mol object:
mol = generate_brancehed_alkane(14)
Next, we can use the function config_from_rdkit():
mol_config = config_from_rdkit(mol, num_conformers=200, calc_normalizers=True, save_file='alkane')
which will create a MolConfig for our branched alkane. The num_conformers parameter specifies the number of conformers we want to generate in each environment episode, in this case 200. By setting calc_normalizeres=True, the function will calculate the normalizing constants which will be later used by the environment for calculating rewards. We also set save_file='alkane', so that the generated MolConfig object is dumped as a binary Pickle file named alkane.pkl, so that it can be reused later.
There are two main benefits for saving the generated MolConfig object. Firstly, the normalizing constants generated by setting calc_normalizers=True are not deterministic and relies on rdkit’s conformer generation functionality (which uses random initialization). As discussed above, if we wish to compare the performance of two separate models on the same environment, the same set of normalizing constants should be used for both models, and a new set of normalizing constants should not be generated. Secondly, the generation of normalizing constants can be time consuming for large molecules, and it is therefore unnecessary to re-generate these constants for the same molecule for multiple experiments.
To use the saved MolConfig object instead of generating one, simply set mol_config as follows:
with open('alkane.pkl', 'rb') as file:
mol_config = pickle.load(file)
Custom Molecules
If you have prepared your own molecule for conformer generation, and it is not in a rdkit mol format, conformer_rl also has functions to create MolConfig for other formats. For example, if your molecule can be expressed as a SMILES string, you can use the config_from_smiles() function, such as in the following example:
mol_config = config_from_smiles('CC(CCC)CC', num_conformers=200, calc_normalizers=True, save_file='alkane')
The molecule can also be in the form of a MOL file, in which the function config_from_molFile() can be used:
mol_config = config_from_molFile('name_of_mol_file.mol', num_conformers=200, calc_normalizers=True, save_file='alkane')
Configuring the Agent
The next step is to configure the training of the agent itself. Similar to the environment, an agent is configured by initializing it with a Config object. There are many things that can be configured for the training of the agent, such as the neural network used, training hyperparameters, logging options, and more. Many of these options have default values found in Config, and we will worry about configuring those in Hyperparameter Tuning and Additional Options. For now, we will focus on the main requirements. We first construct a config object:
config = Config()
config.tag = 'example1'
The tag attribute, along with the time the training script is run, is used as identifiers for log files and data output files for the training script.
Training Environment
Next, we will set the training environment for the agent:
config.train_env = Task('GibbsScorePruningEnv-v0', concurrency=True, num_envs=5, seed=np.random.randint(0,1e5), mol_config=mol_config)
Task() is a function that generates an environment wrapper compatible with the agent. Its main functionality is to generate multiple environments that the agent can interact with concurrently, which speeds up training if there are multiple CPU cores available.
The first parameter, 'GibbsScorePruningEnv-v0', specifies the name of the environment implementation to be used. In this case it represents the class GibbsScorePruningEnv, which has empirically produced good results for several organic molecules. To learn more about how environments are registered and how to create custom environments, see Customizing Environment - Part One and Customizing Environment - Part Two.
We set concurrency=True to utilize multithreading across each of the parallel environments during training. The num_envs parameter specifies the number of environments to be run in parallel. Next we pass in the MolConfig object we created earlier by setting mol_config=mol_config to specify molecule specific parameters when initiating the environments.
Evaluation Environment
Optionally, we can specify an evaluation environment, which is an environment in which the agent will be periodically evaluated on throughout training. This is useful for a number of reasons. Firstly, if the evaluation environment is different from the training environment, we can use the evaluation environment to see how well the agent generalizes to other environments and assess whether the agent is overfitting to the training environment task. Additionally, the training framework implemented in conformer_rl will automatically save the molecules generated by the evaluation environment when evaluating the model, which can be used for analysis in downstream tasks.
For simplicity, we will simply use the same molecule config for the evaluation environment in this example. We specify the evaluation environment in a similar way as the training environment, except that we do not require parallel environments so we use the default values for the concurrency and num_envs parameters:
config.eval_env = Task('GibbsScorePruningEnv-v0', seed=np.random.randint(0,7e4), mol_config=mol_config)
config.eval_episodes=10000
config.eval_episodes specifies how often (in number of episodes) the agent should be evaluated on the evaluation environment. If this is set to 0, the agent will not be evaluated on the evaluation environment.
Training the Agent
Finally we train the agent:
agent = PPORecurrentAgent(config)
agent.run_steps()
PPORecurrentAgent specifies the reinforcement learning algorithm used for training the agent (PPO stands for the ‘proximal policy optimization’ algorithm, and ‘recurrent’ refers to the fact that the agent is designed to be compatible with recurrent neural networks). This algorithm has empirically performed well in the conformer generation task. conformer_rl also includes implementations of other agents, which can be found in Pre-built agents. By calling the run_steps() method, the agent will be trained.
Logs and Output
By default, all logs will be stored in a directory called data from which the training script is run. This can be changed by modifying the data_dir attribute of the Config object.
Tensorboard Logs
During training, several values will be periodically logged and can be viewed using Tensorboard. The list of values that are logged can be found in the API reference for each agent in Pre-built agents. Tensorboard logs will be stored in a subdirectory called tensorboard_log. To view the logs using tensorboard, run:
tensorboard --logdir data/tensorboard_log
All logs will be labeled using the tag attributes of the Config object, as well as the time the training script is run.
Agent Evaluation Logs
If the model is evaluated on an evaluation environment (by setting config.eval_episodes to a non-zero value), data such as the molecules generated and the energies of the molecules will be saved during each evaluation on the evaluation environment. Like the Tensorboard logs, these will be saved in the data directory by default, in a subdirectory called env_data. Within this directory, the logged data is stored in a directory with name corresponding to the tag specified and time the training script was run. After running example1.py, for example, the directory would look something like example1_02-12-2021_20:49:18. Within this directory, there will be a directory for each time the agent was evaluated, specifying the number of training iterations that has passed before that evaluation. Within each of these directories, there is a directory for each episode the agent was evaluated. Within each of these directories is a .pickle file which contains logged values such as the energy and reward for each conformer generated. Additionally, there is a MOL file containing the 3-D representation of the molecule conformer generated at each step within the episode. The data in the .pickle files can be analyzed/visualized using the analysis module. An example of using the analysis module can be found in
examples/example_analysis.ipynb.