Learning Calibratable Policies using Programmatic Style-Consistency

We study the important and challenging problem of controllable generation of long-term sequential behaviors . Solutions to this problem would impact many applications , such as calibrating behaviors of AI agents in games or predicting player trajectories in sports . In contrast to the well-studied areas of controllable generation of images , text , and speech , there are significant challenges that are unique to or exacerbated by generating long-term behaviors : how should we specify the factors of variation to control , and how can we ensure that the generated temporal behavior faithfully demonstrates diverse styles ? In this paper , we leverage large amounts of raw behavioral data to learn policies that can be calibrated to generate a diverse range of behavior styles ( e.g. , aggressive versus passive play in sports ) . Inspired by recent work on leveraging programmatic labeling functions , we present a novel framework that combines imitation learning with data programming to learn style-calibratable policies . Our primary technical contribution is a formal notion of style-consistency as a learning objective , and its integration with conventional imitation learning approaches . We evaluate our framework using demonstrations from professional basketball players and agents in the MuJoCo physics environment , and show that our learned policies can be accurately calibrated to generate interesting behavior styles in both domains . 1 INTRODUCTION . The widespread availability of recorded tracking data is enabling the study of complex behaviors in many domains , including sports ( Chen et al. , 2016a ; Le et al. , 2017 ; Zhan et al. , 2019 ) , video games ( Kurin et al. , 2017 ; Broll et al. , 2019 ) , laboratory animals ( Eyjolfsdottir et al. , 2014 ; 2017 ; Johnson et al. , 2016 ) , facial expressions ( Suwajanakorn et al. , 2017 ; Taylor et al. , 2017 ) , commonplace activities such as cooking ( Nishimura et al. , 2019 ) , and driving ( Bojarski et al. , 2016 ; Chang et al. , 2019 ) . The tracking data is often obtained from multiple experts and can exhibit very diverse styles ( e.g. , aggressive versus passive play in sports ) . Our work is motivated by the opportunity to maximally leverage these datasets by cleanly extracting such styles in addition to modeling the raw behaviors . Our goal is to train policies that can be controlled , or calibrated , to produce different behavioral styles inherent in the demonstration data . For example , Figure 1a depicts demonstrations from real basketball players with variations of many types , including movement speed , desired destinations , tendencies for long versus short passes , and curvature of movement routes , amongst many others . A calibratable policy would be able to generate trajectories consistent with various styles , such as low movement speed as in Figure 1b , or approach the basket as in Figure 1c , or to both styles simultaneously as in Figure 1d . Importantly , we aim to train a single policy that can generate behaviors calibrated across multiple styles . Having such policies would empower many downstream tasks , including behavior discovery ( Eyjolfsdottir et al. , 2014 ) , realistic simulations ( Le et al. , 2017 ) , virtual agent design ( Broll et al. , 2019 ) , and counterfactual behavioral reasoning ( Zhan et al. , 2019 ) . We focus on three research questions . The first question is strategic : what systematic form of domain knowledge can we leverage to quickly and cleanly extract style information from raw behavioral data ? The second question is formulaic : how can we formalize the learning objective to encourage learning style-calibratable policies ? The third question is algorithmic : how do we design practical learning approaches that reliably optimize the learning objective ? To address these challenges , we present a novel framework inspired by data programming ( Ratner et al. , 2016 ) , a paradigm in weak supervision that utilizes automated labeling procedures , called labeling functions , to learn without ground-truth labels . In our setting , labeling functions enable domain experts to quickly translate domain knowledge of diverse styles into programmatically generated style annotations . For instance , it is trivial to write programmatic labeling functions for the two styles—speed and destination—depicted in Figure 1 . Labeling functions also motivate a metric for learning , which we call programmatic style-consistency , to evaluate calibration of policies : rollouts generated for a specific style should return the same style label when fed to the labeling function . Finally , our framework is generic and is easily integrated into conventional imitation learning approaches . To summarize , our contributions are : • We propose a novel framework for learning policies calibrated to diverse behavior styles . • Our framework allows users to express styles as labeling functions , which can be quickly applied to programmatically produce a weak signal of style labels . • Our framework introduces style-consistency as a metric to evaluate calibration to styles . • We present an algorithm to learn calibratable policies that maximize style-consistency of the generated behaviors , and validate it in basketball and simulated physics environments . 2 BACKGROUND : IMITATION LEARNING USING TRAJECTORY VAES . Since our focus is on learning style-calibratable generative policies , for simplicity we develop our approach with the basic imitation learning paradigm of behavioral cloning using trajectory variational autoencoders , which we describe here . Interesting future directions include composing our approach with more advanced imitation learning approaches as well as with reinforcement learning . Notation . Let S andA denote the environment state and action spaces . At each timestep t , an agent observes state st ∈ S and executes action at ∈ A using a policy π : S → A . The environment then transitions to the next state st+1 according to a ( typically unknown ) dynamics function f : S ×A → S. For the rest of this paper , we assume f is deterministic ; a modification of our approach for stochastic f is included in Appendix B . A trajectory τ is a sequence of T state-action pairs and the last state : τ = { ( st , at ) } Tt=1 ∪ { sT+1 } . Let D be a set of N trajectories collected from expert demonstrations . In our experiments , each trajectory in D has the same length T , but in general this does not need to be the case . Learning objective . We begin with the basic imitation learning paradigm of behavioral cloning ( Syed & Schapire , 2008 ) . The goal is to learn a policy that behaves like the pre-collected demonstrations : π∗ = arg min π Eτ∼D [ Limitation ( τ , π ) ] , ( 1 ) where Limitation is a loss function that quantifies the mismatch between the actions chosen by π and those in the demonstrations . Since we are primarily interested in probabilistic or generative policies , we typically use ( variations of ) negative log-likelihood : L ( τ , π ) = ∑T t=1− log π ( at|st ) , where π ( at|st ) is the probability of π choosing action at in state st. Trajectory Variational Autoencoders . A common model choice for instantiating π is the trajectory variational autoencoder ( TVAE ) , which is a sequential generative model built on top of variational autoencoders ( Kingma & Welling , 2014 ) , and have been shown to work well in a range of generative policy learning settings ( Wang et al. , 2017 ; Ha & Eck , 2018 ; Co-Reyes et al. , 2018 ) . In its simplest form , a TVAE introduces a latent variable z ( also called a trajectory embedding ) with prior distribution p , an encoder network qφ , and a policy decoder πθ . Its imitation learning objective is : Ltvae ( τ , πθ ; qφ ) = Eqφ ( z|τ ) [ T∑ t=1 − log πθ ( at|st , z ) ] +DKL ( qθ ( z|τ ) ||p ( z ) ) . ( 2 ) The main shortcoming of TVAEs and related approaches , which we address in Sections 3 & 4 , is that the resulting policies can not be easily calibrated to generate specific styles of behavior . For instance , the goal of the trajectory embedding z is to capture all the styles that exist in the expert demonstrations , but there is no guarantee that the embeddings cleanly encode the desired styles in a calibrated way . Previous work has largely relied on unsupervised learning techniques that either require significant domain knowledge ( Le et al. , 2017 ) , or have trouble scaling to complex styles commonly found in real-world applications ( Wang et al. , 2017 ; Li et al. , 2017 ) . 3 PROGRAMMATIC STYLE-CONSISTENCY . Building upon the basic setup in Section 2 , we focus on the setting where the demonstrations D contain diverse behavior styles . To start , let y ∈ Y denote a single style label ( e.g. , speed or destination , as shown in Figure 1 ) . Our goal is to learn a policy π that can be explicitly calibrated to y , i.e. , trajectories generated by π ( ·|y ) should match the demonstrations in D that exhibit style y . Obtaining style labels can be expensive using conventional annotation methods , and unreliable using unsupervised approaches . We instead utilize easily programmable labeling functions that automatically produce style labels , described next . We then formalize a notion of style-consistency as a learning objective , and in Section 4 describe a practical learning approach . Labeling functions . Introduced in the data programming paradigm ( Ratner et al. , 2016 ) , labeling functions programmatically produce weak and noisy labels to learn models on otherwise unlabeled datasets . A significant benefit is that labeling functions are often simple scripts that can be quickly applied to the dataset , which is much cheaper than manual annotations and more reliable than unsupervised methods . In our framework , we study behavior styles that can be represented as labeling functions , which we denote λ , that map trajectories τ to style labels y . A simple example is : λ ( τ ) = 1 { ‖sT+1 − s1‖2 > c } , ( 3 ) which distinguishes between trajectories with large ( greater than a threshold c ) versus small total displacement . We experiment with a range of labeling functions , as described in Section 6 . Multiple labeling functions can be provided at once , possibly from multiple users . Many behavior styles used in previous work can be represented as labeling functions , e.g. , agent speed ( Wang et al. , 2017 ) . We use trajectory-level labels λ ( τ ) in our experiments , but in general labeling functions can be applied on subsequences λ ( τt : t+h ) to obtain per-timestep labels . We can efficiently annotate datasets using labeling functions , which we denote as λ ( D ) = { ( τi , λ ( τi ) ) } Ni=1 . Our goal can now be phrased as : given λ ( D ) , train a policy π : S × Y 7→ A such that π ( ·|y ) is calibrated to styles y found in λ ( D ) . Style-consistency . A key insight in our work is that labeling functions naturally induce a metric for calibration . If a policy π ( ·|y ) is calibrated to λ , we would expect the generated behaviors to be consistent with the label . So , we expect the following loss to be small : Ey∼p ( y ) , τ∼π ( ·|y ) [ Lstyle ( λ ( τ ) , y ) ] , ( 4 ) where p ( y ) is a prior over the style labels , and τ is obtained by executing the style-conditioned policy in the environment . Lstyle is thus a disagreement loss over labels that is minimized at λ ( τ ) = y , e.g. , Lstyle ( λ ( τ ) , y ) = 1 { λ ( τ ) 6= y } for categorical labels . We refer to ( 4 ) as the style-consistency loss , and say that π ( ·|y ) is maximally calibrated to λ when ( 4 ) is minimized . Our full learning objective incorporating ( 4 ) with ( 1 ) is : π∗ = argmin π E ( τ , λ ( τ ) ) ∼λ ( D ) [ Limitation ( τ , π ( · | λ ( τ ) ) ) ] + Ey∼p ( y ) , τ∼π ( ·|y ) [ Lstyle ( λ ( τ ) , y ) ] . ( 5 ) The simplest choice for the prior distribution p ( y ) is the marginal distribution of styles in λ ( D ) . The first term in ( 5 ) is a standard imitation learning objective and can be tractably estimated using λ ( D ) . To enforce style-consistency with the second term , conceptually we need to sample several y ∼ p ( y ) , then several rollouts τ ∼ π ( · | y ) from the current policy , and query the labeling function for each of them . Furthermore , if λ is a non-differentiable function defined over the entire trajectory , as is the case in ( 3 ) , then we can not simply backpropagate the style-consistency loss . In Section 4 , we introduce differentiable approximations to more easily optimize the challenging objective in ( 5 ) . Multiple styles . Our notion of style-consistency can be easily extended to simultaneously optimize for multiple styles . Suppose we have M labeling functions { λi } Mi=1 and corresponding label spaces { Yi } Mi=1 . Let λ denote ( λ1 , . . . , λM ) and y denote ( y1 , . . . , yM ) . Then style-consistency becomes : Ey∼p ( y ) , τ∼π ( ·|y ) [ M∑ i=1 Lstylei ( λi ( τ ) , yi ) ] . ( 6 ) Note that style-consistency is optimized when the generated trajectory agrees with all labeling functions . Although this can be very challenging to achieve , it describes the most desirable outcome , i.e . π ( ·|y ) is a policy that can be calibrated to all styles simultaneously .

Summary: This paper proposes a method to train style-conditional policies. The method has three components, a dynamics model, a labeling procedure and its approximation, and a policy. The model is trained on real trajectories, while the policy is trained on both real and simulated data coming from the model. The policy is trained to both imitate and be sensitive to the style labelling of states.