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: The paper proposes PerSim, a model-based off-line reinforcement learning algorithm targeting a setting where each recorded trajectory is produced by acting in an environment with slightly different transition dynamics. PerSim first learns a simulator for each of those environments, taking advantage of a low-rank model similar to a recommender system to account for similarities between the environments. It then uses the simulator for a specific environment to perform model predictive control. The experiments compare against strong baselines, both model-free and model-based, and demonstrate gains both in terms of the average reward and in terms of simulator accuracy (in the latter case).