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: The main contribution of this paper is to learn a universal policy that is able to perform near-optimally on test tasks with transition dynamics that were never observed during training. This is achieved by using a "probe policy" to generate short trajectories that are then used to learn a latent encoding to categorise the transition dynamics of the current task. The universal policy is then conditions on both the state and this encoding so that the learned policy can perform well on tasks with different dynamics.