summaries: This paper adopts the framework of unsupervised skill learning. To solve the problem that the discriminator will have low confidence in the unseen data thus providing a low intrinsic reward, the paper derives an information gain auxiliary objective that involves training an ensemble of discriminators and rewarding the policy for their disagreement. The paper conducts extensive experiments on tabular grid world and 57 games of the Atari Suite.

summaries: The paper proposes a novel unsupervised skill discovery algorithm. Beginning with the family of such methods (DIAYN, VIC, etc) that learn a discriminator to distinguish skills given some observations of the trajectory and a policy that executes a skill conditioned on the (discrete) skill random variable Z, the premise of the paper is that such methods on their own will fail when new states are encountered during the skill learning process as the discriminator would not have had sufficient data to learn to distinguish novel states. The paper proposes a novel reward bonus that works in addition to a base method such as DIAYN (DIversity is All You Need), such that this bonus “reimburses” the policy for visiting states where the discriminator uncertainty is high (measured using a form of disagreement across an ensemble of discriminators). Experiments on the pedagogical 4-room environment and the Atari suite of environments demonstrates the benefit of the proposed reward bonus in not only learning more skills (or “empowerment” in the VIC nomenclature), but the learnt skills are also superior for downstream tasks (external reward) and lifetime state coverage.

summaries: This paper identifies a source of pessimism in DIAYN-style methods for exploring new parts of the state space. They argue that this issue is due to using a single point estimator as a discriminator and that capturing the epistemic uncertainty of the discriminator could serve as an additional signal to guide exploration. They achieve this by using an ensemble of discriminators and incorporating the epistemic uncertainty across the ensemble into an additional intrinsic reward to the diversity of skills reward (through a mixing parameter $\lambda$). They examine this method against DIAYN-style methods and count-based methods and show that this new approach broadly outperforms both these classes of methods.

summaries: This paper is concerned with unsupervised RL where an extrinsic reward signal is not available. The objective is for the agent to master the environment by exploring it while learning a diverse set of skills. This is done by simultaneously training a policy conditioned on a latent variable and a discriminator that tries to infer the latent variable from trajectories. The authors identify that the intrinsic rewards used in skill discovery result in the agent being pessimistic towards exploring novel parts of the environment. To alleviate this, the authors propose a new auxiliary objective, which results in a bonus based on the disagreement of an ensemble of discriminators. Empirical results on the grid world and Atari show improvements in skill discovery and solving downstream tasks compared to baselines.

abstractText: Unsupervised skill learning objectives (Eysenbach et al., 2019; Gregor et al., 2016) allow agents to learn rich repertoires of behavior in the absence of extrinsic rewards. They work by simultaneously training a policy to produce distinguishable latent-conditioned trajectories, and a discriminator to evaluate distinguishability by trying to infer latents from trajectories. The hope is for the agent to explore and master the environment by encouraging each skill (latent) to reliably reach different states. However, an inherent exploration problem lingers: when a novel state is actually encountered, the discriminator will necessarily not have seen enough training data to produce accurate and confident skill classifications, leading to low intrinsic reward for the agent and effective penalization of the sort of exploration needed to actually maximize the objective. To combat this inherent pessimism towards exploration, we derive an information gain auxiliary objective that involves training an ensemble of discriminators and rewarding the policy for their disagreement. Our objective directly estimates the epistemic uncertainty that comes from the discriminator not having seen enough training examples, thus providing an intrinsic reward more tailored to the true objective compared to pseudocount-based methods (Burda et al., 2019). We call this exploration bonus discriminator disagreement intrinsic reward, or DISDAIN. We demonstrate empirically that DISDAIN improves skill learning both in a tabular grid world (Four Rooms) and the 57 games of the Atari Suite (from pixels). Thus, we encourage researchers to treat pessimism with DISDAIN.

authors:
Joshua Achiam
Harrison Edwards
Dario Amodei
Pieter Abbeel

title: Variational option discovery algorithms

venue: arXiv preprint arXiv:1807.10299,

authors:
Rishabh Agarwal
Max Schwarzer
Pablo Samuel Castro
Aaron Courville
Marc G Bellemare

title: Deep reinforcement learning at the edge of the statistical precipice

venue: In Neural Information Processing Systems (NeurIPS),

authors:
Ankesh Anand
Evan Racah
Sherjil Ozair
Yoshua Bengio
Marc-Alexandre Côté
R Devon Hjelm

title: Unsupervised state representation learning in Atari

authors:
David Barber
Felix Agakov

title: Information maximization in noisy channels: A variational approach

venue: In Neural Information Processing Systems (NIPS),

authors:
Kate Baumli
David Warde-Farley
Steven Hansen
Volodymyr Mnih

title: Relative variational intrinsic control

venue: In AAAI Conference on Artificial Intelligence,

authors:
M.G. Bellemare
Y. Naddaf
J. Veness
M. Bowling

title: The Arcade Learning Environment: An evaluation platform for general agents

venue: Journal of Artificial Intelligence Research, 47:253–279,

authors:
Marc Bellemare
Sriram Srinivasan
Georg Ostrovski
Tom Schaul
David Saxton
Remi Munos

title: Unifying count-based exploration and intrinsic motivation

authors:
Ronen I Brafman
Moshe Tennenholtz

title: R-max - a general polynomial time algorithm for nearoptimal reinforcement learning

venue: Journal of Machine Learning Research (JMLR),

title: Language models are few-shot learners

authors:
Yuri Burda
Harrison Edwards
Amos Storkey
Oleg Klimov

title: Exploration by random network distillation

venue: In International Conference on Learning Representations (ICLR),

authors:
Richard Y Chen
Szymon Sidor
Pieter Abbeel
John Schulman

title: UCB exploration via Qensembles

venue: arXiv preprint arXiv:1706.01502,

authors:
Ting Chen
Simon Kornblith
Mohammad Norouzi
Geoffrey Hinton

title: A simple framework for contrastive learning of visual representations

venue: In International Conference on Machine Learning (ICML),

authors:
Ting Chen
Simon Kornblith
Kevin Swersky
Mohammad Norouzi
Geoffrey E Hinton

title: Big self-supervised models are strong semi-supervised learners

authors:
Lasse Espeholt
Hubert Soyer
Remi Munos
Karen Simonyan
Vlad Mnih
Tom Ward
Yotam Doron
Vlad Firoiu
Tim Harley
Iain Dunning
Shane Legg
Koray Kavukcuoglu

title: IMPALA: Scalable distributed deep-RL with importance weighted actor-learner architectures

authors:
Benjamin Eysenbach
Abhishek Gupta
Julian Ibarz
Sergey Levine

title: Diversity is all you need: Learning skills without a reward function

authors:
Sebastian Flennerhag
Jane X Wang
Pablo Sprechmann
Francesco Visin
Alexandre Galashov
Steven Kapturowski
Diana L Borsa
Nicolas Heess
Andre Barreto
Razvan Pascanu

title: Temporal difference uncertainties as a signal for exploration

venue: arXiv preprint arXiv:2010.02255,

authors:
Karol Gregor
Danilo Jimenez Rezende
Daan Wierstra

title: Variational intrinsic control

venue: arXiv preprint arXiv:1611.07507,

authors:
Steven Hansen
Will Dabney
Andre Barreto
David Warde-Farley
Tom Van de Wiele
Volodymyr Mnih

title: Fast task inference with variational intrinsic successor features

authors:
Chi Jin
Akshay Krishnamurthy
Max Simchowitz
Tiancheng Yu

title: Reward-free exploration for reinforcement learning

authors:
Steven Kapturowski
Georg Ostrovski
Will Dabney
John Quan
Remi Munos

title: Recurrent experience replay in distributed reinforcement learning

authors:
Saurabh Kumar
Aviral Kumar
Sergey Levine
Chelsea Finn

title: One solution is not all you need: Few-shot extrapolation via structured maxent RL

authors:
Balaji Lakshminarayanan
Alexander Pritzel
Charles Blundell

title: Simple and scalable predictive uncertainty estimation using deep ensembles

authors:
Anuj Mahajan
Tabish Rashid
Mikayel Samvelyan
Shimon Whiteson

title: MAVEN: Multi-agent variational exploration

authors:
OpenAI
Ilge Akkaya
Marcin Andrychowicz
Maciek Chociej
Mateusz Litwin
Bob McGrew
Arthur Petron
Alex Paino
Matthias Plappert
Glenn Powell
Raphael Ribas
Jonas Schneider
Nikolas Tezak
Jerry Tworek
Peter Welinder
Lilian Weng
Qiming Yuan
Wojciech Zaremba
Lei Zhang

title: Solving Rubik’s Cube with a robot hand

venue: arXiv preprint arXiv:1910.07113,

authors:
Ian Osband
Charles Blundell
Alexander Pritzel
Benjamin Van Roy

title: Deep exploration via bootstrapped DQN

authors:
Ian Osband
John Aslanides
Albin Cassirer

title: Randomized prior functions for deep reinforcement learning

authors:
Ian Osband
Benjamin Van Roy
Daniel J. Russo
Zheng Wen

title: Deep exploration via randomized value functions

authors:
Deepak Pathak
Dhiraj Gandhi
Abhinav Gupta

title: Self-supervised exploration via disagreement

authors:
Jing Peng
Ronald J Williams

title: Incremental multi-step Q-learning

authors:
Alec Radford
Jeff Wu
Rewon Child
David Luan
Dario Amodei
Ilya Sutskever

venue: Technical report,

authors:
Roberta Raileanu
Tim Rocktäschel

title: RIDE: Rewarding impact-driven exploration for procedurallygenerated environments

authors:
Ramanan Sekar
Oleh Rybkin
Kostas Daniilidis
Pieter Abbeel
Danijar Hafner
Deepak Pathak

title: Planning to explore via self-supervised world models

authors:
Younggyo Seo
Lili Chen
Jinwoo Shin
Honglak Lee
Pieter Abbeel
Kimin Lee

title: State entropy maximization with random encoders for efficient exploration

venue: arXiv preprint arXiv:2102.09430,

authors:
H. Sebastian Seung
Manfred Opper
Haim Sompolinsky

title: Query by committee

venue: In Conference on Learning Theory (COLT),

authors:
Pranav Shyam
Wojciech Jaśkowski
Faustino Gomez

title: Model-based active exploration

authors:
Richard S. Sutton
Doina Precup
Satinder Singh

title: Between MDPs and semi-MDPs: A framework for temporal abstraction in reinforcement learning

venue: Artificial Intelligence,

authors:
Yuhuai Wu
Roman Ring
Dani Yogatama
Dario Wünsch
Katrina McKinney
Oliver Smith
Tom Schaul
Timothy Lillicrap
Koray Kavukcuoglu
Demis Hassabis
Chris Apps
David Silver

title: Grandmaster level in StarCraft II using multi-agent reinforcement learning

authors:
David Warde-Farley
Tom Van de Wiele
Tejas Kulkarni
Catalin Ionescu
Steven Hansen
Volodymyr Mnih

title: Unsupervised control through non-parametric discriminative rewards

authors:
Chuheng Zhang
Yuanying Cai
Longbo Huang Jian Li

title: Exploration by maximizing Rényi entropy for reward-free RL framework

authors:
Yunzhi Zhang
Pieter Abbeel
Lerrel Pinto

title: Automatic curriculum learning through value disagreement

heading: 1 INTRODUCTION

text:
Reinforcement learning (RL) has proven itself capable of learning useful skills when clear taskspecific rewards are available (OpenAI et al., 2019a;b; Vinyals et al., 2019). However, truly intelligent agents should, like humans, be able to learn even in the absence of supervision in order to acquire repurposable task-agnostic knowledge. Such unsupervised pre-training has seen recent success in language (Radford et al., 2019; Brown et al., 2020) and vision (Chen et al., 2020a;b), but its potential has yet to be fully realized in the learning of behavior.
The most promising class of algorithms for unsupervised skill discovery is based upon maximizing the discriminability of skills represented by latent variables on which a policy is conditioned. Objectives are typically derived from variational approximations to the mutual information between latent variables and states visited (Gregor et al., 2016; Eysenbach et al., 2019; Warde-Farley et al., 2019; Hansen et al., 2020; Baumli et al., 2021), employing a learned parametric skill discriminator. Both the policy and discriminator have the objective of strong predictive performance by the discriminator, though the discriminator is trained with supervised learning rather than RL. The end result is a policy capable of producing consistently distinguishable behavioral motifs, or “skills.” These skills can be evaluated zero-shot, fine-tuned, or composed in a hierarchical RL setup to maximize task reward when one is introduced (Eysenbach et al., 2019; Hansen et al., 2020). Unsupervised skill learning objectives can also be maximized in conjunction with task reward, in order to promote robustness of learned behavior to environment perturbations (Mahajan et al., 2019; Kumar et al., 2020).
∗equal contribution
The degree to which unsupervised skill discovery methods are useful in such downstream applications depends on how many skills they are able to learn. Indeed, Eysenbach et al. (2019) showed that as more skills are learned, the performance obtained using the learned skills in a hierarchical reinforcement learning setup improves. In follow up work, Achiam et al. (2018) showed that methods like DIAYN can struggle to learn large numbers of skills and proposed gradually increasing the number of skills according to a curriculum to make skill discovery easier.
The aim of this work is to improve the ability of skill discovery methods to learn more skills. We highlight an exploration problem intrinsic to the entire class of variational skill discovery algorithms which can inhibit discovery of new skills. During skill learning, the policy will necessarily need to explore new states of the environment. The discriminator must then make latent predictions for states it has never seen before, resulting in incorrect and/or low-confidence predictions. The policy will in turn be penalized for this poor discriminator performance, and discouraged from seeking out new states. We refer to this problem as “pessimistic exploration” and describe it further in section 2.
To motivate our solution, we argue that it is important for the policy to distinguish between two kinds of uncertainty in the discriminator - aleatoric uncertainty that comes from the policy producing similar trajectories for different skills, and epistemic uncertainty that comes from a lack of training data. The former indicates poor policy performance, but the latter is in fact desirable, and serves as a signal for potential discriminator learning. Unfortunately, skill discovery algorithms treat both types of uncertainty the same and thus ignore this important signal. We propose to capture it.
Our primary contribution, which we present in section 3, is an exploration bonus tailored to skill discovery algorithms, designed to overcome pessimistic exploration. We identify states of high epistemic uncertainty in the discriminator by training an ensemble of discriminators and measuring their disagreement. We call this exploration bonus discriminator disagreement intrinsic reward, or DISDAIN. Intuitively, the ensemble members may disagree in novel states, but must come to agree in frequently visited ones. More formally, we derive DISDAIN from a Bayesian perspective that 1) represents the posterior over discriminator parameters using an ensemble, and 2) encourages the policy to maximize information gain (i.e. reduce uncertainty) about the parameters of the discriminator. In section 4, we demonstrate empirically that DISDAIN improves skill learning over unbonused skill discovery algorithms in both an illustrative grid world (Four Rooms, Sutton et al. (1999)) and the Atari 2600 learning environment (ALE, Bellemare et al. (2013)) more so than augmenting with popular exploration bonuses not tailored to skill discovery.

heading: 2 UNSUPERVISED SKILL LEARNING THROUGH VARIATIONAL INFOMAX

heading: 2.1 INTRODUCTION

text:
We now formalize our setting of interest: unsupervised skill learning through variational information maximization (Gregor et al., 2016; Eysenbach et al., 2019; Warde-Farley et al., 2019; Hansen et al., 2020; Baumli et al., 2021). We consider a Markov decision process (MDP), defined by the tuple M = (S,A, pE, ρ, r, γ), where S and A are state and action spaces, respectively, the environment dynamics pE(s′ | s, a) specifies the probability of transitioning to state s′ ∈ S when taking action a ∈ A in state s ∈ S , ρ(s) denotes the probability of starting an episode in state s, and γ ∈ [0, 1) is a discount factor. Since we focus on unsupervised training, we ignore the environmental reward r.
Our agents seek to learn a repertoire of skills, indexed by the latent variable Z and represented by the policy πθ(a | s, z), which is parameterized by θ and maps from states and latent variables to distributions over actions. The latent variables are sampled z ∼ p(Z) at the beginning of each trajectory and then fixed, so each z represents a temporally extended behavior. The skill trajectory length T may differ from the episode length, thus a new skill might be resampled within an episode.
For conciseness, we will denote trajectories sampled from the policy by τ ∼ π(z) when conditioning on a particular skill z, and τ ∼ π when collecting trajectories across skills. To simplify our discussion, and because it is the most common case in practice, we will assume that Z is categorical with cardinality NZ , although much of the discussion carries over to continuous Z.
A large and growing class of objectives for unsupervised skill discovery are derived from maximizing the mutual information between the latent skill Z and some feature of the resulting trajectories O(τ):
F(θ) ≡ I(Z,O) = H(Z)−H(Z | O) = Ez∼p(z),τ∼π(z)[log p(z | o(τ))− log p(z)] . (1)
For example, variational intrinsic control (VIC, Gregor et al. (2016)) maximizes I(Z;S0, ST ) - the mutual information between the skill and initial and final states (oT = (s0, sT )).1 Diversity is all you need (DIAYN, Eysenbach et al. (2019)), on the other hand, maximizes I(Z, S) - the mutual information between the skill and each state along the trajectory (ot = st for t ∈ 1 : T ). Intuitively, VIC produces skills that vary in the destination reached (without regard for the path taken), while DIAYN produces skills that vary in the path taken (with less emphasis on the destination reached).
In practice, maximizing equation 1 is not straightforward, because it requires calculating the conditional distribution p(Z | O). In general, it is necessary to estimate it with a learned parametric model qφ(Z | O). We refer to this model as the discriminator, since it is trained to discriminate between skills. Fortunately, replacing p with q still yields a lower bound on F(θ) (Barber and Agakov, 2004), and we may instead maximize the proxy objective F̃(θ):
F(θ) ≥ F̃(θ) = Ez∼p(z),τ∼π(z)[log qφ(z | o(τ))− log p(z)] . (2)
Optimizing F̃(θ) with respect to the policy parameters θ corresponds to RL on the reward:
rskill = log qφ(z | o)− log p(z) . (3)
Since the intent is for the agent to learn a full repertoire of skills, the skill prior p(z) is typically fixed to be uniform (Eysenbach et al., 2019; Achiam et al., 2018; Baumli et al., 2021), in which case − log p(z) = logNZ . If the discriminator simply ignores the trajectory and guesses skills uniformly as well, then log qφ(z | o) = − logNZ and the reward will be zero. If the discriminator instead guesses perfectly, then log qφ(z | o) = 0, and the reward will be logNZ . More generally, the expected reward is an estimate of the logarithm of the effective number of skills. Thus, measuring the reward in bits (i.e. using log2 in equation 3), we can estimate the number of skills learnt as:
nskills = 2 E[rskill]. (4)
We adopt this quantity as our primary performance metric for our experiments.
To make sure the bound in equation 2 is as tight as possible, the discriminator qφ(Z | O) must also be fit to its target p(Z | O) through supervised learning (SL) on the negative log likelihood loss:
L(φ) ≡ −Ez∼p(z),τ∼π(z)[log qφ(z | o(τ))] . (5)
The loss is minimized, and the bound in equation 2 tight, when qφ(Z | O) = p(Z | O).
The joint optimization of F̃(θ) by RL and L(φ) by SL forms a cooperative communication game between policy and discriminator. The agent samples a skill z ∼ p(Z) and generates the “message” τ ∼ π(z). The discriminator receives the message τ and attempts to decode the original skill z. When the policy produces trajectories for different skills that do not overlap in the features O(τ), the discriminator will easily learn to label trajectories, and when the discriminator makes accurate and confident predictions, the reward in equation 3 will be high. Ideally, the end result is a policy exhibiting a maximally diverse set of skills. This joint-training system is depicted in Figure 2a.

heading: 2.2 PESSIMISTIC EXPLORATION IN UNSUPERVISED SKILL DISCOVERY

text:
A conflict arises between the exploration necessary for skill diversification and rewards supplied by an imperfect discriminator, trained only on past policy-generated experience. Without data (and in the absence of perfect generalization), the discriminator is likely to make poor predictions when presented with trajectories containing previously unseen states, resulting in low reward for the policy. Importantly, this penalization occurs regardless of whether the policy produces distinguishable skill-conditioned trajectories if the current discriminator is a locally poor approximation to p(Z | O) for a region of state space represented in O. We note that this is distinct from issues of pessimism in exploration that arise more generally, including with stationary reward functions (Osband et al., 2019), wherein naive exploration strategies fail to adequately position an agent for further acquisition of information. In the scenario we examine here, the agent’s sole source of supervision directly sabotages the learning process when the discriminator extrapolates poorly.
1More accurately, Gregor et al. (2016) conditioned on initial state and used I(Z, ST | S0) = H(Z | S0)− H(Z | S0, ST ). However, it has subsequently become common not to condition the skill sampling distribution on s0 (Eysenbach et al., 2019), in which case H(Z | S0) = H(Z) and I(Z, ST | S0) = I(Z;S0, ST ).
Aleatoric uncertainty is created by indistinguishable states for different skills.
Epistemic uncertainty is created by exploring novel states. Buffer of previously seen states.
(a) Past trajectories.
τ ∼ π (z1) τ ∼ π (z2)
q(Z | τ) = z1 vs z2? (b) Aleatoric uncertainty arises when different skills visit similar states.
τ ∼ π (z1)
q(Z | τ) = z1 vs z2? (c) Epistemic uncertainty arises from exploring novel states.
Figure 1: The pessimistic exploration problem. Because skill discovery objectives do not distinguish between aleatoric and epistemic uncertainty, they penalize exploration.
We argue that in order to overcome this inherent pessimism, we must distinguish between two kinds of uncertainty in the discriminator: aleatoric uncertainty (figure 1b) that is due to the policy producing overlapping skills (i.e. high H(Z | O)), and epistemic uncertainty (figure 1c) that is due to a lack of training data (i.e. poor match between qφ(Z | O) and p(Z | O)). Naively, both kinds of uncertainty contribute to low reward for the policy (equation 3), but while reduction of aleatoric uncertainty requires changes by the policy, epistemic uncertainty can be reduced (and thus reward increased) simply by creating more data (i.e. visiting the same states again). Thus, we argue that we should “reimburse” the policy for epistemic uncertainty in the discriminator, and in fact encourage the policy to visits states of high epistemic uncertainty.
Put another way, when the policy only maximizes the reward of equation 3, it maximizes the lower bound F̃(θ) ≤ F(θ) without regard for its tightness, despite that a looser bound means a more pessimistic reward. The job of keeping the bound tight is left entirely to the discriminator (equation 5). By incentivizing the policy to visit states of high epistemic uncertainty, valuable training data is provided to the discriminator that allows it to better approximate its target and close the gap between F̃(θ) and F(θ). In this way, we encourage the policy to help keep the bound F̃(θ) ≤ F(θ) tight. In the next section, we formalize the intuitions outlined in the last two paragraphs. The result is an exploration bonus measured as the disagreement among an ensemble of discriminators.

heading: 3 DISDAIN: DISCRIMINATOR DISAGREEMENT INTRINSIC REWARD

text:
To formalize the notion of discriminator uncertainty, we take a Bayesian approach and replace the point estimate of discriminator parameters φ with a posterior p(φ). Maintaining a posterior allows us to quantify the information gained about those parameters. Specifically, we will incentivize the policy to produce trajectories in which observing the paired skill label z provides maximal information about the discriminator parameters φ:
I(Z; Φ | O) = H(Z | O)−H(Z | O,Φ) . (6)
Rewriting the entropies as expectations over trajectories, we have: I(Z; Φ | O) = Eτ∼π [ H [∫ p(φ) qφ(Z | o(τ)) dφ ] − ∫ p(φ)H[qφ(Z | o(τ))] dφ ] (7)
where we have used that p(φ) does not depend on the present trajectory and so p(φ | o(τ)) = p(φ). Note that unlike in equations 2 and 5, the expectation over trajectories is not skill-conditioned. The marginalization over Z happens in the entropy and is over the discriminator’s posterior qφ(z | s) rather than the agent’s prior p(z). This is because the discriminator parameters Φ are now part of the probabilistic model and do not just enter through a variational approximation.
How should we represent the posterior over discriminator parameters p(φ)? There is considerable work in Bayesian deep learning that offers possible answers, but here we take an ensemble approach (Seung et al., 1992; Lakshminarayanan et al., 2017). We train N discriminators with parameters φi for the ith discriminator. The discriminators are independently initialized and in theory could also be trained on different mini-batches, though in practice, we found it both sufficient and
simpler to train them on the same mini-batches, as others have also found (Osband et al., 2016). The posterior p(φ) is then represented as a mixture of point masses at the φi:
p(φ) = 1
N N∑ i=1 δ(φ− φi) . (8)
Substituting the posterior in equation 8 into equation 7, we have:
I(Z; Φ | O) = Eτ∼π [ H [ 1
N N∑ i=1 qφi(Z | o(τ)) ] − 1 N N∑ i=1 H[qφi(Z | o(τ))] ] . (9)
Maximizing equation 9 with RL corresponds to adding the following auxiliary reward to the policy:
rDISDAIN(t) ≡ H
[ 1
N N∑ i=1 qφi(Z | ot) ] − 1 N N∑ i=1 H[qφi(Z | ot)] . (10)
In words, this is the entropy of the mean discriminator minus the mean of the entropies of the discriminators. By Jensen’s inequality, entropy increases under averaging, and thus rDISDAIN ≥ 0. For trajectories on which there has been ample training data for the discriminators, the ensemble members should agree and qφi(z | o) ≈ qφj (z | o) for all i, j. Thus the two terms in equation 10 will be equal and this reward will vanish. For states of high discriminator disagreement, however, this reward will be positive, encouraging exploration. Therefore, we call equation 10 the Discriminator Disagreement Intrinsic reward, or DISDAIN (figure 2b).
DISDAIN is simple to calculate for discrete Z, as is common in the relevant literature (Gregor et al., 2016; Eysenbach et al., 2019; Achiam et al., 2018; Baumli et al., 2021). DISDAIN augments any discriminator-based unsupervised skill learning algorithm with two changes. First, an ensemble of discriminators is trained instead of just one, and rskill should be calculated using the ensembleaveraged prediction qφ(Z | O) = 1N ∑N i=1 qφi(Z | O). Second, the DISDAIN reward is combined with rskill, which we do through simple addition with a tunable multiplier λ. Pseudocode for DISDAIN is provided in Algorithm 1. With N = 1 and λ = 0, this is standard unbonused skill discovery.

heading: 4 EXPERIMENTS

text:
We validate DISDAIN by testing its ability to increase skill learning in an illustrative grid world (Four Rooms) as well as a more challenging pixel-based setting requiring function approximation (the 57 Atari games of the Arcade Learning Environment (Bellemare et al., 2013)). In addition to comparing performance to unbonused skill learning, we also compare to using popular off-the-shelf exploration bonuses that are not tailored to skill discovery. In Four Rooms, we compare to using count-based bonuses, which are known to perform well in these settings (Brafman and Tennenholtz, 2002). In Atari, where count-based bonuses become untenable due to the enormous state space, we compare to random network distillation (RND; Burda et al. (2019)), one of the most commonly used
Algorithm 1: Skill discovery with DISDAIN
Input: policy πθ, discriminator ensemble {qφi} N i=1, skill features O(τ), skill distribution p(Z), skill trajectory length T , DISDAIN reward weight λ while not converged do
Reset environment, sampling initial state s0 while episode not ended do
Sample skill, z ∼ p(Z) Sample trajectory of length T from s0, τ ∼ π(z) Form average discriminator from ensemble, qφ = 1N ∑N i=1 qφi rskill = log qφ(z | O(τ))− log p(z) rDISDAIN = H[qφ(· | O(τ))]− 1N ∑N i=1H[qφi(· | O(τ))] r = rskill + λrDISDAIN Update θ with RL to maximize r Update {qφi} N i=1 with SL to maximize log qφi(z | O(τ)) s0 = sT
pseudo-count based methods (Bellemare et al., 2016). In addition, we validate that any advantages of DISDAIN are not purely due to using an ensemble of discriminators by evaluating an ablation that includes the same ensemble but removes the DISDAIN reward (i.e. Algorithm 1 with λ = 0). In all cases, our primary metric for comparison is the effective number of skills learnt, nskills (equation 4).
The effective number of skills is just an interpretable transformation of the mutual information objective shared by a large number of unsupervised skill learning algorithms (e.g. Gregor et al. (2016); Achiam et al. (2018); Eysenbach et al. (2019); Hansen et al. (2020)). As such, the fact that DISDAIN helps better maximize this objective should be worthwhile in its own right, considering the various use cases that motivated this objective in the existing literature. That said, we recognize that it is not obvious what utility comes with increasing the number of effective skills. To address this, we also measure three surrogates of skill utility: downstream goal achievement, unsupervised reward attainment, and unsupervised state coverage.
Distributed training We use a distributed actor-learner setup similar to R2D2 (Kapturowski et al., 2019), except we do not use replay prioritization or burn-in, and Q-value targets are computed with Peng’s Q(λ) (Peng and Williams, 1994) rather than n-step double Q-learning. In all cases, we found it important to train separate Q-functions for the skill learning rewards and exploration bonuses (DISDAIN, RND, or count-based). This follows from Burda et al. (2019), who also found that this setup helped stabilize learning. Unlike that work, we need to specify a value target for each Q-function, which we take to be the value of the action that maximizes the composite value function, where the two values are added with a tunable weight λ. The discriminator is trained on the same mini-batches sampled from replay to train the Q-functions. For Atari experiments, the Q-networks process batches of state, skill, and action tuples to produce scalar Q-values for each, and the ResNet state embedding network used in Espeholt et al. (2018) is shared by both of the Q-networks and discriminator. For the Four Rooms grid world, we use tabular Q-functions and discriminators. Further implementation details can be found in appendix A.
Skill discovery Our skill discovery baseline deviates slightly from prior work in order to make it representative of the skill discovery literature as a whole. We utilize a simple discriminator that receives only the final state of a trajectory (so O(τ) = sT ), omitting the state-conditional prior and action entropy bonuses used in specific algorithms (Gregor et al., 2016; Eysenbach et al., 2019). One generally helpful change that we had not previously encountered was to clip negative skill rewards to 0 (rskill = max(rskill, 0)). This yields a strictly tighter lower bound on the mutual information (equation 2), since negative rewards imply the discriminator’s performance is worse than chance.
Hyperparameters Most of the RL hyperparameters were not tuned, but rather taken from standard values known to be reasonable for Peng’s Q(λ). The skill discovery specific hyperparameters were tuned for the basic algorithm without exploration bonuses, and then reused in all conditions. Our RND implementation was first tuned without skill learning to achieve Atari performance competitive with the original paper. For all exploration bonuses (DISDAIN, RND, count), we tuned their reward weighting (e.g. λ in algorithm 1), while for DISDAIN we additionally swept the ensemble size (N ).
Baselines As with the skill discovery reward, we apply exploration bonuses only to the terminal states of each skill trajectory. For the count-based bonus, we track the number of times an agent ends a skill trajectory in each state n(s) and apply the exploration bonus rT = 1/ √ n(sT ). For RND, we follow the details of Burda et al. (2019) as closely as possible. For our target and predictor networks, we use the same ResNet architecture as the policy observation embedding described above, and then project to a latent dimension of 128. Rather than normalizing observations based on running statistics, we found it more reliable to use the standard 1255 normalization of Atari observations.
DISDAIN Our ensemble-based uncertainty estimator required many design choices, including the size of ensembles, and to what extent the ensemble members shared training data and parameters. In all of the domains we tested, we found training ensemble members on different batches to be unnecessary, similar to Osband et al. (2016). In the tabular case (Four Rooms), parameter-sharing is not a concern, and we found an ensemble size of N = 2 to be sufficient. For Atari, the ensemble of discriminators reuse the single ResNet state embedding network which is shared by the value function. The input to the ensemble will inevitably drift, even if the data distribution remains constant, since the ResNet representations evolve according to a combination of the value function and discriminator updates. Expressive ensembles (e.g. with hidden layers) never converged in practice. By contrast, large linear ensembles (N = 40) were reliably convergent, with convergence time increasing with ensemble size. We follow Osband et al. (2016) in scaling the gradients passed backward through the embedding network by 140 to account for the increase in effective learning rate.

heading: 4.1 FOUR ROOMS

text:
First, we evaluate skill learning in the illustrative grid world seen in figure 3. There are 4 rooms and 104 states. The agent begins each episode in the top left corner and at each step chooses an action from the set: left, right, up, down, or no-op. Episodes are 20 steps and we sample one skill per episode (i.e. T = 20). The episodes are long enough to reach all but one state, allowing for a maximum of 103 skills. For each method, we set NZ = 128 to make this theoretically possible.
Results As seen in figure 3, even in this simple task, unbonused agents are unable to exceed 30 skills and barely leave the first room (see figures 18 and 19 for example skill rollouts). With both DISDAIN and a count bonus, agents explore all four rooms and learn approximately triple the number of skills, with the best seeds learning approximately 90 skills. Both bonuses do slow learning and add variance due to the addition of a separate learned Q-function (see figure 20 for individual seeds). The ensemble-only ablation provides a small benefit to skill learning, but far less than DISDAIN, demonstrating that the DISDAIN exploration bonus (rDISDAIN), rather than the ensembling, drives the increased performance. As a simple demonstration of the usefulness of our learnt skills for downstream tasks, we also evaluate each methods’ skills on the original Four Rooms reward (i.e. reaching specified goal states) without additional finetuning. Specifically, for each method, we sample each accessible state of the environment as a goal, pass it through the trained discriminator qφ, choose the highest probability skill z, rollout the policy for one episode conditioned on z, and track the fraction of goal states successfully reached (similar to the imitation learning evaluation of Eysenbach et al. (2019)). As seen in Figure 3b, more learnt skills leads to better downstream task performance.

heading: 4.2 ATARI

text:
Next we consider skill learning on a standard suite of 57 Atari games, where prior work has shown that learning discrete skills is quite difficult (Hansen et al., 2020). In addition, it is a non-trivial test for our ensemble uncertainty based method to work in the function approximation setting, since this depends on how each ensemble member generalizes to unseen data. Here we use NZ = 64 and T = 20. Since Atari episodes vary in length, skills may be resampled within an episode.
The count-based baseline used in the Four Rooms experiments cannot be directly applied to a nontabular environment like Atari. While pseudo-count methods have been used here (Bellemare et al., 2016), Random Network Distillation (RND) similarly induces long term exploration (Burda et al., 2019) and has been used for this purpose in a state-of-the-art Atari agent (Badia et al., 2020). While newer approaches surpass RND in some domains (Raileanu and Rocktäschel, 2020; Seo et al., 2021), it is unclear if this is the case across the full Atari suite. So, at present, we believe RND is the most compelling baseline to compare against DISDAIN.
Results As pointed out in Agarwal et al. (2021), making statistically sound conclusions can be challenging in the few-seed, many-task setting of Atari. Thus, we follow their recommendations and focus our results primarily on statistically robust distributional claims here. Individual learning curves and game-by-game results are available in the appendix.
As shown in figure 4, DISDAIN increases the effective number of skills learnt across the Atari suite, with only modest damage to sample effiency. Additionally, we found that DISDAIN’s performance
boosts are robust to its key hyperparameters, namely the bonus weight λ and ensemble size N (see figure 8), with significant gains over unbonused skill learning maintained over more than an order of magnitude of variation in both.
Notably, our results show that RND fails to significantly aid in skill learning. A sweep over the RND bonus weight is shown in figure 11, but the summary is that as the RND bonus becomes similar in magnitude to the skill learning reward, it damages skill learning rather than helping, and so the “best" RND bonus weight for skill learning is approximately zero. This is perhaps unsurprising: RND was designed in the context of stationary task rewards (as are most other exploration methods, such as pseudo-counts), whereas skill learning objectives produce a highly non-stationary reward function. This highlights the importance of using an exploration bonus tailored to the skill discovery setting. Additionally, the failure of the “ensemble-only" baseline to significantly increase skill learning demonstrates that is the rDISDAIN exploration bonus that is crucial to DISDAIN’s success, and not just the ensembling of the discriminator.
0.50 0.55 0.60 0.65 P(higher reward)
Zero-shot reward improvement
-50% 0% +50% +100% +150% Lifetime coverage boost ( )
0.00
0.25
0.50
0.75
1.00
Fr ac
tio n
of ru
ns w
ith b
oo st
>
Lifetime coverage improvement RND Ensemble DISDAIN
Figure 5: Qualitative analysis of skills learnt on Atari. Both plots depict improvement over Unbonused skill learning aggregated across seeds and games. Left: Probability of improvement on zeroshot reward evaluation. Right: Distribution of percentile improvements on lifetime coverage. Error bars depict 95% stratified boostrap CIs.
To further probe the utility of our learnt skills, we measure the unsupervised reward attainment of a policy that randomly switches between them. First reported in Hansen et al. (2020), the idea is that while this policy is certainly far from optimal, this metric indicates whether or not the skill space is sufficient to perform the various reward collecting behaviors involved in each game. Additionally, we measure lifetime state coverage as an indication of exploration throughout learning (on a subset of games supporting this metric; see Appendix B for details). Figure 5 confirms that DISDAIN leads to both increased reward attainment and lifetime coverage (see figures 14, 15, and 16 for additional lifetime and episodic state coverage results).

heading: 5 DISCUSSION & LIMITATIONS

text:
We introduced DISDAIN, an enhancement for unsupervised skill discovery algorithms which increases the effective number of skills learnt across a diverse range of tasks, by using ensemble-based uncertainty estimation to counteract a bias towards pessimistic exploration.
The connection between ensemble estimates of uncertainty and infomax exploration dates back to at least Seung et al. (1992), who use it to select examples in an active supervised learning setting. These ideas have more recently found use in the RL literature, with recent work using disagreement among ensembles of value functions (Chen et al., 2017; Flennerhag et al., 2020; Zhang et al., 2020) and forward models of the environment (Pathak et al., 2019; Shyam et al., 2019; Sekar et al., 2020) to drive exploration. In a closely related line of work, ensembles of Q-functions have been used to drive exploration without an explicit bonus based on disagreement (Osband et al., 2016; 2018). To our knowledge, DISDAIN represents the first application of these ideas to unsupervised skill discovery.
We focused on discrete skill learning methods due to their relative prevalence (Gregor et al., 2016; Eysenbach et al., 2019; Achiam et al., 2018; Baumli et al., 2021). In some cases, continuous or structured skill spaces might make more sense (Hansen et al., 2020; Warde-Farley et al., 2019). While the principles behind DISDAIN should still apply, further approximations may be necessary, e.g. for the entropy of the ensemble-averaged discriminator, which may be unavailable in closed form.
Designing agents that explore and master their environment in the absence of task reward remains an open problem, for which there exist many different families of approaches (e.g. reward-free exploration (Jin et al., 2020; Zhang et al., 2021)). In this paper, we focus on improving one such family - unsupervised skill learning through variational infomax (Section 2). By treating agents with DISDAIN, we empower them to better maximize their objective and learn more skills. Leveraging these skills for rapid task reward maximization remains an important direction for future research.

heading: ACKNOWLEDGEMENTS

text: The authors would like to thank Stephen Spencer for engineering and technical support, Ian Osband for feedback on an early draft, Rishabh Agarwal for suggestions on statistical analysis and plotting, and David Schwab for pointing us to the original Query by Committee work (Seung et al., 1992).

heading: A COMPUTE REQUIREMENTS AND HYPERPARAMETERS

text:
The compute cluster we performed experiments on is rather heterogeneous, and has features such as host-sharing, adaptive load-balancing, etc. It is therefore hard to give precise details regarding compute resources – however, the following is a best-guess estimate.
A full experimental training run for Atari lasted 4 days on average. Our distributed reinforcement learning setup (Espeholt et al., 2018) used 100 CPU actors and a single V100 GPU learner. Thus, we required approximately 9600 CPU hours and 96 V100 GPU hours per seed, with 3 seeds and 3 conditions per game.
Tuning required approximately 10 different hyper-parameters combinations on 6 games, amounting to 864,000 CPU hours and 8,640 V100 GPU hours. The results on the full suite of 57 Atari games required 4,924,800 CPU hours and 49,248 V100 GPU hours. Combining these, we get a total compute budget of 5,788,800 CPU hours and 57,888 V100 GPU hours.
It is worth remembering that the above is likely quite a loose upper-bound, as this estimate assumes 100 percent up time, which is far from the truth given the host-sharing and load-balancing involved in our setup. Additionally, V100 GPUs were chosen based on what was on hand; our models are small enough to fit on much cheaper cards without much slowdown.

heading: B COVERAGE METRICS

text: We calculate two related notions of coverage: lifetime and episodic. Lifetime coverage corresponds to the number of unique states encountered during an agent’s lifetime, whereas episodic coverage corresponds to the number of unique states encountered during each episode. Both rely on the notion of a unique state. The subset of games chosen for these metrics were those where a good notion of a unique state is simple: they all involve a controllable avatar that moves in a coordinate system, so unique avatar coordinates are used. This information is exposed in the RAM state of the Atari emulator, as shown in Anand et al. (2019).

title: LEARNING MORE SKILLS THROUGH OPTIMISTIC EXPLORATION


summaries: This paper presents a method to construct a molecular graph, which is inspired by a spanning tree. In specific, a molecule is constructed by a sequence of actions, each of which adds an atom, adds a bond, starts/ends a branch, adds residual edges for a circular structure, and terminates. The generative process is controlled by a transformer-based neural network, which is specialized to tree construction procedure. The empirical studies show that the proposed method can learn the distribution of molecules comparably or better than the existing methods, and can be used to molecular optimization.

summaries: This paper proposes a spanning tree based generative model (STGG) for molecular graphs. STGG sequentially generate a molecule's spanning tree and fill in the residual edges on the way. The spanning tree construction is similar to the standard SMILES representation, but the model operates on a molecular graph rather than a SMILES string. STGG adopts a tree-based transformer with relative positional encoding for tree generation, and a attention-based predictor for residual edge prediction. The method is evaluated on standard ZINC250K, QM9, and MOSES benchmarks and outperform existing baselines.

summaries: The paper proposes a novel spanning tree-based graph generation (STGG) framework. The key idea is to represent the molecular graph as a sequence of decisions according to a **novel spanning tree-based grammar**, and then model these decision sequences using the transformer.  To accommodate such a novel algorithm, several interesting techniques are involved, e.g., representing molecules as bipartite graphs, using tree-based positional encoding for transformer. The authors claim that the proposed method allows generating valid molecular structures and inferring the intermediate graph structure.

summaries: This paper proposes a new paradigm of graph generation, a transformer network sequentially generates a sequence of decisions to generate a spanning tree for a bipartite graph. These decisions have 7 forms such as attach atom, attach bond, branch start,  etc. After the spanning tree, they add residual edges. The authors also have a focus on generating valid graphs (which is seen in the results) in which they mask out invalid decisions during the generation process itself.

abstractText: In this paper, we explore the problem of generating molecules using deep neural networks, which has recently gained much interest in chemistry. To this end, we propose a spanning tree-based graph generation (STGG) framework based on formulating molecular graph generation as a construction of a spanning tree and the residual edges. Such a formulation exploits the sparsity of molecular graphs and allows using compact tree-constructive operations to define the molecular graph connectivity. Based on the intermediate graph structure of the construction process, our framework can constrain its generation to molecular graphs that satisfy the chemical valence rules. We also newly design a Transformer architecture with tree-based relative positional encodings for realizing the tree construction procedure. Experiments on QM9, ZINC250k, and MOSES benchmarks verify the effectiveness of the proposed framework in metrics such as validity, Fréchet ChemNet distance, and fragment similarity. We also demonstrate the usefulness of STGG in maximizing penalized LogP value of molecules.

authors:
Sungsoo Ahn
Junsu Kim
Hankook Lee
Jinwoo Shin

title: Guiding deep molecular optimization with genetic exploration

venue: Advances in Neural Information Processing Systems,

authors:
Lili Chen
Kevin Lu
Aravind Rajeswaran
Kimin Lee
Aditya Grover
Michael Laskin
Pieter Abbeel
Aravind Srinivas
Igor Mordatch

title: Decision transformer: Reinforcement learning via sequence modeling

venue: arXiv preprint arXiv:2106.01345,

title: Defining and exploring chemical spaces

venue: Trends in Chemistry,

authors:
Hanjun Dai
Yingtao Tian
Bo Dai
Steven Skiena
Le Song

title: Syntax-directed variational autoencoder for structured data

venue: In International Conference on Learning Representations,

authors:
P Erdös
A Rényi

title: On random graphs

venue: In Publicationes mathematicae 6,

authors:
Wenhao Gao
Connor W Coley

title: The synthesizability of molecules proposed by generative models

venue: Journal of Chemical Information and Modeling,

authors:
Rafael Gómez-Bombarelli
Jorge Aguilera-Iparraguirre
Timothy D Hirzel
David Duvenaud
Dougal Maclaurin
Martin A Blood-Forsythe
Hyun Sik Chae
Markus Einzinger
Dong-Gwang Ha
Tony Wu

title: Design of efficient molecular organic light-emitting diodes by a high-throughput virtual screening and experimental approach

venue: Nature materials,

authors:
Shion Honda
Hirotaka Akita
Katsuhiko Ishiguro
Toshiki Nakanishi
Kenta Oono

title: Graph residual flow for molecular graph generation

authors:
John J Irwin
Teague Sterling
Michael M Mysinger
Erin S Bolstad
Ryan G Coleman

title: Zinc: a free tool to discover chemistry for biology

venue: Journal of chemical information and modeling,

authors:
Michael Janner
Qiyang Li
Sergey Levine

title: Reinforcement learning as one big sequence modeling problem

venue: arXiv preprint arXiv:2106.02039,

authors:
Wengong Jin
Regina Barzilay
Tommi Jaakkola

title: Junction tree variational autoencoder for molecular graph generation

venue: In International Conference on Machine Learning,

authors:
Wengong Jin
Kevin Yang
Regina Barzilay
Tommi Jaakkola

title: Learning multimodal graph-tograph translation for molecule optimization

title: Hierarchical generation of molecular graphs using structural motifs

venue: In International Conference on Machine Learning, pp. to appear,

title: Molecular hypergraph grammar with its application to molecular optimization

authors:
Hyunseung Kim
Jonggeol Na
Won Bo Lee

title: Generative chemical transformer: attention makes neural machine learn molecular geometric structures via text

venue: arXiv preprint arXiv:2103.00213,

authors:
Aviral Kumar
Sergey Levine

title: Model inversion networks for model-based optimization, 2020

authors:
Matt J Kusner
Brooks Paige
José Miguel Hernández-Lobato

title: Grammar variational autoencoder

authors:
Phillip Lippe
Efstratios Gavves

title: Categorical normalizing flows via continuous transformations

authors:
Qi Liu
Miltiadis Allamanis
Marc Brockschmidt
Alexander Gaunt

title: Constrained graph variational autoencoders for molecule design

venue: In Advances in Neural Information Processing Systems,

authors:
Tairan Liu
Misagh Naderi
Chris Alvin
Supratik Mukhopadhyay
Michal Brylinski

title: Break down in order to build up: decomposing small molecules for fragment-based drug design with e molfrag

authors:
Ilya Loshchilov
Frank Hutter

title: Decoupled weight decay regularization

authors:
Denis Lukovnikov
Asja Fischer

title: Insertion-based tree decoding. In Findings of the Association for Computational Linguistics: ACL-IJCNLP

authors:
Youzhi Luo
Keqiang Yan
Shuiwang Ji

title: GraphDF: A discrete flow model for molecular graph generation

authors:
Kaushalya Madhawa
Katushiko Ishiguro
Kosuke Nakago
Motoki Abe

title: Graphnvp: An invertible flow model for generating molecular graphs

authors:
Daniil Polykovskiy
Alexander Zhebrak
Benjamin Sanchez-Lengeling
Sergey Golovanov
Oktai Tatanov
Stanislav Belyaev
Rauf Kurbanov
Aleksey Artamonov
Vladimir Aladinskiy
Mark Veselov
Artur Kadurin
Simon Johansson
Hongming Chen
Sergey Nikolenko
Alan AspuruGuzik
Alex Zhavoronkov

title: Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Models

venue: Frontiers in Pharmacology,

authors:
Mariya Popova
Mykhailo Shvets
Junier Oliva
Olexandr Isayev

title: MolecularRNN: Generating realistic molecular graphs with optimized properties

venue: arXiv preprint arXiv:1905.13372,

authors:
Kristina Preuer
Philipp Renz
Thomas Unterthiner
Sepp Hochreiter
Günter Klambauer

title: Fréchet ChemNet distance: a metric for generative models for molecules in drug discovery

authors:
Oleksii Prykhodko
Simon Viet Johansson
Panagiotis-Christos Kotsias
Josep Arús-Pous
Esben Jannik Bjerrum
Ola Engkvist
Hongming Chen

title: A de novo molecular generation method using latent vector based generative adversarial network

venue: Journal of Cheminformatics,

authors:
Raghunathan Ramakrishnan
Pavlo O Dral
Matthias Rupp
O Anatole Von Lilienfeld

title: Quantum chemistry structures and properties of 134 kilo molecules

venue: Scientific data,

title: Reinforcement learning upside down: Don’t predict rewards–just map them to actions

venue: arXiv preprint arXiv:1912.02875,

authors:
Marwin HS Segler
Thierry Kogej
Christian Tyrchan
Mark P Waller

title: Generating focused molecule libraries for drug discovery with recurrent neural networks

venue: ACS central science,

authors:
Chence Shi
Minkai Xu
Zhaocheng Zhu
Weinan Zhang
Ming Zhang
Jian Tang

title: GraphAF: a flow-based autoregressive model for molecular graph generation

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Łukasz Kaiser
Illia Polosukhin

title: Attention is all you need

venue: In Advances in neural information processing systems,

authors:
Johannes Villmow
Adrian Ulges
Ulrich Schwanecke

title: A structural transformer with relative positions in trees for code-to-sequence tasks

venue: In 2021 International Joint Conference on Neural Networks (IJCNN),

title: SMILES, a chemical language and information system. 1. introduction to methodology and encoding rules

venue: Journal of chemical information and computer sciences,

authors:
Chengxuan Ying
Tianle Cai
Shengjie Luo
Shuxin Zheng
Guolin Ke
Di He
Yanming Shen
Tie-Yan Liu

title: Do transformers really perform bad for graph representation

venue: arXiv preprint arXiv:2106.05234,

authors:
Jiaxuan You
Bowen Liu
Zhitao Ying
Vijay Pande
Jure Leskovec

title: Graph convolutional policy network for goal-directed molecular graph generation

authors:
Chengxi Zang
Fei Wang

title: Moflow: an invertible flow model for generating molecular graphs

venue: In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining,

heading: 1 INTRODUCTION

text:
Researchers have extensively studied graph generative models, dating back to the early works of Erdös Rényi (Erdös et al., 1959). Recently, models based on deep neural networks (DNNs) have gained much attraction due to their expressive power in learning a graph dataset. The molecule-generating DNNs stand out among them for their success in the task of drug discovery.
Recent works have proposed molecule-generating DNNs based on string-based and graph-based representations (Segler et al., 2018; Jin et al., 2018; You et al., 2018; Shi et al., 2020; Jin et al., 2020). For example, Segler et al. (2018) proposed to train language models on the domain-specific linear string representation of molecules, i.e., simplified molecular-input line-entry system (SMILES, Weininger 1988). Since the string-based models ignore the inherent graph structure, recent works explore the graph-based generation that use (a) atom-by-atom (You et al., 2018; Shi et al., 2020; Luo et al., 2021) or (b) substructure-based (Jin et al., 2018; 2019; 2020) operations.
Notably, the substructure-based generative models (Jin et al., 2018; 2019; 2020) successfully exploit the molecular prior knowledge: the graphs are sparsely connected and can be represented as a junction tree with molecular substructure as building blocks. Based on such knowledge, the models use the junction tree construction operators which (a) require a fewer number of steps to generate the whole molecular graph and (b) guarantee generating molecules that satisfy the chemical valence rules. However, despite such advantages, a recent benchmark (Polykovskiy et al., 2020) suggests that they do not outperform the existing methods in terms of learning the data distribution, even when compared with the simple SMILES-based language models. We hypothesize that this is due to the models using a coarse-grained representation of the molecule and they may lack the ability to learn the inner semantics of each substructure-based building block.
Contribution. In this work, we propose a novel framework, coined spanning tree-based graph generation (STGG), for fine-grained generation of molecules while exploiting their sparsity.1 Mainly inspired from the SMILES representation of molecules, our idea is to generate the molecular graph as a composition of a spanning tree and the corresponding residual edges with atoms and bonds as building blocks. Such a formulation allows our framework to utilize compact tree-constructive
1While our framework is designed for general sparse graphs, we focus on the molecular graphs in this paper.
operations to define the molecular graph connectivity. See Figure 1 for an illustration of how we formulate the generation of a molecular graph as a sequence of tree-constructive operations.
Since our framework maintains the molecular graph structure during construction, it can pre-determine decisions that (a) violate the graph construction rule and (b) lead to molecules that violate the chemical valence rule. Such criteria allow control over the generative model to guarantee generating valid molecular graphs by forbidding invalid actions. This is in contrast to prior works (Shi et al., 2020; Luo et al., 2021) that generate the molecular graph atom-by-atom but determines the validity of construction operations through a sample-rejection scheme.
To recognize the spanning tree-based representation used in our STGG framework, we propose a Transformer architecture (Vaswani et al., 2017) with tree-based relative encoding. Inspired by recent works (Villmow et al., 2021; Lukovnikov & Fischer, 2021; Ying et al., 2021) on tree-based and graph-based Transformers, our framework expresses the relative position between two vertices as the number of forward and reverse edges in the shortest path between them. We also introduce an attention-based mechanism for constructing residual edges.
We experiment on popular graph generation benchmarks of QM9, ZINC250K, and MOSES to validate the effectiveness of our algorithm. In the experiments on QM9 and ZINC, our STGG framework outperforms the existing graph-based generative models by a large margin. In the MOSES benchmark, our algorithm achieves superior performance compared to both string-based and graphbased methods for majority of the metrics, e.g., Fréchet ChemNet distance (Preuer et al., 2018) and fragment-based similarity. We also conduct experiments on the offline optimization task for high penalized octanol-water partition coefficient and achieve competitive results.

heading: 2 SPANNING TREE-BASED GENERATION OF GRAPHS (STGG)

heading: 2.1 OVERVIEW

text:
In this section, we introduce our spanning tree-based graph generation (STGG) framework to sequentially generate a molecule as a composition of a spanning tree and residual edges. To this end, we propose compact tree-constructive operations inspired by the simplified molecular-input line-entry system (SMILES, Weininger, 1988). In contrast to the existing SMILES-based molecular generative methods, our framework (a) allows inferring the intermediate graph structure and (b) is generally applicable to graph types other than molecules. In particular, (a) further enables our framework to control the construction process such that the sequential operations comply with tree-constructive grammar and only generate molecules satisfying the chemical valence rule.
Molecular graph representation. To apply our framework, we represent a molecule as a bipartite graph G = (A,B, E) where A and B are the set of vertices associated with atoms and bonds of the molecule, respectively.2 Each edge {a, b} ∈ E is assigned for each adjacent pair of atom and bond. We assign attributes xa ∈ Xatom and xb ∈ Xbond for vertices a ∈ A and b ∈ B to indicate the corresponding atom type and bond order, respectively. For example, {"C", "N", "O"} ⊆ Xatom and {"-", "="} ⊆ Xbond. See Figure 1 for an example of such a molecular graph representation. Molecular graph from sequence of decisions. To generate the molecular graph G = (A,B, E), our framework makes a sequence of decisions d1, . . . , dT to generate a spanning tree T = (AT ,BT , ET ) and a set of residual edges ER = E \ET . At each iteration, seven types of decisions are applicable, i.e., attach atom, attach bond, branch start, branch end, res atom, res bond, and terminate.
2Many existing works, e.g., (Shi et al., 2020), use a non-bipartite graph with bonds assigned to edges.
See Table 1 for examples of decisions and the corresponding operations. We provide a detailed description of the graph construction process in Section 2.2.
Generating valid molecular graphs. Without any control, a model may generate decisions that (a) do not comply with the grammar of STGG or (b) leads to a molecule violating the chemical valence rule. To prevent this scenario, we conduct two criteria for determining validity of the given decision for (a) and (b). We further elaborate this in Section 2.3.

heading: 2.2 DECISION PROCESS FOR SPANNING TREE-BASED GRAPH GENERATION

text:
We now explain how our STGG framework incorporates the decisions d1, . . . , dT to build the spanning tree T = (AT ,BT , ET ) and residual edges ER from scratch. To this end, our framework introduces the state information of (a) a pointer vertex ipoint ∈ AT ∪ BT for specifying the target of the next operation, (b) a stack Sbranch that stores vertices to use later as a starting point of a “branch” in the spanning tree, and (c) a list Lres that stores vertices to use later for constructing residual edges.
Algorithm 1 Tree-based generation of molecular graphs 1: Input: sequence of decisions d1, . . . , dT . 2: Output: graph G = (A,B, E), atom attributes {xa}a∈A, and bond attributes {xb}b∈B 3: Set AT ← ∅, BT ← ∅, ET ← ∅, ER ← ∅, and T ← (AT ,BT , ET ). . Initialize the empty graph. 4: Set Lres as an empty list and Sbranch as an empty stack. 5: for t = 1, . . . , T do 6: if dt ∈ Xatom then . Add a new atom vertex. 7: Create a new atom vertex a and set A ← A∪ {b} and xa ← dt. 8: If |BT | > 0, set ET ← ET ∪ {{a, ipoint}}. . Edge is added when tree is not empty. 9: Set ipoint ← a.
10: if dt ∈ Xbond then . Add a new bond vertex. 11: Create a new bond vertex b and set ET ← ET ∪ {{b, ipoint}}, BT ← BT ∪ {b}, and xb ← dt. 12: Set ipoint ← b. 13: if dt = "*" then Insert ipoint into Lres. . Add pointer vertex into the queue. 14: if dt ∈ Lres then Pop dt from Lres and update ER ← ER ∪ {{ipoint, dt}}. . Add a new residual edge. 15: if dt = "(" then Insert ipoint into Sbranch. . Add pointer vertex into the stack. 16: if dt = ")" then Set ipoint ← pop(Sbranch) . Update pointer vertex from the stack. 17: Set A ← AT , B ← BT , and E ← ET ∪ ER.
In what follows, we describe the seven types of operations, i.e., attach atom, attach bond, branch start, branch end, res atom, res bond, and terminate, corresponding to decision values d ∈ Xatom ∪ Xbond ∪ Lres ∪ {"(", ")", "*", "[eos]"} in detail. See Table 1 for the pairs of operations and the corresponding decisions. We also provide an example of the graph construction process in Figure 2.
Attaching atom and bond vertices to the spanning tree. If the decision d specifies one of the atom or bond attributes, i.e., d ∈ Xatom or d ∈ Xbond, it applies the corresponding attach atom and attach bond operations, respectively. To be specific, the attach atom operation adds a new atom vertex a into the spanning tree T as a neighbor of the pointer vertex ipointer, i.e., AT ← AT ∪ {a}, ET ← ET ∪ {a, ipointer}. The value d is set as the new atom attribute, i.e., xa ← d. The newly added vertex is set as the next pointer vertex, i.e., ipointer ← a. The attach bond operation similarly adds a new bond vertex. For example, a line graph can be expressed as a sequence of attach atom and attach bond operations, e.g., C-C-C where "C" ∈ XA and "-" ∈ XB. Branching out the spanning tree. To express graph structures with vertices of degree larger than two, our framework utilizes pairs of the branch start and the branch end operations with decision values of "(" and ")", respectively. To be specific, the branch start operation inserts the current pointer vertex into a stack Sbranch of vertices. Then the branch end operation pops a vertex from the stack Sbranch and sets it as the new pointer vertex. For an example, a graph with one atom vertex of degree three is constructed from a sequence of decisions C-C(-C)(-C).
Adding residual edges. To construct cyclic molecular graphs, our framework generates residual edges based on pairs of res atom operation and res bond operation, corresponding to decision values of "*" and d ∈ Lres, respectively. To be specific, the res atom operation inserts the current (atom) pointer vertex into a list Lres. Next, when the decision value d ∈ Lres is received for the res bond operation, the corresponding vertex d is popped from the list Lres and forms a new residual edge with the current (bond) pointer vertex, i.e., ER ← ER ∪ {{d, ipointer}}. For an example, a cyclic molecular graph is constructed from a sequence of decisions C*-O-O-1, where "1" indicates res bond operation with decision of the first atom vertex with attribute "C".
Termination. The decision "[eos]" applies the terminate operation to finish the construction. We provide the full algorithm in Algorithm 1. We also provide an algorithm to extract a sequence of decisions for constructing a given graph in Appendix A. Such an algorithm is used to obtain sequence of decisions as targets for training the generative model under the STGG framework.

heading: 2.3 MASKING OUT INVALID DECISIONS FOR A VALID MOLECULAR GRAPH

text:
Based on Algorithm 1, we develop two criteria for determining whether if a sequence of decisions leads to (a) valid generation of a molecular graph and (b) generation of a molecule satisfying the valence rule. Such criteria are used to mask out invalid decisions to guarantee generating a valid molecular graph.
Validity of graph generation. To determine whether if a sequence of decisions lead to a valid generation of a molecular graph, we propose an algorithm that outputs a set of valid decisions given the previous decision d, stack of pointer vertices Sbranch, and list of atom vertices Lres during execution of Algorithm 1. In what follows, we provide a brief description of grammars enforced by the algorithm. We also provide the detailed algorithm in Appendix B.
• The branch end operation only appears when the stack of pointer vertex Sbranch is non-empty. • The operations res atom and res bond are atom-specific and bond-specific, hence they only
appear when the pointer vertex is located at an atom vertex and a bond vertex, respectively. • All the bond vertices have degree of two, hence branch start and branch end operations only
appear when the pointer vertex is located at an atom vertex. • The stack do not contain duplicates of the a pointer vertex at the same time.
Here, we note that our criteria for valid molecular graph generation do not enforce branches and rings to be closed, e.g., C*=C-C#N is allowed by our criteria. This does not violate the validity since our Algorithm 1 may still define a valid molecular graph by ignoring the open branches and the open rings during construction, i.e., C*=C-C#N generates a molecule identical to that of C=C-C#N.
Validity of satisfying the valence rule. To consider chemical validity of molecules, our framework offers the ability to constrain the its generation on molecules that satisfy the valence rule for each atom. That is, the generated graph G = (A,B, E) satisfies the constraint v(xa) ≥ ∑ b∈N (a) o(xb) for every atom vertex a ∈ A where v(xa) denotes the valence of an atom type xa and o(xb) denotes the bond order. To this end, we keep a record r(a) of available valence for each atom a ∈ A and update them for each decision. For example, when a bond vertex b is newly added, record of the neighboring atom vertex a is updated by r(a) ← v(a) − o(xb). The main idea is to forbid actions that lead to negative values of r(a). We provide a detailed algorithm in Appendix C.

heading: 3 TRANSFORMER ARCHITECTURE FOR TREE-BASED GENERATION

text: In this section, we describe our deep neural network architecture for generating sequence of decisions d1, . . . , dT under the STGG framework. To accurately recognize the decision process, we employ the tree-based relative positional encodings on the intermediate spanning tree T . We also introduce an attention mechanism to express a probability distribution over Lres which depends on intermediate state of the algorithm.

heading: 3.1 TREE-BASED POSITIONAL ENCODING FOR MULTI-HEAD ATTENTION LAYERS

text:
Each intermediate layer in our model is a combination of a multi-head self-attention module and a position-wise feed-forward neural network similar to that of Vaswani et al. (2017). The main difference is on how we modify the architecture to incorporate tree-based positional encodings. To be specific, let H = [h>1 , . . . , h > T ] ∈ RT×` denote the input of a self-attention module where d is the hidden dimension and ht ∈ R1×` is the hidden representation at position t. The input H is projected by three matrices WQ ∈ R`×`K ,WK ∈ R`×`K and WV ∈ R`×`V to the corresponding representations Q,K and V , respectively. A single self-attention head is then calculated as
Q = HWQ, K = HWK , V = HWV , (1) A = QK>√ `K + P, Pt1,t2 = z (1) φforward(t1,t2) + z (2) φbackward(t1,t2) + z (3) φseq(t1,t2) , (2)
Attention(H) = SoftMax(M ◦A)V, (3) where Attention(H) is output of the attention head,M is the triangular mask to forbid the model from accessing future information while making a prediction, and ◦ denotes the element-wise multiplication between matrices.
Furthermore, P is the newly introduced relative positional encoding. It is a summation over the trainable embedding vectors z(1), z(2), z(3) indexed by relative position values of φforward(t1, t2), φbackward(t1, t2), and φseq(t1, t2). To be specific, the tree-based relative positions φforward(t1, t2) and φbackward(t1, t2) denotes the number of forward and backward edges in the spanning tree path between pointer vertices at the t1-th and t2-th time step. The direction of edge is decided by order of generation in the STGG framework. Such an encoding was inspired from recent works (Villmow et al., 2021; Lukovnikov & Fischer, 2021; Ying et al., 2021) using Transformers to recognize graphs and trees. Finally, the sequence-based relative position φseq(t1, t2) = t1 − t2 denotes the relative difference of time-steps for the decisions.

heading: 3.2 ATTENTION FOR UPDATING RESIDUAL EDGES.

text:
Our model generates a categorical distribution over the space of X = Xatom ∪ Xbond ∪ Lres ∪ {"(", ")", "*", "[eos]"}. It is relatively straight-forward to output an unnormalized probability over values of Xatom ∪ Xbond ∪ {"(", ")", "*"} using a linear classifier on top of the Transformer model. However, it is non-trivial to assign probability values for res bond operation, i.e., decisions values of d ∈ Lres, since Lres varies between different time-steps. To handle this case, we use an attention-based mechanism for assigning unnormalized probability to decision values in the list Lres. To be specific, at the final layer of our model, we obtain the following probability distribution p(d).
p(d) ∝ { mg(d) ·mv(d) · exp(w>d h) ∀d ∈ Xatom ∪ Xbond ∪ {"(", ")", "*"}, mg(d) ·mv(d) · exp ( h>dW1W > 2 h ) ∀d ∈ Lres, (4) where wd ∈ R1×` is a decision-specific vector, W1,W2 ∈ R`× ˜̀ are weight matrices, and h is the decision embedding, i.e., output of the Transformer layer corresponding to the previously made decision. Furthermore, hd is the embedding corresponding to a past decision d ∈ Lres. Finally, mv(d),mg(d) are the mask for excluding invalid decisions that violate the validity of graph generation and valence rule, respectively. The masks are obtained using the criteria explained in Section 2.3. We use the mask during both training and evaluation of the model; this differs from existing graphgenerative models which forbid invalid decisions only at evaluation using a sample-rejection scheme.

heading: 4 RELATED WORKS

text:
SMILES-based molecular generative models. Several studies proposed to generate a SMILES representation of molecules using string-based (Gómez-Bombarelli et al., 2016; Segler et al., 2018; Kim et al., 2021) or grammar-based (Kusner et al., 2017; Dai et al., 2018) models. While our STGG is largely inspired from such works, our STGG allows realizing the intermediate graph structure of the molecule being constructed while the SMILES-based models cannot. This difference allows the adoption of structure-aware deep neural networks to STGG. To be specific, the difference between STGG and the SMILES-based models appears from our newly introduced graph construction procedure using a pointer vertex ipoint, a vertex-list L, and a vertex-stack S . They allow recognizing an incomplete sequence of decisions as a graph and assigning positions to each decision. In contrast,
an incomplete SMILES string does not define a graph structure and assigning positions to each character is non-trivial.
Graph-based molecular generative models. Researchers have developed a large variety of molecular graph generation frameworks based on atom-wise and bond-wise operations (You et al., 2018; Kajino, 2019; Popova et al., 2019; Madhawa et al., 2019; Honda et al., 2019; Shi et al., 2020; Zang & Wang, 2020; Luo et al., 2021). Our STGG framework simplifies the decision space of such models by exploiting the tree-like graph structures of molecules. To be specific, STGG requires O(|A|+ |B|) decisions for constructing a molecule while the existing atom-by-atom graph generative models typically require O(|A|2) decisions. This implies that our generative model requires a smaller number of decisions for sparse graphs like molecules, i.e., when |B| is small. Furthermore, our work is the first to successfully train a Transformer architecture (Vaswani et al., 2017) for graph-based molecule generation.
In another line of research, several works (Jin et al., 2018; 2019; 2020) proposed generative models based on using the junction-tree representation with molecular substructures as building blocks. Based on such a representation, such works utilize tree-constructive operations to generate the full graph. Since they operate on such a coarse-grained molecular representation, they typically require a fewer number of building blocks to generate the whole molecule. In comparison, our STGG framework utilizes a more fine-grained molecular representation and may additionally learn the inner semantics of substructures that are used as building blocks for the junction tree.

heading: 5 EXPERIMENT

text: In this section, we report the experimental results of the proposed spanning tree-based graph generation (STGG) framework. In Section 5.1 and 5.2, we compare with the existing graph generative models in the ZINC250K (Irwin et al., 2012) and QM9 (Ramakrishnan et al., 2014). We provide ablation studies on each component of our method using the ZINC250K dataset. In Section 5.2, we compare with the existing molecule generative models using the MOSES benchmark (Polykovskiy et al., 2020). Finally, in Section 5.3, we provide our results on the molecular optimization task with respect to the penalized octanol-water partition coefficient function (PLOGP). We provide the implementation details and illustration of the generated molecules in Appendix D and E, respectively.

heading: 5.1 MOLECULE GENERATION ON ZINC250K AND QM9 DATASETS

text:
We first compare to the literature standard for the molecular generation task in the ZINC250K and the QM9 datasets. To this end, we train our generative model on the respective datasets and sample 10,000 molecules to measure (a) the ratio of valid molecules (VALID), (b) the ratio of unique molecules (UNIQUE), and (c) the ratio of novel molecules with respect to the training dataset (NOVEL). We compare with the numbers reported by recently proposed graph generative models (Shi et al., 2020; Luo et al., 2021). We also provide an additional baseline of a transformer architecture trained to generate the SMILES representation for the molecule (SMILES-TRANSFORMER).
We mark CORRECTABLE for methods which can optionally use a sample-rejection scheme to forbid decisions that violate the chemical rules at evaluation. Note that our framework can train the generative model under the valence correction mask during training, while existing graph generative
models use the valence correction only at evaluation. However, for comparison, we do not use the valence correction mask during training in this experiment.
We report the experimental results in Table 2. In the table, we observe that our STGG framework outperforms all the existing molecular graph generative models for high VALID at the cost of relatively lower NOVEL. In particular, our generative model can achieve a 100% ratio of valid molecules in the QM9 dataset even without any correction procedure. Such a result highlights the how our model can effectively learn the chemical rules and model the underlying distribution.
Finally, our STGG framework performing better than the SMILES-based transformer implies how the performance of our generative model stems from the STGG framework, rather than using the Transformer architecture.
Ablation studies. We also conduct ablation studies on the ZINC250k dataset to verify the effectiveness of our method. To this end, we report the experimental results of our method without specific components. To be specific, we ablate the effects of using sequential relative positional encoding (S), treebased relative positional encoding (T), graph-construction mask (G), and valence rule mask (V). We also consider an additional baseline of using the absolute positional encoding (A) as in the original Transformer architecture (Vaswani et al., 2017). In Table 3 and Figure 4, one can observe how each component of
our algorithm is crucial for achieving high VALID. In particular, the tree encoding is essential for the performance, showing the importance of tree-based representation that we use in our model.

heading: 5.2 MOLECULE GENERATION ON THE MOSES BENCHMARK

text:
We also compare our method on the MOSES benchmark with the existing models. The MOSES benchmark offers a large collection of metrics to access the overall quality of generated molecules. To be specific, in addition to VALID, UNIQUE, NOVEL, we consider internal diversity of molecules (INTDIV), ratio of samples being accepted to chemical filters (FILTERS), Frétchet ChemNet Distance (FCD), nearest neighborhood similarity (SNN), frament similarity (FRAG), and Scaffold similarity
(SCAF). The similarity metrics of FCD, SNN, FRAG, SCAF are measured with respect to the test dataset of molecules and the scaffolds extracted from them.
In Table 4 and 5, we provide our experimental result. Here, one can observe how our algorithm outperforms the existing works for 10 out of 15 metrics including FILTERS, FCD-TEST, FCD-TESTSF, SNN-TEST, SNN-TESTSF, FRAG-TEST, FRAG-TESTSF, and SCAF-TEST. This highlights the ability of our STGG framework to successfully learn the training distribution.

heading: 5.3 MOLECULAR OPTIMIZATION FOR PENALIZED OCTANOL-WATER PARTITION COEFFICIENT

text:
Finally, we demonstrate the usefulness of our STGG framework for the task of molecular optimization. To this end, we consider the literature standard of maximizing the penalized octanol-water partition coefficient (PLOGP). However, several works (Gao & Coley, 2020; Coley, 2020) have noted how the existing algorithms on this benchmark may not be practical, since PLOGP is ill-defined as a scoring function for molecules; this scoring function may assign high values to “unrealistic” molecules that are unstable and hard to synthesize in practice.
To consider this aspect, we propose a new algorithm which can control the quality of molecules by trading off scores and realistic-ness of molecules. Using this algorithm, we demonstrate how our STGG is capable of generating both (a) high-scoring molecules and (b) realistic molecules with a reasonably high score. At a high-level, we train a conditional generative model pθ(m|γ) under the STGG framework with PLOGP as the condition γ. At the test time, we sample from a high value γ to obtain high-scoring molecules. Such an algorithm is inspired from the recent offline reinforcement learning algorithms (Schmidhuber, 2019; Kumar & Levine, 2020; Chen et al., 2021; Janner et al., 2021). We fully describe our molecular optimization algorithm in Appendix F.
In Table 5 and Figure 6, we report the result of our molecular optimization experiment. We provide additional illustrations of the generated molecules in Appendix G. Here, our STGG model is able to generate molecules with considerably high PLOGP scores outside the training distribution. Furthermore, in Figure 6 and Appendix G, one can observe how increasing γ gradually changes the optimized molecule from realistic structures to large, chain-like, and unrealistic structures.3 Given such results, one may conclude that our STGG combined with the offline optimization algorithm can successfully make a trade-off between high PLOGP and realistic-ness of the generated molecules. However, we also remark that our results do not imply our optimization results to be strictly better than the baselines; we believe it is necessary to develop and incorporate quantitative measures for realistic-ness of molecules to fairly evaluate the molecular optimization algorithms. We believe such a research to be a important future dirction.

heading: 6 CONCLUSION

text:
In this paper, we propose STGG which is the first spanning tree-based framework for the generation of molecules using the Transformer architecture. The key idea of using the spanning tree for graph generation applies to any graph type outside the molecules; we believe such an extension of our work to be both promising and interesting. We also propose an offline algorithm for molecular optimization which allows the trade-off between the high score and the realistic-ness of molecules. We leave more investigation of the newly proposed optimization algorithm as future work.
3This is in agreement with prior works (Shi et al., 2020; Ahn et al., 2020; Luo et al., 2021).

heading: 7 REPRODUCIBILITY STATEMENT

text: We provide explicit description of our algorithm in Algorithm 1, Appendix 2, 3, and 4. We list the hyper-parameters, the hardware used for the experiments, and the data-processing information in Appendix D. We provide illustrations of the molecules generated for the experiments in Figure 6, and Appendix E, G. We submit the full implementation of our STGG framework and the baselines used in our experiments as a supplementary material.

heading: A EXTRACTING SEQUENCE OF DECISIONS FROM A MOLECULAR GRAPH

text:
In this section, we explain our algorithm for finding a sequence of decisions to construct a given molecular graph G = (A,B, E). The high-level idea is to first perform a depth-first search on G to find a spanning tree T = (A,B, ET ) and the corresponding set of residual edges ER = E\ET . Then the algorithm traverses the spanning tree T according to the depth-first search tree while (a) allocating branch start and branch end for vertices with degree higher than two and (b) adding res atom and res bond operations for any vertex covered by a residual edge {a, b} ∈ ER. To this end, we utilize a stack Sdfs that stores the list of vertices in G and branching tokens {"(", ")"} to visit. At each iteration, an element i of the stack Sdfs is popped. If i is a vertex, the algorithm adds the corresponding decision for attach atom and attach bond operations. If the vertex has more than two successors with respect to the spanning tree T , the successors are inserted into the stack Sdfs with surrounding "(" and ")" tokens. If the vertex has only one successor, the successor is inserted into the stack without an additional operation. When the branching tokens {"(", ")"} are popped from the stack, the algorithm adds the corresponding decision value to the sequence of decisions. We describe the full scheme in Algorithm 2.
Algorithm 2 Generating sequence of decisions for a molecular graphs 1: Input: graph G = (A,B, E), atom attributes {xa}a∈A, and bond attributes {xb}b∈B 2: Find a spanning tree T = (A,B, ET ) of G based on depth-first order and set ER ← E \ ET . 3: Initialize an empty sequence of decisions D. 4: Choose the root a ∈ A of T to insert in an empty stack Sbranch. 5: do 6: Pop i from Sbranch. 7: if i ∈ A ∪ B then 8: Append xi to D. . Decision to attach the atom vertex. 9: for j ∈ {j|j ∈ N (i), {i, j} ∈ ER} do . Decisions for residual edges. 10: If i ∈ A, append "*" to D. 11: If i ∈ B, append j to D. 12: Let V denote {j|j ∈ N (i), j 6∈ AT ∪ BT }. . Successors of i in depth-first order. 13: If |V| > 1, insert "(" to Sbranch. . Allocate decision to record pointer vertex. 14: Insert vertices in V to Sbranch. . Allocate successors to visit later. 15: If |V| > 1, insert ")" to Sbranch. . Allocate decision to return to the pointer vertex. 16: if i ∈ {"(", ")"} then 17: Append i to D. 18: while |Sbranch| > 0 19: Output: sequence of decisions D = d1, . . . , dT to reconstruct G.

heading: B ALGORITHMS FOR GRAPH MASKING

text:
To determine whether if a sequence of decisions lead to a valid generation of a molecular graph, we propose an algorithm that outputs a set of valid decisions given the current decision d, stack of pointer vertices S , and list of atom vertices L during execution of Algorithhm 1. We provide the full description in Algorithm 3.
Algorithm 3 Determination of grammar violation 1: Input: current decision d, stack Sbranch, and list Lres. 2: Output: List of candidate decisions D that are valid. 3: if d ∈ Xatom then 4: Set D ← Xbond. . When the atom vertex is followed by bond vertex. 5: Set D ← D ∪ {"("} . When the atom vertex has more than one successors. 6: Set D ← D ∪ {"*"} . When the atom vertex has a neighboring residual edge. 7: If |Sbranch| > 0, set D ← D ∪ {")"}. . The ")" decision appears only when Sbranch is non-empty 8: Set D ← D ∪ {"[eos]"}. . Allow termination. 9: if d ∈ Xbond then 10: Set D ← Xatom. . When the bond vertex is followed by atom vertex. 11: Set D ← D ∪ Lres . When the bond vertex has a neighboring residual edge. 12: if d = "*" then 13: Set D ← Xbond. . When the atom vertex is followed by bond vertex. 14: Set D ← D ∪ {"("} . When the atom vertex has more than one successors. 15: Set D ← D ∪ {"*"} . When the atom vertex has a neighboring residual edge. 16: If |Sbranch| > 0, set D ← D ∪ {")"}. . The ")" decision appears only when Sbranch is non-empty 17: Set D ← D ∪ {"[eos]"}. . Allow termination. 18: if d ∈ Lres then 19: Set D ← {)}. . Residual edge is only constructed at end of each branch. 20: if d = "(" then 21: Set D ← Xbond. . Branch always starts with a bond vertex. 22: Set D ← D ∪ {"[eos]"}. . Allow termination. 23: if d = ")" then 24: Set D ← Xbond ∪ {"(", ")"}. . Branch is followed by start or end of another branch. 25: Set D ← D ∪ {"[eos]"}. . Allow termination.
We also establish theoretical result on how the sequence of decisions generated from Algorithm 3 is always a valid sequence of decisions for Algorithm 1. To this end, we define a valid molecular graph as follows.
Definition 1. A valid molecular graph G = (A,B, E) is a connected bipartite graph where the number of vertices adjacent to any bond vertex b ∈ B is exactly two, i.e., |N (b)| = 2.
Such a definition implies how a molecule should have exactly two atoms connected to a bond. Combined with additional conditions to guarantee the well-behavior of Algorithm 1 on sequence of decisions, we obtain the following result.
Theorem 1. Let G = (A,B, E), S , and L be a graph, a stack of vertices, and a list of vertices being updated by Algorithm 1 and a sequence of decisions d1, . . . , dT . If the sequence of decisions satisfies the criteria defined by Algorithm 3, the following properties are satisfied.
P1 At the t-th step of Algorithm 1, |S| > 0 if dt = ")".
P2 At the t-th step of Algorithm 1, dt ∈ L if dt ∈ A ∪ B.
P3 When dT = [eos], the graph G is a valid molecular graph.
Here, P1 and P2 implies how the operations in Algorithm 1 are well-defined for d1, . . . , dT .
Proof. First, P1 is enforced by the step in Algorithm 3 which forbids the decision value of ")" when the stack S is empty. Next, P2 is enforced by the step selecting decision values from the current list of vertices L. To enforce P3, when dT = [eos], G (a) has to be a connected bipartite graph and (b) the number of vertices adjacent to any bond vertex has to be exactly two. For (a), Algorithm 3 allows the decision of dt ∈ Xatom ∪ L only when the pointer vertex is a bond vertex, i.e., d ∈ Xbond. Similarly, dt ∈ Xbond is allowed only when d ∈ Xatom ∪ {"*", "("}. For (b), the algorithm does not allow adding a bond vertex b ∈ B to the list of vertices L, which is required for any vertex with degree higher than two. Termination is not allowed when there exists a bond vertex with degree smaller than two.

heading: C ALGORITHM FOR VALENCE MASKING

text:
To consider chemical validity of molecules, our framework offers the ability to constrain its generation on molecules that satisfy the valence rule for each atom. That is, the generated graph G = (A,B, E) satisfies the constraint v(xa) ≥ ∑ b∈N (a) o(xb) for every atom vertex a ∈ A where v(xa) denotes the valence of an atom type xa and o(xb) denotes the bond order.
To this end, we propose an algorithm which iteratively updates a record r(a) of available valence for each atom vertex a ∈ A. The key idea is to (a) update the record accordingly for each addition of atom and bond orders and (b) pre-allocate valence for the branch start and res atom operations by the amount of minimum bond order minx∈Xbond o(x). The second part (b) is required since the branch start and res atom operations indicate future bond vertices to be added as a neighbor of the current atom vertex. We provide the full description in Algorithm 4.
Algorithm 4 Determination of valence rule violation
1: Input: Intermediate tree T = (AT ,BT , ET ), current pointer vertex ipoint, previous pointer vertex ĩpoint, current decision d, previous decision d̃, and record r(·) of available valence.
2: Output: Newly updated r and the list D of decisions that violates the valence rule. 3: if d ∈ Xatom then 4: Set r(ipoint)← v(d). . Initialize record by atom valence. 5: Set r(ipoint)← r(ipoint)− o(d̃). . Update record using previously added bond vertex. 6: Set D ← {x|x ∈ Xbond, o(x) > r(ipoint)}. . Reject bond orders higher than the record. 7: if r(ipoint) < minx∈Xbond o(x) then 8: Set D ← D ∪ {"(", "*"}. . Reject decisions requiring minimal amount of valence. 9: if d ∈ Xbond then 10: if d̃ 6= "(" then 11: Set r(̃ipoint)← r(̃ipoint)− o(d). . Update record of previously added atom vertex. 12: else 13: Set r(̃ipoint)← r(̃ipoint)− o(d) + minx∈Xbond o(x).
. Update previously added atom vertex considering pre-allocated valence.
14: Set D ← {x|x ∈ Xatom, v(x) < r(ipoint))}. . Reject atom valence lower than bond order. 15: Set D ← D ∪ {x|x ∈ Lres, r(x) < o(d)−minx∈Xbond o(x)}.
. Reject residual edge candidates with valence lower than bond order.
16: if d = "(" then 17: Set r(ipoint)← r(ipoint)−minx∈Xbond o(x). . Pre-allocate minimum bond order. 18: Set D ← {x|x ∈ Xbond, o(x) > r(ipoint)}. . Reject bond orders higher than the record. 19: if d = ")" then 20: Set D ← ∅ 21: if d = "*" then 22: Set r(ipoint)← r(ipoint)−minx∈Xbond o(x). . Pre-allocate minimum bond order. 23: Set D ← {x|x ∈ Xbond, o(x) > r(ipoint)}. . Reject bond orders higher than the record. 24: if d ∈ Lres then 25: Set r(d)← r(d) + minx∈Xbond o(x)− o(xipoint). . Update record of previously added atom vertex. 26: Set D ← ∅.
Given Algorithm 3 and 4, we establish the following theoretical guarantee. Definition 2. A valid molecular graph G = (A,B, E) satisfies the valency rule if v(xa) ≥∑ b∈N (a) o(xb) for every atom vertex a ∈ A.
Theorem 2. Let G = (A,B, E) be a graph being updated by Algorithm 1 and a sequence of decisions d1, . . . , dT . If the sequence of decisions satisfies the criteria defined by Algorithm 3 and 4, the corresponding graph G satisfies the valency rule.
Proof. To prove the validity of our algorithms, we show that (1) r(a) ≤ v(xa)− ∑ b∈N (a) o(xb) and (2) r(a) > 0 at any time-step applying Algorithm 4 to the sequence of decisions d1, . . . , dT .
For (1), we note that a record of an atom a is initialized as v(xa)− ∑ b∈N (a) o(xb) whenever it is newly added to the graph G by a decision. Furthermore, whenever a new edge is added by attach bond and res bond operation, the corresponding bond order o(xb) is deducted from the record. Importantly, a minimum bond order minx∈Xbond o(x) is also added to the record for a attach bond operation consecutive to a branch start operation or a res bond operation. We note how this does not harm (1) since the minimum bond order has already been deducted (or pre-allocated) by the corresponding branch start and res atom operations, respectively.
For the case of (2), one can observe that Algorithm 4 filters out the atoms and bonds that will deduct the record of the corresponding atom to be negative. This completes proves the correctness of our theorem.
D IMPLEMENTATION DETAILS
In this section, we provide specific details on how we implement the STGG framework for our experiments.
Training detail. For all the experiments, we train the Transformer under STGG framework for 100 epochs with batch size of 128 for all the dataset. We use the AdamW (Loshchilov & Hutter, 2019) optimizer with constant learning rate of 10−4. We use three and six Transformer layers for {QM9, ZINC250K} and MOSES, respectively. The rest of Transformer-related configurations follow that of the original work (Vaswani et al., 2017); we use the attention module with embedding size of 1024 with eight heads, MLP with dimension of 2048, and dropout with probability of 0.1. Using a single Quadro RTX 6000 GPU, it takes approximately three, ten, and 96 hours to fully train the models on QM9, ZINC250K, and MOSES datasets, respectively.
Pre-processing. For the whole dataset, we use the following set of atom vocabularies Xatom: { "CH", "CH2", "CH-", "CH2-", "C", "N-", "NH-", "N", "NH", "N+", "NH+", "NH2+", "NH3+", "O-", "O", "O+", "OH+", "F", "P", "PH", "PH2", "P+", "PH+", "S-", "S", "S+", "SH", "SH+", "Cl", "Br", "I", }. Note that we assign different features for the same atom numbers with different number of explicit hydrogens and formal charges. This allows our algorithm to properly allocate maximum valence for each atom feature. Next, we use the bond vocabularyXbond = {"-", "=", "#"}, corresponding to bond orders of single, double, and triple, respectively. For explicit calculation of the atom valence during molecular construction, we train our models on kekulized molecules, i.e., aromatic bonds are fixed to single or double bonds.

heading: E EXAMPLE OF GENERATED MOLECULES

heading: F OFFLINE OPTIMIZATION OF MOLECULES

text:
In this section, we describe our offline molecular optimization algorithm mainly inspired by existing works in offline reinforcement learning (Schmidhuber, 2019; Chen et al., 2021; Janner et al., 2021) and offline model-based optimization (Kumar & Levine, 2020).
For maximizing a reward function defined on a molecule, our algorithm consists of two simple steps. First, our offline optimization algorithm trains a conditional generative model pθ(m|γ) where m is the molecule and γ is the reward function evaluated on the offline dataset of molecules. Next, the reward-conditional generative model samples highly-rewarding molecules by generation conditioned on high values of γ. In particular, we set the value of γ to extrapolate outside the training dataset. Based on the highly expressive power of Transformer architecture, our algorithm can successfully generate highly-rewarding molecules.

heading: G ADDITIONAL EXPERIMENTAL RESULTS ON MOLECULAR OPTIMIZATION

heading: H COMPARISON WITH CG-VAE

text:
In this section, we additionally compare our STGG framework with the CG-VAE model (Liu et al., 2018), which is another atom-by-atom graph generative model that allows masking out the action space to generate molecules satisfying the valence rules. Compared to Table 2, we additionally use the FCD, SNN, Frag, Scaf metrics to measure the faithfulness of the generative models for learning the underlying distribution of molecules. Note that the VALID metric used in Table 2 is insufficient to compare the faithfulness of STGG and CG-VAE since they both enjoy the guarantee to generate molecules satisfying the valence rule.
In Table 6 and 7, one can observe that our algorithm highly outperforms the CG-VAE in terms of faithfully learning the underlying distribution of molecules at the cost of relatively lower UNIQUE. For example, FCD score of our STGG for learning the ZINC dataset is 0.2775 while that of the CG-VAE is 11.33. This highlights the expressive power of our STGG framework.