Single Episode Policy Transfer in Reinforcement Learning

1 INTRODUCTION . One salient feature of human intelligence is the ability to perform well in a single attempt at a new task instance , by recognizing critical characteristics of the instance and immediately executing appropriate behavior based on experience in similar instances . Artificial agents must do likewise in applications where success must be achieved in one attempt and failure is irreversible . This problem setting , single episode transfer , imposes a challenging constraint in which an agent experiences—and is evaluated on—only one episode of a test instance . As a motivating example , a key challenge in precision medicine is the uniqueness of each patient ’ s response to therapeutics ( Hodson , 2016 ; Bordbar et al. , 2015 ; Whirl-Carrillo et al. , 2012 ) . Adaptive therapy is a promising approach that formulates a treatment strategy as a sequential decision-making problem ( Zhang et al. , 2017 ; West et al. , 2018 ; Petersen et al. , 2019 ) . However , heterogeneity among instances may require explicitly accounting for factors that underlie individual patient dynamics . For example , in the case of adaptive therapy for sepsis ( Petersen et al. , 2019 ) , predicting patient response prior to treatment is not possible . However , differences in patient responses can be observed via blood measurements very early after the onset of treatment ( Cockrell and An , 2018 ) . As a first step to address single episode transfer in reinforcement learning ( RL ) , we propose a general algorithm for near-optimal test-time performance in a family of environments where differences in dynamics can be ascertained early during an episode . Our key idea is to train an inference model and a probe that together achieve rapid inference of latent variables—which account for variation in a family of similar dynamical systems—using a small fraction ( e.g. , 5 % ) of the test episode , then deploy a universal policy conditioned on the estimated parameters for near-optimal control on the new instance . Our approach combines the advantages of robust transfer and adaptation-based transfer , as we learn a single universal policy that requires no further training during test , but which is adapted to the new environment by conditioning on an unsupervised estimation of new latent dynamics . In contrast to methods that quickly adapt or train policies via gradients during test but assume access to multiple test rollouts and/or dense rewards ( Finn et al. , 2017 ; Killian et al. , 2017 ; Rakelly et al. , 2019 ) , we explicitly optimize for performance in one test episode without accessing the reward function at test time . Hence our method applies to real-world settings in which rewards during test are highly delayed or even completely inaccessible—e.g. , a reward that depends on physiological factors that are accessible only in simulation and not from real patients . We also consider computation time a crucial factor for real-time application , whereas some existing approaches require considerable computation during test ( Killian et al. , 2017 ) . Our algorithm builds on variational inference and RL as submodules , which ensures practical compatibility with existing RL workflows . Our main contribution is a simple general algorithm for single episode transfer in families of environments with varying dynamics , via rapid inference of latent variables and immediate execution of a universal policy . Our method attains significantly higher cumulative rewards , with orders of magnitude faster computation time during test , than the state-of-the-art model-based method ( Killian et al. , 2017 ) , on benchmark high-dimensional domains whose dynamics are discontinuous and continuous in latent parameters . We also show superior performance over optimization-based meta-learning and favorable performance versus baselines for robust transfer . 2 SINGLE EPISODE TRANSFER IN RL : PROBLEM SETUP . Our goal is to train a model that performs close to optimal within a single episode of a test instance with new unknown dynamics . We formalize the problem as a family ( S , A , T , R , γ ) , where ( S , A , R , γ ) are the state space , action space , reward function , and discount of an episodic Markov decision process ( MDP ) . Each instance of the family is a stationary MDP with transition function Tz ( s′|s , a ) ∈ T . When a set Z of physical parameters determines transition dynamics ( Konidaris and Doshi-Velez , 2014 ) , each Tz has a hidden parameter z ∈ Z that is sampled once from a distribution PZ and held constant for that instance . For more general stochastic systems whose modes of behavior are not easily attributed to physical parameters , Z is induced by a generative latent variable model that indirectly associates each Tz to a latent variable z learned from observed trajectory data . We refer to “ latent variable ” for both cases , with the clear ontological difference understood . Depending on application , Tz can be continuous or discontinuous in z . We strictly enforce the challenging constraint that latent variables are never observed , in contrast to methods that use known values during training ( Yu et al. , 2017 ) , to ensure the framework applies to challenging cases without prior knowledge . This formulation captures a diverse set of important problems . Latent space Z has physical meaning in systems where Tz is a continuous function of physical parameters ( e.g. , friction and stiffness ) with unknown values . In contrast , a discrete set Z can induce qualitatively different dynamics , such as a 2D navigation task where z ∈ { 0 , 1 } decides if the same action moves in either a cardinal direction or its opposite ( Killian et al. , 2017 ) . Such drastic impact of latent variables may arise when a single drug is effective for some patients but causes serious side effects for others ( Cockrell and An , 2018 ) . Training phase . Our training approach is fully compatible with RL for episodic environments . We sample many instances , either via a simulator with controllable change of instances or using off-policy batch data in which demarcation of instances—but not values of latent variables—is known , and train for one or more episodes on each instance . While we focus on the case with known change of instances , the rare case of unknown demarcation can be approached either by preprocessing steps such as clustering trajectory data or using a dynamic variant of our algorithm ( Appendix C ) . Single test episode . In contrast to prior work that depend on the luxury of multiple experience rollouts for adaptation during test time ( Doshi-Velez and Konidaris , 2016 ; Killian et al. , 2017 ; Finn et al. , 2017 ; Rakelly et al. , 2019 ) , we introduce the strict constraint that the trained model has access to—and is evaluated on—only one episode of a new test instance . This reflects the need to perform near-optimally as soon as possible in critical applications such as precision medicine , where an episode for a new patient with new physiological dynamics is the entirety of hospitalization . 3 SINGLE EPISODE POLICY TRANSFER . We present Single Episode Policy Transfer ( SEPT ) , a high-level algorithm for single episode transfer between MDPs with different dynamics . The following sections discuss specific design choices in SEPT , all of which are combined in synergy for near-optimal performance in a single test episode . 3.1 POLICY TRANSFER THROUGH LATENT SPACE . Our best theories of natural and engineered systems involve physical constants and design parameters that enter into dynamical models . This physicalist viewpoint motivates a partition for transfer learning in families of MDPs : 1. learn a representation of latent variables with an inference model that rapidly encodes a vector ẑ of discriminative features for a new instance ; 2. train a universal policy π ( a|s , z ) to perform near-optimally for dynamics corresponding to any latent variable in Z ; 3. immediately deploy both the inference model and universal policy on a given test episode . To build on the generality of model-free RL , and for scalability to systems with complex dynamics , we do not expend computational effort to learn a model of Tz ( s′|s , a ) , in contrast to model-based approaches ( Killian et al. , 2017 ; Yao et al. , 2018 ) . Instead , we leverage expressive variational inference models to represent latent variables and provide uncertainty quantification . In domains with ground truth hidden parameters , a latent variable encoding is the most succinct representation of differences in dynamics between instances . As the encoding ẑ is held constant for all episodes of an instance , a universal policy π ( a|s , z ) can either adapt to all instances when Z is finite , or interpolate between instances when Tz is continuous in z ( Schaul et al. , 2015 ) . Estimating a discriminative encoding for a new instance enables immediate deployment of π ( a|s , z ) on the single test episode , bypassing the need for further fine-tuning . This is critical for applications where further training complex models on a test instance is not permitted due to safety concerns . In contrast , methods that do not explicitly estimate a latent representation of varied dynamics must use precious experiences in the test episode to tune the trained policy ( Finn et al. , 2017 ) . In the training phase , we generate an optimized1 dataset D : = { τ i } Ni=1 of short trajectories , where each τ i : = ( si1 , a i 1 , . . . , s i Tp , aiTp ) is a sequence of early state-action pairs at the start of episodes of instance Ti ∈ T ( e.g . Tp = 5 ) . We train a variational auto-encoder , comprising an approximate posterior inference model qφ ( z|τ ) that produces a latent encoding ẑ from τ and a parameterized generative model pψ ( τ |z ) . The dimension chosen for ẑ may differ from the exact true dimension when it exists but is unknown ; domain knowledge can aid the choice of dimensionality reduction . Because dynamics of a large variety of natural systems are determined by independent parameters ( e.g. , coefficient of contact friction and Reynolds number can vary independently ) , we consider a disentangled latent representation where latent units capture the effects of independent generative parameters . To this end , we bring β-VAE ( Higgins et al. , 2017 ) into the context of families of dynamical systems , choosing an isotropic unit Gaussian as the prior and imposing the constraint DKL ( qφ ( z|τ i ) ‖p ( z ) ) < . The β-VAE is trained by maximizing the variational lower bound L ( ψ , φ ; τ i ) for each τ i across D : max ψ , φ log pψ ( τ i ) ≥ L ( ψ , φ ; τ i ) : = −βDKL ( qφ ( z|τ i ) ‖p ( z ) ) + Eqφ ( z|τ i ) [ log pψ ( τ i|z ) ] ( 1 ) This subsumes the VAE ( Kingma and Welling , 2014 ) as a special case ( β = 1 ) , and we refer to both as VAE in the following . Since latent variables only serve to differentiate among trajectories that arise from different transition functions , the meaning of latent variables is not affected by isometries and hence the value of ẑ by itself need not have any simple relation to a physically meaningful z even when one exists . Only the partition of latent space is important for training a universal policy . Earlier methods for a family of similar dynamics relied on Bayesian neural network ( BNN ) approximations of the entire transition function st+1 ∼ T̂ ( BNN ) z ( st , at ) , which was either used to perform computationally expensive fictional rollouts during test time ( Killian et al. , 2017 ) or used indirectly to further optimize a posterior over z ( Yao et al. , 2018 ) . Our use of variational inference is more economical : the encoder qφ ( z|τ ) can be used immediately to infer latent variables during test , while the decoder pψ ( τ |z ) plays a crucial role for optimized probing in our algorithm ( see Section 3.3 ) . In systems with ground truth hidden parameters , we desire two additional properties . The encoder should produce low-variance encodings , which we implement by minimizing the entropy of qφ ( z|τ ) : min φ H ( qφ ( z|τ ) ) : = − ∫ z qφ ( z|τ ) log qφ ( z|τ ) dz = D 2 log ( 2π ) + 1 2 D∑ d=1 ( 1 + log σ2d ) ( 2 ) under a diagonal Gaussian parameterization , where σ2d = Var ( qφ ( z|τ ) ) and dim ( z ) = D. We add −H ( qφ ( z|τ ) ) as a regularizer to equation 1 . Second , we must capture the impact of z on higher- 1In the sense of machine teaching , as explained fully in Section 3.3 order dynamics . While previous work neglects the order of transitions ( st , at , st+1 ) in a trajectory ( Rakelly et al. , 2019 ) , we note that a single transition may be compatible with multiple instances whose differences manifest only at higher orders . In general , partitioning the latent space requires taking the ordering of a temporally-extended trajectory into account . Therefore , we parameterize our encoder qφ ( z|τ ) using a bidirectional LSTM—as both temporal directions of ( st , at ) pairs are informative—and we use an LSTM decoder pψ ( τ |z ) ( architecture in Appendix E.2 ) . In contrast to embedding trajectories from a single MDP for hierarchical learning ( Co-Reyes et al. , 2018 ) , our purpose is to encode trajectories from different instances of transition dynamics for optimal control .

Summary: This paper presents a strategy for single-trajectory transfer of a reinforcement learned policy.  They follow a typical approach in the few-shot supervised-learning community, of assuming that the plausible set of solutions may be modelled as a much lower dimensional latent variable, and then try to quickly infer that latent variable at test time.  In this case, the latent variable is ‘Z’.