Tutorial - RL Train Cycle Overview

RL Train Cycle Overview

The goal of this tutorial is to walk through the RL fit cycle to familiarize ourselves with the Events cycle and get a better understanding of how Callback and Environment classes work.

Performance Notes

The workflow in this notebook is more CPU-constrained than GPU-constrained due to the need to evaluate samples on CPU. If you have a multi-core machine, it is recommended that you uncomment and run the set_global_pool cells in the notebook. This will trigger the use of multiprocessing, which will result in 2-4x speedups.

This notebook may run slow on Collab due to CPU limitations.

If running on Collab, remember to change the runtime to GPU

High Level Overview

The Environment

At the highest level, we have the Environment class. The Environment holds together several sub-modules and orchestrates them during the fit loop. The following are contained in the Environment:

• agent - This is the actual model we're training
• template_cb - this holds a Template class that we use to define our chemical space
• samplers - samplers generate new samples to train on
• buffer - the buffer collects and distributes samples from all the samplers
• rewards - rewards score samples
• losses - losses generate values we can backpropagate through
• log - the log holds a record of all samples in the training process

Callbacks and the Event Cycle

Each one of the above items is a Callback. A Callback is a a general class that can hook into the Environment fit cycle at a number of pre-defined Events. When the Environment calls a specific Event, the event name is passed to every callback in the Environment. If a given Callback has a defined function named after the event, that function is called. This creates a very flexible system for customizing training loops.

We'll be looking more at Events later. For now, we'll just list them in brief. These are the events called during the RL training cycle in the order they are executed:

• setup - called when the Environment is created, used to set up values
• before_train - called before training is started
• build_buffer - draws samples from samplers into the buffer
• filter_buffer - filters samples in the buffer
• after_build_buffer - called after buffer filtering. Used for cleanup, logging, etc
• before_batch - called before a batch starts, used to set up the batch state
• sample_batch - samples are drawn from sampers and buffer into the batch state
• before_filter_batch - allows preprocessing of samples before filtering
• filter_batch - filters samples in batch state
• after_sample - used for calculating sampling metrics
• before_compute_reward - used to set up any values needed for reward computation
• compute_reward - used by rewards to compute rewards for all samples in the batch state
• after_compute_reward - used for logging reward metrics
• reward_modification - modify rewards in ways not tracked by the log
• after_reward_modification - log reward modification metrics
• get_model_outputs - generate necessary tensors from the model
• after_get_model_outputs - used for any processing required prior to loss calculation
• compute_loss - compute loss values
• zero_grad - zero grad
• before_step - used for computation before optimizer step (ie gradient clipping)
• step - step optimizer
• after_batch - compute batch stats
• after_train - final event after all training batches

Getting Started

We start by creating all the components we need to train a model

Agent

The Agent is the actual model we want to train. For this example, we will use the LSTM_LM_Small_ZINC model, which is a LSTM_LM model trained on a chunk of the ZINC database.

The agent will actually contain two versions of the model. The main model that we will train with every update iteration, and a baseline model which is updated as an exponentially weighted moving average of the main model. Both models are used in the RL training algorithm we will set up later

Template

The Template class is used to conrol the chemical space. We can set parameters on what molecular properties we want to allow. For this example, we set the following:

• Hard Filters - must have qualities
  • ValidityFilter - must be a valid chemical structure
  • SingleCompoundFilter - samples must be single compounds
  • RotBondFilter - compounds can have at most 8 rotatable bonds
  • ChargeFilter - compounds must have no net charge
• Soft Filters - nice to have qualities
  • QEDFilter - Compounds get a score bonus of +1 if their QED value is greater than 0.5
  • SAFilter - compounds get a score bonus of + if their SA score is less than 5

We then pass the Template to the TemplateCallback which integrates the template into the fit loop. Note that we pass prefilter=True to the TemplateCallback, which ensures compounds that don't meet our hard filters are removed from training

Reward

For the reward, we will load a scikit-learn linear regression model. This model was trained to predict affinity against erbB1 using molecular fingerprints as inputs

This score function is extremely simple and likely won't translate well to real affinity. It is used as a lightweight example

We can think of the score function as a black box that takes in samples (SMILES strings) and returns a single numeric score for each sample. Any score function that follows this paradigm can be integrated into MRL

Loss Function

For our loss, we will use the PPO reinforcement learning algorithm. See the PPO paper for full details.

The gist of it is the loss function takes a batch of samples and directs he model to increase the probability of above-average samples (relative to the batch mean) and decrease he probability of below-average samples.

Samplers

Samplers fill the role of generating samples to train on. We will use four samplers for this run:

• sampler1: ModelSampler - this sampler will draw samples from the main model in the Agent. We set buffer_size=1000, which means we will generate 1000 samples every time we build the buffer. We set p_batch=0.5, which means during training, 50% of each batch will be sampled on the fly from the main model and the rest of the batch will come from the buffer
• sampler2: ModelSampler - this sampler is the same as sampler1, but we draw from the baseline model instead of the main model. We set p_batch=0., so this sampler will only contribute to the buffer
• sampler3: LogSampler - this sampler looks through the log of previous samples. Based on our input arguments, it grabs the top 95 percentile of samples in the log, and randomly selects 100 samples from that subset
• sampler4: DatasetSampler - this sampler is seeded wih erbB1 training data used to train the score function. This sampler will randomly select 4 samples from the dataset to add to the buffer

Other Callbacks

We'll add three more callbacks:

• MaxCallback: this will grab the max reward within a batch that came from the source live. live is the name we gave to sampler1 above. This means the max callback will grab all outputs from sampler1 corresponding to samples from the live model and add the largest to the batch metrics
• PercentileCallback: this does the same as MaxCallback but instead of printing the maximum score, it prints the 90th percentile score
• NoveltyReward: this is reward modification that gives a bonus score of 0.05 to new samples (ie samples that haven't appeared before in training)

Training Walkthrough

Now we will step through the training cycle looking at how each callback event is used

Setup

The first event occurs when we create our Environment using the callbacks we set up before. Instantiating the Environment registers all callbacks and runs the setup event. Many callbacks use the setup event to add terms to the batch log or the metrics log.

Inside the environment, we just created a Buffer and a Log.

The Buffer holds a list of samples, which is currently empty

The Log holds a number of containers for tracking training outputs

• metrics: dictionary of batch metrics. Each key maps to a list where each value in the list is the metric term for given batch
• batch_log: dictionary of batch items. Each key maps to a list. Each element in that list is a list containing the batch values for that key in a given batch
• unique_samples: dictionary of unique samples and the rewards for those samples. Useful for looking up if a sample has been seen before
• df: dataframe of unique samples and all associated values stored in the batch_log

We can see that these log terms have already been populated during the setup event

The keys in the above dictionaries were added by the associated callbacks. For example, look at the setup method in ModelSampler, the type of sampler we used for sampler1:

    def setup(self):
        if self.p_batch>0. and self.track:
            log = self.environment.log
            log.add_metric(f'{self.name}_diversity')
            log.add_metric(f'{self.name}_valid')
            log.add_metric(f'{self.name}_rewards')
            log.add_metric(f'{self.name}_new')

We gave sampler1 the name live. As a result, the terms live_diversity, live_valid, live_rewards and live_new were added to the metrics.

We can also look at the setup method of our loss function loss:

    def setup(self):
        if self.track:
            log = self.environment.log
            log.add_metric(self.name)
            log.add_log(self.name)

This is responsible for the PPO terms in the batch_log and the metrics. The PPO metrics term will store the average PPO loss value across a batch, while the PPO batch log term will store the PPO value for each item in a batch

The Fit Cycle

At this point, we could start training using Environment.fit. We could call env.fit(200, 90, 10, 2) to train for 10 batches with a batch size of 200. For this tutorial, we will step through each part of the fit cycle and observe what is happening

Before Train

The first stage of the fit cycle is the before_train stage. This sets the batch size and sequence length based on the inputs to Environment.fit (which we will set manually) and prints the top of the log

Build Buffer

The next stage of the cycle is the build_buffer stage. This consists of the following events:

• build_buffer: samplers add items to the buffer
• filter_buffer: the buffer is filtered
• after_build_buffer: use as needed

Going into this stage, our buffer is empty:

build_buffer

By calling the build_buffer event, our samplers will add items to the buffer

Now we have 2004 items in the buffer.

We can use the buffer_sources attribute to see where each item came from. We have 1000 items from live_buffer which corresponds to sampler1, sampling from the main model.

We have 1000 items from base_buffer which corresponds to sampler2, sampling from the baseline model.

We have 4 items from erbB1_data_buffer, our dataset sampler (sampler4).

Our log sampler, sampler3 was set to start sampling after 10 training iterations, so we don't currently have any samples from that sampler

filter_buffer

It's likely some of these samples don't match our compound requirements defined in the Template we used, so we want to filter the buffer for passing compounds. This is what the filter_buffer does. For this current example, the only callback doing any buffer filtering is the template callback. However, the filter_buffer can be used to implement any form of buffer filtering.

Any callback that passes a list of boolean values to Buffer._filter_buffer can filter the buffer.

After filtering, we have 1829 remaining samples

after_build_buffer

Next is the after_build_buffer event. None of our current callbacks make use of this event, but it exists to allow for evaluation/postprocessing/whatever after buffer creation.

Sample Batch

The next event stage is the sample_batch stage. This consists of the following events:

• before_batch: set up/refresh any required state prior to batch sampling
• sample_batch: draw one batch of samples
• before_filter_batch: evaluate unfiltered batch
• filter_batch: filter batch
• after_sample: compute sample based metrics

before_batch

This event is used to create a new BatchState for the environment. The batch state is a container designed to hold any values required by the batch

Currently the batch state only has placeholder values for commonly generated terms

sample_batch

Now we actually draw samples to form a batch. All of our Sampler objects have a p_batch value, which designated what percentage of the batch should come from that sampler. Batch sampling is designed such that individual sampler p_batch values are respected, and any remaining batch percentage comes from the buffer.

Only sampler1 has p_batch>0., with a value of p_batch=0.5. This means 50% of the batch will be sampled on he fly from sampler1, and the remaining 50% of the batch will come from the buffer.

Using a hybrid of live sampling and buffer sampling seems to work best. That said, it is possible to have every batch be 100% buffer samples (like offline RL), or have 100% be live samples (like online RL)

Now we can see we've populated several terms in the batch state. BatchState.samples now has a list of samples. BatchState.sources has the source of each sample.

We also added BatchState.live_raw and BatchState.base_raw. These terms hold the outputs of sampler1 and sampler2. When we filter BatchState.samples, we can refer to the _raw terms to see what samples were removed.

Note that BatchState.base_raw is an empty list since sampler2.p_batch=0.

BatchState.sources holds the source of each sample. We have 100 samples from live, which corresponds to our on the fly samples from sampler1. The remaining 100 samples come from live_buffer and base_buffer. This means they came from either sampler1 (live) or sampler2 (base) by way of being sampled from the buffer

before_filter_batch

This event is not used by any of our current callbacks. It provides a hook to influence the batch state prior to filtering

filter_batch

Now the batch will be filtered by our Template, as well as any other callbacks with a filter_batch method

We can see that 13 of our 200 samples were removed by filtering

We can compare the values in BatchState.samples and BatchState.live_raw to see what was filtered

after_sample

The after_sample event is used to calculate metrics related to sampling

We can see that several values have been added to Environment.log.metrics

• new: percent of samples that have not been seen before
• diversity: number of unique samples relative to the number of total samples
• bs: true batch size after filtering
• valid: percent of samples that passed filtering
• live_diversity: number of unique samples relative to the number of total samples from sampler1
• live_valid: percent of samples that passed filtering from sampler1
• live_new: percent of samples that have not been seen before from sampler1

Compute Reward

After we sample a batch, we enter the compute_reward stage. This consists of the following events:

• before_compute_reward - used to set up any values needed for reward computation
• compute_reward - used by rewards to compute rewards for all samples in the batch state
• after_compute_reward - used for logging reward metrics
• reward_modification - modify rewards in ways not tracked by the log
• after_reward_modification - log reward modification metrics

before_compute_reward

This event can be used to set up any values needed for reward computation. Most rewards only need the raw samples as inputs, but rewards can use other inputs if needed. The only requirement for a reward is that it returns a tensor with one value per batch item.

By default, the Agent class will tensorize the samples present at this step. Our PPO loss will also add placeholder values for the terms needed by that function

A number of new items have populated the batch state

compute_reward

This step actually computes rewards. The BatchState has a tensor of 0s as a placeholder for reward values. Rewards will compute a numeric score for each item in the batch and add it to BatchState.rewards

So where did these rewards come from?

One reward term comes from our Template. We specified soft rewards for compounds with QED>=0.5 and SA<=5. Compounds could score a maximum of 2 from the template.

We also have the reward from the erbB1 regression model we set up earlier.

The specific rewards from each of these sources are logged in the BatchState

For the Template, we have BatchState.template and BatchState.template_passes

Template scores:

BatchState.template_passes shows which samples passed the hard filters. Since we decided to prefilter with our template earlier, all remaining samples are passing

And here we have the erbB2 regression scores

after_compute_reward

This event is used to calculate metrics on the rewards

reward_modification

The reward modification event can be thought of as a second reward that isn't logged. The reason for including this is to allow for transient, "batch context" rewards that don't affect logged values.

When we set up our callbacks earlier, we had a term

new_cb = NoveltyReward(weight=0.05)

Which would add a bonus score of 0.05 to new, never before seen samples. The point of this callback is to give the model a soft incentive to generate novel samples.

We want this score to impact our current batch. However, if we treated it the same as our actual rewards, the samples would be saved into env.log with their scores inflated by 0.05. Later, when our LogSampler samples from the log, the sampling would be influenced by a score that was only supposed to be given once.

Separating out rewards and reward modifications lets us avoid this

after_reward_modification

Similar to after_compute_reward, this event can be used to compute stats on reward modifications

Get Model Outputs

After computing rewards, we move to set up our loss calculation. The get_model_outputs stage is based on generating the values that we will be backpropagating through. This stage consists of the following events:

• get_model_outputs - generate necessary tensors from the model
• after_get_model_outputs - used for any processing required prior to loss calculation

get_model_outputs

This is where we generate tensor values used for loss computation.

The specifics of what happens here depends on the type of model used. For autoregressive models, this step involves taking the x and y tensors we generated during the before_compute_reward event and doing a forward pass.

x is a tensor of size (bs, sl). Running x through the model will give a set of log probabilities of size (bs, sl, d_vocab). We then use y to gather the relevant log probs to get a gathered log prob tensor of size (bs, sl).

We generate these values from both the main model and the baseline model

after_get_model_outputs

This event is not used by any of our current callbacks, but can be used for any sort of post-processing needed before loss computation

Compute Loss

Now we actually compute a loss value and do an optimizer update. See the PPO class for a description of the policy gradient algorithm used.

Loss computation consists of the following steps:

• compute_loss - compute loss values
• zero_grad - zero grad
• before_step - used for computation before optimizer step (ie gradient clipping)
• step - step optimizer

compute_loss

When we first created our BatchState, there was a placehoder value for loss. This is the value that will ulimately be backpropagated through. This means we can run any sort of loss configuration, so long as the final values end up in BatchState.loss.

For example, the PPO policy gradient algorithm we are using involved a ValueHead that predicts values at every time step. This model is held in the PolicyLoss callback that holds the PPO class. During the compute_loss event, PPO computes an additional loss for the value head that is added to BatchState.loss. PolicyLoss also holds an optimizer for the ValueHead parameters.

zero_grad

This is an event to zero gradients of all optimizers in play. We currently have one optimizer in Agent for our generative model and one in PolicyLoss for the ValueHead of our policy gradient algorithm.

before_step

This is an event before the actual optimizer step. This is used for things like gradient clipping

step

This is the actual optimizer step. This will step both the Agent and PolicyLoss optimizers

After Batch

The after_batch stage consists of a single after_batch event. This is used for any updates at the end of the batch.

In particular, the Log will update Log.df and the Agent will update he baseline model

After Train

The after_train event can be used to calculate any final statistics or other values as desired

Conclusions

Hopefully walking through the training process step by step has made he process more understandable. We conclude by simply running Environment.fit so we don't have to go through things step by step anymore