PerSim: Data-Efficient Offline Reinforcement Learning with Heterogeneous Agents via Personalized Simulators

1 Introduction . Reinforcement learning ( RL ) coupled with expressive deep neural networks has now become a generic yet powerful solution for learning complex decision-making policies for an agent of interest ; it provides the key algorithmic foundation underpinning recent successes such as game solving [ 34 , 45 , 44 ] and robotics [ 30 , 22 ] . However , many state-of-the-art RL methods are data hungry and require the ability to query samples at will , which is infeasible for numerous settings such as healthcare , autonomous driving , and socio-economic systems . As a result , there has been a rapidly growing literature on “ offline RL ” [ 31 , 24 , 32 , 15 ] , which focuses on leveraging existing datasets to learn decision-making policies . Within offline RL , we consider a regime of severe data scarcity : there are multiple agents and for each agent , we only observe a single historical trajectory generated under an unknown , potentially sub-optimal policy ; further , the agents are heterogeneous , i.e. , each agent has unique state transition dynamics . Importantly , the characteristics of the agents that make their transition dynamics heterogeneous are latent . Using such limited offline data to simultaneously learn a good “ personalized ” policy for each agent is a challenging setting that has received limited attention . Below , we use a prevalent example 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . from healthcare to motivate and argue that tackling such a challenge is an important and necessary step towards building personalized decision-making engines via RL in a variety of important applications . A Motivating Example . Consider a pre-existing clinical dataset of patients ( agents ) . Our goal is to design a personalized treatment plan ( policy ) for each patient moving forward . Notable challenges include the following : First , each patient only provides a single trajectory of their medical history . Second , each patient is heterogeneous in that they may have a varied response for a given treatment under similar medical conditions ; further , the underlying reason for this heterogeneous response across patients is likely unknown and thus not recorded in the dataset . Third , in the absence of an accurate personalized “ forecasting ” model for a patient ’ s medical outcome given a treatment plan , the treatment assigned is likely to be sub-optimal . This is particularly true for complicated medical conditions like T-Cell Lymphoma . We aim to address the challenges laid out above ( offline scarce data , heterogeneity , and suboptimal policies ) so as to develop a personalized “ forecasting ” model for each patient under any given treatment plan . Doing so will then naturally enable ideal personalized treatment policies for each patient . Key question . Tackling a scenario like the one described above in a principled manner is the focus of this work . In particular , we seek to answer the following question : “ Can we leverage scarce , offline data collected across heterogeneous agents under unknown , sub-optimal policies to learn a personalized policy for each agent ? ” Our Contributions . As our main contribution , we answer this question in the affirmative by developing a structured framework to tackle this challenging yet meaningful setting . Next , we summarize the main methodological , theoretical , algorithmic , and experimental contributions in our proposed framework . Methodological—personalized simulators . We propose a novel methodological framework , PerSim , to learn a policy given the data availability described above . Taking inspiration from the model-based RL literature , our approach is to first build a personalized simulator for each agent . Specifically , we use the offline data collectively available across heterogeneous agents to learn the unique transition dynamics for each agent . We do this without requiring access to the covariates or features that drive the heterogeneity amongst the agents . Having constructed a personalized simulator , we then learn a personalized decision-making policy separately for each agent by simply using online model predictive control ( MPC ) [ 16 , 9 ] . Theoretical—learning across agents . As alluded to earlier , the challenge in building a personalized simulator for each agent is that we only have access to a single offline trajectory for any given agent . Hence , each agent likely explores a very small subset of the entire state-action space . However , by viewing the trajectories across the multitude of agents collectively , we potentially have access to a relatively larger and more diverse offline dataset that covers a much richer subset of the state-action space . Still , any approach that augments the data of an agent in this manner must address the possibly large heterogeneity amongst the agents , which is challenging as we do not observe the characteristics that make agents heterogeneous . Inspired by the literature on collaborative filtering for recommendation systems , we posit that the transition dynamics across agents can be represented as a latent function of latent factors associated with agents , states , and actions . In doing so , we establish that this function is well-approximated by a “ low-rank ” decomposition of separable agent , state , and action latent functions ( Theorem 1 ) . Hence , for any finite sampling of the state and action spaces , accurate model learning for each agent with offline data—generated from any policy—can be reduced to estimating a low-rank tensor corresponding to agents , states , and actions . As such , low-rank tensors can provide a useful algorithmic lens to enable model learning with offline data in RL and we hope this work leads to further research studying the relationship between these seemingly disparate fields . Algorithmic—regularizing via a latent factor model . As a consequence of our low-rank representation result , we propose a natural neural network architecture that respects the constraints induced by the factorization across agents , states , and actions ( Section 4 ) . It is this principled structure , which accounts for agent heterogeneity and regularizes the model learning , that ensures the success of our approach despite access to only scarce and heterogeneous data . Further , we propose a natural extension of our framework which generalizes to unseen agents , i.e. , agents that are not observed in the offline data . Experimental—extensive benchmarking . Using standard environments from OpenAI Gym , we extensively benchmark PerSim against four state-of-the-art methods : a model-based online RL method ( CaDM ) [ 29 ] , two model-free offline RL method ( BCQ and CQL ) [ 15 , 25 ] , and a model-based offline RL method ( MOReL ) [ 23 ] . Below , we highlight six conclusions we reach from our experiments . ( i ) Despite access to only a single trajectory from each agent ( and no access to the covariates that drive agent heterogeneity ) , PerSim produces accurate personalized simulators for each agent . ( ii ) All benchmarking algorithms perform sub-optimally for the data availability we consider , even on simple baseline environments such as MountainCar and Cartpole , which are traditionally considered to be “ solved ” . ( iii ) PerSim is able to robustly extrapolate outside the policy used to generate the offline dataset , even if the policy is highly sub-optimal ( e.g. , actions are sampled uniformly at random ) . ( iv ) To corroborate our latent factor representation , we find that across all environments , the learned agent-specific latent factors correspond very closely with the latent source of heterogeneity amongst agents ; we re-emphasize this is despite PerSim not getting access to the agent covariates . ( v ) We find that augmenting the training data of an offline model-free method ( e.g. , BCQ ) with PerSim-generated synthetic trajectories results in a significantly better average reward . ( vi ) As an ablation study , if we decrease the number of observed trajectories , PerSim consistently achieves a higher reward than the other baselines across most agents , indicating its robustness to data scarcity . For a visual depiction of the conclusions , see Figure 1 for the learned latent agent factors and Figure 2 for the relative prediction accuracy of the learned model using PerSim versus [ 29 ] . Related Work . Due to space constraints , we present a short overview of the related literature . A more detailed review is provided in Appendix B . There are two sub-fields within RL that are of particular relevance : ( i ) modelbased online RL and ( ii ) model-free offline RL . ( i ) In the model-based RL literature , the transition dynamics ( simulator ) is learnt and subsequently utilized for policy learning . These methods have been found to have far better data efficiency compared to their model-free counterparts [ 50 , 10 , 11 , 26 , 21 , 17 ] . However , the current model-based literature mostly focuses on the setting where one can adaptively sample trajectories in an online manner during model-learning . A few recent works [ 29 , 36 ] have considered agent heterogeneity . We compare with [ 29 ] given its strong performance in handling heterogeneous agents . ( ii ) In the model-free offline RL literature , one uses a pre-recorded dataset to directly learn a policy , i.e. , without first learning a model . Thus far , the vast majority of offline RL methods are model-free and designed for settings that allow access to numerous trajectories from a single agent , i.e. , no agent heterogeneity [ 15 , 24 , 28 , 32 , 51 , 4 , 25 ] . To study how much offline methods suffer if agent heterogeneity is introduced , we compare with both [ 15 ] and [ 25 ] given their strong performance with offline data . We choose these baselines as they come from a well-established literature but their abilities to simultaneously handle agent heterogeneity and sparse offline data has yet to be studied . Model-based offline RL is still a relatively nascent field . Two recent works [ 23 , 53 ] have shown that , in certain settings , model-based offline methods can outperform their model-free counterparts on benchmark environments . In [ 53 ] , authors exhibit this using existing online model-based methods with minimal changes . However , both works restrict attention to the setting where there is only one agent ( or environment ) and a large number of observations from that agent are available . Hence , we study how these methods perform when given only sparse data from heterogeneous agents , and we find their performance does suffer . Extending these model-based offline RL methods to work with sparse data from heterogeneous agents , possibly by building upon the latent low-rank functional representation we propose , remains an interesting future work . 2 Problem Statement . We consider the standard RL framework with N heterogeneous agents . We index agents with n2 [ N ] 1 . Formally , we describe our problem as a Markov Decision Process ( MDP ) defined by the tuple ( S , A , Pn , Rn , n , µn ) . Here , S and A denote the state and action spaces , respectively , which are common across agents . For every agent n , Pn ( s0|s , a ) is the unknown transition kernel , Rn ( s , a ) is the immediate reward received , n2 ( 0,1 ) is the discounting factor , and µn is the initial state distribution . Observations . We consider an offline RL setting where we observe a single trajectory of length T for each of the N heterogeneous agents . Formally , for each agent n and time step t , let sn ( t ) 2 S , an ( t ) 2A , and rn ( t ) 2R denote the observed state , action , and reward . We denote our observations as D= { ( sn ( t ) , an ( t ) , rn ( t ) ) : n2 [ N ] , t2 [ T ] } . Goals . We state our two primary goals below . “ Personalized ” trajectory prediction . For a given agent n and state-action pair ( s , a ) 2 S⇥A , we would like to estimate E [ s0n| ( sn , an ) = ( s , a ) ] , i.e. , given the observations D , we would like to build a “ personalized ” simulator ( i.e. , a model of the transition dynamic ) for each agent n. “ Personalized ” model-based policy learning . To test the efficacy of the personalized simulator , we would like to subsequently use it to learn a good decision-making policy for agent n , denoted as ⇡n : S ! A , which takes as input a given state and produces a corresponding action . A causal inference lens . We note that this problem can be thought of as a counterfactual prediction problem . Our goal is to build a personalized simulator that answers the following question : “ what would have happened had the agent took a sequence of actions other than the one we observe in the dataset ? ” . Broadly speaking , our aim in PerSim is to answer such questions by observing how other agents behave under varied sequences of actions . More generally , we hope PerSim serves to further link causal inference with offline RL , which is fundamentally regarded as a counterfactual inference problem [ 31 ] .

Summary: This paper proposes a new model-based offline RL approach that tackles a particular offline RL problem where the dataset is generated by multiple agents and the data for each agent is scarce. To learn good policies from such heterogeneous dataset with limited per-agent data, the authors propose to learn per-agent personalized simulators (per-agent dynamics models) using latent low-rank factor representations with three encoders that encode the agent information, the state, and the action respectively. The authors show that the transition dynamics can be approximated with the low-rank order-three functional tensor representation with mild assumptions. Through experiments on gym tasks with heterogeneous agent data, the proposed method is shown to outperform prior online MBRL (adapted to the offline setting) and offline model-free RL method (BCQ).