Learning to Synthesize Programs as Interpretable and Generalizable Policies

1 Introduction . Recently , deep reinforcement learning ( DRL ) methods have demonstrated encouraging performance on a variety of domains such as outperforming humans in complex games [ 1–4 ] or controlling robots [ 5–11 ] . Despite the recent progress in the field , acquiring complex skills through trial and error still remains challenging and these neural network policies often have difficulty generalizing to novel scenarios . Moreover , such policies are not interpretable to humans and therefore are difficult to debug when these challenges arise . To address these issues , a growing body of work aims to learn programmatic policies that are structured in more interpretable and generalizable representations such as decision trees [ 12 ] , statemachines [ 13 ] , and programs described by domain-specific programming languages [ 14 , 15 ] . Yet , the programmatic representations employed in these works are often limited in expressiveness due to constraints on the policy spaces . For example , decision tree policies are incapable of naïvely generating repetitive behaviors , state machine policies used in [ 13 ] are computationally complex to ∗Contributed equally . †Work partially done as a visiting scholar at USC . ‡AI Advisor at NAVER AI Lab . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . scale to policies representing diverse behaviors , and the programs of [ 14 , 15 ] are constrained to a set of predefined program templates . On the other hand , program synthesis works that aim to represent desired behaviors using flexible domain-specific programs often require extra supervision such as input/output pairs [ 16–20 ] or expert demonstrations [ 21 , 22 ] , which can be difficult to obtain . In this paper , we present a framework to instead synthesize human-readable programs in an expressive representation , solely from rewards , to solve tasks described by Markov Decision Processes ( MDPs ) . Specifically , we represent a policy using a program composed of control flows ( e.g . if/else and loops ) and an agent ’ s perceptions and actions . Our programs can flexibly compose behaviors through perception-conditioned loops and nested conditional statements . However , composing individual program tokens ( e.g . if , while , move ( ) ) in a trial-and-error fashion to synthesize programs that can solve given MDPs can be extremely difficult and inefficient . To address this problem , we propose to first learn a latent program embedding space where nearby latent programs correspond to similar behaviors and allows for smooth interpolation , together with a program decoder that can decode a latent program to a program consisting of a sequence of program tokens . Then , when a task is given , this embedding space allows us to iteratively search over candidate latent programs to find a program that induces desired behavior to maximize the reward . Specifically , this embedding space is learned through reconstruction of randomly generated programs and the behaviors they induce in the environment in an unsupervised manner . Once learned , the embedding space can be reused to solve different tasks without retraining . To evaluate the proposed framework , we consider the Karel domain [ 23 ] , featuring an agent navigating through a gridworld and interacting with objects to solve tasks such as stacking and navigation . The experimental results demonstrate that the proposed framework not only learns to reliably synthesize task-solving programs but also outperforms program synthesis and deep RL baselines . In addition , we justify the necessity of the proposed two-stage learning scheme as well as conduct an extensive analysis comparing various approaches for learning the latent program embedding spaces . Finally , we perform experiments which highlight that the programs produced by our proposed framework can both generalize to larger state spaces and unseen state configurations as well as be interpreted and edited by humans to improve their task performance . 2 Related Work . Neural program induction and synthesis . Program induction methods [ 20 , 24–36 ] aim to implicitly induce the underlying programs to mimic the behaviors demonstrated in given task specifications such as input/output pairs or expert demonstrations . On the other hand , program synthesis methods [ 16– 19 , 21 , 37–58 ] explicitly synthesize the underlying programs and execute the programs to perform the tasks from task specifications such input/output pairs , demonstrations , language instructions . In contrast , we aim to learn to synthesize programs solely from reward described by an MDP without other task specifications . Similarly to us , a two-stage synthesis method is proposed in [ 46 ] . Yet , the task is to match truth tables for given test programs rather than solve MDPs . Their first stage requires the entire ground-truth table for each program synthesized during training , which is infeasible to apply to our problem setup ( i.e . synthesizing imperative programs for solving MDPs ) . Learning programmatic policies . Prior works have also addressed the problem of learning programmatic policies [ 59–61 ] . Bastani et al . [ 12 ] learns a decision tree as a programmatic policy for pong and cartpole environments by imitating an oracle neural policy . However , decision trees are incapable of representing repeating behaviors on their own . Silver et al . [ 49 ] addresses this by including a loop-style token for their decision tree policy , though it is still not as expressive as synthesized loops . Inala et al . [ 13 ] learns programmatic policies as finite state machines by imitating a teacher policy , although finite state machine complexity can scale quadratically with the number of states , making them difficult to scale to more complex behaviors . Another line of work instead synthesizes programs structured in Domain Specific Languages ( DSLs ) , allowing humans to design tokens ( e.g . conditions and operations ) and control flows ( e.g . while loops , if statements , reusable functions ) to induce desired behaviors and can produce human interpretable programs . Verma et al . [ 14 , 15 ] distill neural network policies into programmatic policies . Yet , the initial programs are constrained to a set of predefined program templates . This significantly limits the scope of synthesizable programs and requires designing such templates for each task . In contrast , our method can synthesize diverse programs , without templates , which can flexibly represent the complex behaviors required to solve various tasks . 3 Problem Formulation We are interested in learning to synthesize a program structured in a given DSL that can be executed to solve a given task described by an MDP , purely from reward . In this section , we formally define our definition of a program and DSL , tasks described by MDPs , and the problem formulation . Program and Domain Specific Language . The programs , or programmatic policies , considered in this work are defined based on a DSL as shown in Figure 1 . The DSL consists of control flows and an agent ’ s perceptions and actions . A perception indi- cates circumstances in the environment ( e.g . frontIsClear ( ) ) that can be perceived by an agent , while an action defines a certain behavior that can be performed by an agent ( e.g . move ( ) , turnLeft ( ) ) . Control flow includes if/else statements , loops , and boolean/logical operators to compose more sophisticated conditions . A policy considered in this work is described by a program ρ which is executed to produce a sequence of actions given perceptions from the environment . MDP . We consider finite-horizon discounted MDPs with initial state distribution µ ( so ) and discount factor γ . For a fixed sequence { ( s0 , a0 ) , ... , ( st , at ) } of states and actions obtained from a rollout of a given policy , the performance of the policy is evaluated based on a discounted return ∑T t=0 γ trt , where T is the horizon of the episode and rt = R ( st , at ) the reward function . Objective . Our objective is maxρ Ea∼EXEC ( ρ ) , s0∼µ [ ∑T t=0 γ trt ] , where EXEC returns the actions induced by executing a program policy ρ in the environment . Note that one can view this objective as a special case of the standard RL objective , where the policy is represented as a program which follows the grammar of the DSL and the policy rollout is obtained by executing the program . 4 Approach . Our goal is to develop a framework that can synthesize a program ( i.e . a programmatic policy ) structured in a given DSL that can be executed to solve a task of interest . This requires the ability to synthesize a program that is not only valid for execution ( e.g . grammatically correct ) but also describes desired behaviors for solving the task from only the reward . Yet , learning to synthesize such a program from scratch for every new task can be difficult and inefficient . To this end , we propose our Learning Embeddings for lAtent Program Synthesis framework , dubbed LEAPS , as illustrated in Figure 2 . LEAPS first learns a latent program embedding space that continuously parameterizes diverse behaviors and a program decoder that decodes a latent program to a program consisting of a sequence of program tokens . Then , when a task is given , we iteratively search over this embedding space and decode each candidate latent program using the decoder to find a program that maximizes the reward . This two-stage learning scheme not only enables learning to synthesize programs to acquire desired behaviors described by MDPs solely from reward , but also allows reusing the learned embedding space to solve different tasks without retraining . In the rest of this section , we describe how we construct the model and our learning objectives for the latent program embedding space in Section 4.1 . Then , we present how a program that describes desired behaviors for a given task can be found through a search algorithm in Section 4.2 . 4.1 Learning a Program Embedding Space . To learn a latent program embedding space , we propose to train a variational autoencoder ( VAE ) [ 62 ] that consists of a program encoder qφ which encodes a program ρ to a latent program z and a program decoder pθ which reconstructs the program from the latent . Specifically , the VAE is trained through reconstruction of randomly generated programs and the behaviors they induce in the environment in an unsupervised manner . Architectural details are listed in Section L.6 . Since we aim to iteratively search over the learned embedding space to achieve certain behaviors when a task is given , we want this embedding space to allow for smooth behavior interpolation ( i.e . programs that exhibit similar behaviors are encoded closer in the embedding space ) . To this end , we propose to train the model by optimizing the following three objectives . 4.1.1 Program Reconstruction . To learn a program embedding space , we train a program encoder qφ and a program decoder pθ to reconstruct programs composed of sequences of program tokens . Given an input program ρ consisting of a sequence of program tokens , the encoder processes the input program one token at a time and produces a latent program embedding z . Then , the decoder outputs program tokens one by one from the latent program embedding z to synthesize a reconstructed program ρ̂ . Both the encoder and the decoder are recurrent neural networks and are trained to optimize the β-VAE [ 63 ] loss : LPθ , φ ( ρ ) = −Ez∼qφ ( z|ρ ) [ log pθ ( ρ|z ) ] + βDKL ( qφ ( z|ρ ) ‖pθ ( z ) ) . ( 1 )

Summary: - Authors use a two-stage method for program synthesis: (1) They   create an embedding space capturing the semantics of the program   space; (2) They search through the embedding space for a program   that solves the task in hand. - They use variational autoencoder to encode programs, where an   encoder maps a sequence of program tokens into a latent space, while   a decoder maps from the latent space into the original program. They   use a combination of loss functions to ensure that not only   syntactically similar programs are close to each other in the latent   space but also semantically close programs. Loss function consists   of a composition of 3 losses. (A) Classical beta VAE loss which   helps that syntactically close programs are close to each other. (B)   Program behavior reconstruction loss which helps that semantically   similar programs are close to each other in the embedding space. For   example, two syntactically different programs can make same actions   under same situations. This loss is defined by difference between   execution traces of original and reconstructed program. (C) Latent   behavior reconstruction loss. This loss is constructed by learning a   policy to predict correct action based on the latent program   embedding. This allows for backpropagating gradient through the   policy and the encoder to ensure better quality of the embeddings. - After training the variational autoencoder, authors use a cross   entropy method to find a program with maximal performance using the   latent space. They sample the latent space distribution of programs,   decode the latent program into an actual program, execute the   program to obtain corresponding rewards, and repeat the process. - Authors did a nice ablation study where they measured effectiveness   of their individual design decisions: evaluating each of the loss   function components, as well as the effectiveness of the latent   space search method for finding the best performing program. They   compared they method to other methods, deep reinforcement learning   (DRL) based as well as Viper which is a decision tree policy   obtained by mimicking a DRL agent.