Targeted Environment Design from Offline Data

1 INTRODUCTION . Recent years have witnessed the proliferation of reinforcement learning ( RL ) as a way to train agents to solve a variety of tasks , encompassing game playing ( Brown & Sandholm , 2019 ; Silver et al. , 2016 ) and some robotics applications ( Akkaya et al. , 2019 ; Schulman et al. , 2015 ; Chiang et al. , 2019 ) , among others . These successes often rely on the use simulators , providing several benefits to online RL agent training : large amounts of low cost and more diverse data not readily available in reality ( Tolani et al. , 2021 ) , and safe training that avoiding physical damage ( Krishnan et al. , 2021 ) or human privacy concerns ( Tolani et al. , 2021 ) . In many instances , for example in game playing , an accurate simulator is easy to devise . In others , rough approximations of the real target environment are sufficient ( e.g . Chiang et al . ( 2019 ) ) to close sim2real gap . However , in other applications such as robot manipulation ( Kalashnikov et al. , 2018 ; Todorov et al. , 2012 ) , health ( Ernst et al. , 2006 ) , or web navigation ( Gur et al. , 2021 ) designing the simulator to match a desired target environment with sufficient fidelity can be a highly manual process ( Nachum et al. , 2019a ; Krishnan et al. , 2021 ) , potentially requiring repeated online access to the target environment ( Chebotar et al. , 2019 ; Ramos et al. , 2019 ; Mehta et al. , 2020 ) . Due to the difficulty in matching simulators to target environments and motivated by the successes of supervised learning on large static datasets , a recent line of works Fujimoto et al . ( 2019 ) ; Agarwal et al . ( 2020 ) ; Wu et al . ( 2019 ) ; Levine et al . ( 2020 ) ; Fu et al . ( 2020 ) ; Gulcehre et al . ( 2020 ) has proposed to circumvent online training altogether ( whether on the simulator or in the target environment ) and instead perform RL on an offline dataset of logged experience in the target environment , typically collected by an unknown behavior policy . This approach avoids any potential mismatch between training and test environment . However , existing algorithms for offline RL are highly liable to generalization errors , and so these methods often impose a strong regularization on the learned policy to maintain proximity to the offline dataset ( Fujimoto et al. , 2019 ; Wu et al. , 2019 ; Kostrikov et al. , 2021 ; Yu et al. , 2020 ) . Thus , such methods can rarely approach the performance of online RL on an accurate simulator . Here , we propose to leverage both approximate simulators and offline datasets to train agents . Namely , we consider settings in which one has access to an offline dataset collected from a target environment as well an approximate ( but nevertheless inaccurate ) simulator . How can we leverage both of these to learn a near-optimal policy for the target environment ? To tackle this problem , we propose an approach which , at a high level , is composed of two stages : ( 1 ) use the offline dataset to learn a distribution over simulator parameters ; ( 2 ) train an RL agent in an online fashion on the simulator using any off-the-shelf RL algorithm . The key question is then how to appropriately learn the simulator parameters from an offline dataset , which we term offline targeted environment design ( OTED ) . For this problem , we devise an objective quantifying the divergence between the state-action distribution appearing in the offline dataset and the state-action distribution induced by a learned behavior policy acting in the simulator . We show that the simulator parameters and the target behavior policy can be jointly learned such that even a suboptimal behavior policy , that would otherwise perform poorly on the target environment , can help learn simulator parameters very effectively . We further show how the performance of the final returned policy can be improved by employing OTED not only for policy training but also for offline policy selection ( Yang et al. , 2020 ; Fu et al. , 2021 ) . Namely , we show how one may use the OTED-learned simulators to both provide a set of candidate policies as well as evaluate those policies in order to rank and choose the best policy to return , reminiscent of cross-validation techniques popular in supervised learning ( Arlot & Celisse , 2010 ) but so far elusive in RL contexts ( Paine et al. , 2020 ) . We apply our method to a variety of domains , including D4RL ( Fu et al. , 2020 ) and MiniGrid ( Chevalier-Boisvert et al. , 2018 ) , showing that our method is able to recover the target environment parameter using as few as 5 demonstrations and learn from high-dimensional image observations . We also show that an online agent trained in designed simulators yields improved performance compared to methods which leverage only simulators or only offline datasets , as well as existing methods for domain transfer ( Tobin et al. , 2017 ) . On the D4RL ( Fu et al. , 2020 ) benchmark , we demonstrate that our method can even outperform an online SAC agent trained on the ground-truth target environment . Our method achieves up to 17 times higher score compared to previous state-of-the-art offline RL and behavior cloning ( BC ) methods as well as domain randomization . On datasets with medium-level behavior that can be collected without domain experts , we show more relative improvements where our model even reaches the performance of the offline RL methods trained on expert-level behaviors . 2 RELATED WORK . Our work focuses on learning simulator parameters to match a target environment and is thus related to a large and diverse literature on domain transfer in RL . One of the most common and simple approaches to handling unknown simulator parameters is domain randomization ( Sadeghi & Levine , 2016 ; Tobin et al. , 2017 ; Matas et al. , 2018 ; Chiang et al. , 2019 ; Tolani et al. , 2021 ; Krishnan et al. , 2021 ) , which randomizes the setting of those unknown parameters during training , thus learning an agent which is robust to a wide distribution of possible simulators . While domain randomization has demonstrated successes , the process of choosing appropriate ranges for unknown parameters can be a highly manual process ( Andrychowicz et al. , 2020 ) , in which one must balance between either too wide range that hampers agent training , and too narrow range that may exclude the target environment parameterization ( Nachum et al. , 2019a ) . While some works have proposed heuristic mechanisms to better automate this process ( Akkaya et al. , 2019 ; Krishnan et al. , 2021 ) , others have explored incorporating online interaction with the environment into this tuning process ( Chebotar et al. , 2019 ; Ramos et al. , 2019 ; Mehta et al. , 2020 ) . In this work , we avoid any online interactions with the target environment during training , and only assume access to a static offline dataset . The proposed OTED method , which uses an offline dataset to learn a simulator , is related to recent model-based approaches to offline RL ( Matsushima et al. , 2020 ; Yu et al. , 2020 ; Argenson & Dulac-Arnold , 2020 ) . In these works , the offline dataset learns approximate dynamics and reward models of the environment , and the agent optimizes behavior in this approximate model of the environment . Although this makes minimal assumptions on the target environment , using reward and dynamics models as proxies for the environment leads to extrapolation issues which necessitate strong regularizers on the learned policy . Moreover , these existing approaches mostly ignore the fact that in many scenarios , one already has an approximate simulator , albeit not differentiable ; e.g. , the physical laws in a continuous control environments are straightforward to implement ( Todorov et al. , 2012 ) , although parameters more specific to the environment ( friction coefficient , actuator gain , gravity ) are unknown . Leveraging these simulators and focusing learning on the unknown parameters , as performed by our approach , can thus improve over modelling the full dynamics from scratch . Our proposed objective is based on a distribution matching loss . Similar losses have appeared in the past , especially in generative models ( Goodfellow et al. , 2014 ) , off-policy evaluation ( Nachum et al. , 2019b ) , off-policy learning ( Nachum et al. , 2019c ) , and imitation learning ( Ho & Ermon , 2016 ; Stadie et al. , 2017 ; Kostrikov et al. , 2020 ) . These predominantly focus on learning a policy via an adversarial objective and ignore the mismatch between simulators and real environments . Even when the distribution matching objective is optimal , this mismatch might give very suboptimal policies , due to potential poor support of the data distribution . This is a very realistic scenario since tuning simulator parameters accurately is a very tedious process and requires domain-experts . Finally , we emphasize that our proposed approach is distinct from previous work which suggests to combine offline learning on static datasets with online learning on the target environment ( Nair et al. , 2020 ; Ajay et al. , 2020 ; Yang & Nachum , 2021 ; Xie et al. , 2021 ; Lee et al. , 2021 ) . Specifically , we avoid any online interactions with the target environment , and only allow online training on the simulator , which may or may not be a good approximation to the target environment . 3 BACKGROUND . We consider an RL setting in which we are given access to both an offline dataset of experience from the target environment and an approximate simulator , in which the agent can gather more experience . We elaborate on the notation and background relevant to this setting . Reinforcement Learning We define an environment as a Markov Decision Process ( MDP ) given by a tuple ( S , A , ⇢0 , T , r , ) where S is a state space , A is an action space , ⇢0 is an initial state distribution , T ( s0|s , a ) is a state transition distribution , r : S ⇥ A ! R is a reward function over state and actions , and 2 [ 0 , 1 ) is a discount factor . A stationary policy ⇡ in this environment is a function from states to distributions over actions . In reinforcement learning , we are interested in finding a policy ⇡ that maximizes the cumulative discounted returns : E [ P1 t=0 tr ( st , at ) |s0 ⇠ ⇢0 , at ⇠ ⇡ ( st ) , st+1 ⇠ T ( st , at ) ] . We will use d⇡ ( s , a ) to denote the state-action distribution of ⇡ : d⇡ ( s , a ) : = ( 1 ) P1 t=0 tPr [ st = s , at = a|s0 ⇠ ⇢0 , at ⇠ ⇡ ( st ) , st+1 ⇠ T ( st , at ) ] . Simulator vs. Target Environment Typically in online reinforcement learning , a policy is trained based on experience collected by the agent in a simulator . While this implicitly assumes that the simulator can perfectly mimic the target environment , in practice this is often not the case . The simulator typically only provides an approximation to the true initial state distribution ⇢0 and state transition distribution T . For example in robotics applications , while physics simulators can be built , it is infeasible to accurately set all appropriate parameters ( e.g. , friction coefficients , actuator gains , wind speed ) . In this work , we assume such a parameterized simulator exists , and it is given by M ( z ) : = ( S , A , ⇢0 ( ·|z ) , T ( ·|· , · , z ) , r ( z ) , ) , where z denotes the parameters of the simulated environment . Naive approaches would either set z heuristically or using randomization methods . In this work , we aim to use offline datasets to learn z so that M ( z ) approximates the target MDP , which we denote as M⇤ . We use d⇡ ( s , a|z ) to denote the state-action distribution of ⇡ in the parameterized simulator M ( z ) . Offline Reinforcement Learning Offline RL , as an alternative to online RL on simulators , assumes that an agent can not collect any online experience and is restricted to a dataset D of ( s , a , s0 ) collected by an unknown behavior policy µ ; i.e. , ( s , a ) ⇠ dµ , s0 ⇠ p ( s , a ) . We use dD ( s , a ) to denote the distribution over state-action pairs appearing in D. When trained on a fixed dataset , previous state-ofthe-art online reinforcement learning algorithms such as SAC ( Haarnoja et al. , 2019 ) underperform due to extrapolation errors ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ) , which stem from sparsity of the samples in the dataset . Accordingly , many RL algorithms designed specifically for the offline setting make use of strong regularizations to maintain proximity to the offline dataset ( Wu et al. , 2019 ; Kostrikov et al. , 2021 ; Yu et al. , 2020 ) .

Summary: The paper introduces an approach to tune a parameterized simulator to produce data samples that mimic a given offline dataset of trajectories from real data. The core technical contributions involve defining a KL divergence to quantify the difference of the simulator data to a given offline dataset, a training algorithm that interleaves simulator tuning and agent training on simulator data, and selection algorithms to choose a hyperparameterization among simulators and a policy among a set trained on a simulator. Evaluations assess how well the training procedure can recover known simulator parameters and the performance of policies trained on a learned simulator to policies trained with offline RL algorithms on the offline data directly. Results show the simulator can accurately recover parameters and is able to obtain similar performance to offline RL algorithms when tuning certain parameters.