Autoregressive Dynamics Models for Offline Policy Evaluation and Optimization

1 INTRODUCTION . Model-based Reinforcement Learning ( RL ) aims to learn an approximate model of the environment ’ s dynamics from existing logged interactions to facilitate efficient policy evaluation and optimization . Early work on Model-based RL uses simple tabular ( Sutton , 1990 ; Moore and Atkeson , 1993 ; Peng and Williams , 1993 ) and locally linear ( Atkeson et al. , 1997 ) dynamics models , which often result in a large degree of model bias ( Deisenroth and Rasmussen , 2011 ) . Recent work adopts feedforward neural networks to model complex transition dynamics and improve generalization to unseen states and actions , achieving a high level of performance on standard RL benchmarks ( Chua et al. , 2018 ; Wang et al. , 2019 ) . However , standard feedforward dynamics models assume that different dimensions of the next state and reward are conditionally independent given the current state and action , which may lead to a poor estimation of uncertainty and unclear effects on RL applications . In this work , we propose a new family of autoregressive dynamics models and study their effectiveness for off-policy evaluation ( OPE ) and offline policy optimization on continuous control . Autoregressive dynamics models generate each dimension of the next state conditioned on previous dimensions of the next state , in addition to the current state and action ( see Figure 1 ) . This means that to sample the next state from an autoregressive dynamics model , one needs n sequential steps , where n is the number of state dimensions , and one more step to generate the reward . By contrast , standard feedforward dynamics models take current state and action as input and predict the distribution of the next state and reward as a multivariate Gaussian with a diagonal covariance structure ( e.g. , Chua et al . ( 2018 ) ; Janner et al . ( 2019 ) ) . This modeling choice assumes that different state dimensions are conditionally independent . ∗Work done as an intern at Google Brain . Autoregressive generative models have seen success in generating natural images ( Parmar et al. , 2018 ) , text ( Brown et al. , 2020 ) , and speech ( Oord et al. , 2016 ) , but they have not seen use in Model-based RL for continuous control . We find that autoregressive dynamics models achieve higher log-likelihood compared to their feedforward counterparts on heldout validation transitions of all DM continuous control tasks ( Tassa et al. , 2018 ) from the RL Unplugged dataset ( Gulcehre et al. , 2020 ) . To determine the impact of improved transition dynamics models , we primarily focus on OPE because it allows us to isolate contributions of the dynamics model in value estimation vs. the many other factors of variation in policy optimization and data collection . We find that autoregressive dynamics models consistently outperform existing Model-based and Model-free OPE baselines on continuous control in both ranking and value estimation metrics . We expect that our advances in model-based OPE will improve offline policy selection for offline RL ( Paine et al. , 2020 ) . Finally , we show that our autoregressive dynamics models can help improve offline policy optimization by model predictive control , achieving a new state-of-the-art on cheetah-run and fish-swim from RL Unplugged ( Gulcehre et al. , 2020 ) . Key contributions of this paper include : • We propose autoregressive dynamics models to capture dependencies between state dimensions in forward prediction . We show that autoregressive models improve log-likelihood over nonautoregressive models for continuous control tasks from the DM Control Suite ( Tassa et al. , 2018 ) . • We apply autoregressive dynamics models to Off-Policy Evaluation ( OPE ) , surpassing the performance of state-of-the art baselines in median absolute error , rank correlation , and normalized top-5 regret across 9 control tasks . • We show that autoregressive dynamics models are more useful than feedforward models for offline policy optimization , serving as a way to enrich experience replay by data augmentation and improving performance via model-based planning . 2 PRELIMINARIES . Here we introduce relevant notation and discuss off-policy ( offline ) policy evaluation ( OPE ) . We refer the reader to Lange et al . ( 2012 ) and Levine et al . ( 2020 ) for background on offline RL , which is also known as batch RL in the literature . A finite-horizon Markov Decision Process ( MDP ) is defined by a tupleM = ( S , A , T , d0 , r , γ ) , where S is a set of states s ∈ S , A is a set of actions a ∈ A , T defines transition probability distributions p ( st+1|st , at ) , d0 defines the initial state distribution d0 ≡ p ( s0 ) , r defines a reward function r : S × A → R , and γ is a scalar discount factor . A policy π ( a | s ) defines a conditional distribution over actions conditioned on states . A trajectory consists of a sequence of states and actions τ = ( s0 , a0 , s1 , a1 , . . . , sH ) of horizon length H . We use st , i to denote the i-th dimension of the state at time step t ( and similarly for actions ) . In reinforcement learning , the objective is to maximize the expected sum of discounted rewards over the trajectory distribution induced by the policy : Vγ ( π ) = Eτ∼pπ ( τ ) [ H∑ t=0 γtr ( st , at ) ] . ( 1 ) The trajectory distribution is characterized by the initial state distribution , policy , and transition probability distribution : pπ ( τ ) = d0 ( s0 ) H−1∏ t=0 π ( at|st ) p ( st+1|st , at ) . ( 2 ) In offline RL , we are given access to a dataset of transitions D = { ( sit , ait , rit+1 , sit+1 ) } Ni=1 and a set of initial states S0 . Offline RL is inherently a data-driven approach since the agent needs to optimize the same objective as in Eq . ( 1 ) but is not allowed additional interactions with the environment . Even though offline RL offers the promise of leveraging existing logged datasets , current offline RL algorithms ( Fujimoto et al. , 2019 ; Agarwal et al. , 2020 ; Kumar et al. , 2019 ) are typically evaluated using online interaction , which limits their applicability in the real world . The problem of off-policy ( offline ) policy evaluation ( OPE ) entails estimating Vγ ( π ) , the value of a target policy π , based on a fixed dataset of transitions denoted D , without access to the environment ’ s dynamics . Some OPE methods assume that D is generated from a known behavior ( logging ) policy µ and assume access to µ in addition to D. In practice , the logged dataset D may be the result of following some existing system that does not have a probabilistic form . Hence , in our work , we will assume no access to the original behavior policy µ for OPE . That said , for methods that require access to µ , we train a behavior cloning policy on D . 3 PROBABILISTIC DYNAMICS MODELS . Feedforward dynamics model . In the context of our paper , we use the term “ model ” to jointly refer to the forward dynamics model ps ( st+1|st , at ) and reward model pr ( rt+1|st , at ) . We use neural nets to parameterize both distributions since they are powerful function approximators that have been effective for model-based RL ( Chua et al. , 2018 ; Nagabandi et al. , 2018 ; Janner et al. , 2019 ) . Let θ denote the parameters of a fully connected network used to model pθ ( st+1 , rt+1 | st , at ) . We expect joint modeling of the next state and reward to benefit from sharing intermediate network features . Similar to prior work ( Janner et al. , 2019 ) , our baseline feedforward model outputs the mean and log variance of all state dimensions and reward simultaneously , as follows : pθ ( st+1 , rt+1 | st , at ) = N ( µ ( st , at ) , Diag ( exp { l ( st , at ) } ) ) , ( 3 ) where µ ( st , at ) ∈ Rn+1 denotes the mean for the concatenation of the next state and reward , l ( st , at ) ∈ Rn+1 denotes the log variance , and Diag ( v ) is an operator that creates a diagonal matrix with the main diagonal specified by the vector v. During training , we seek to minimize the negative log likelihood of the parameters given observed transitions in the dataset D : ` ( θ | D ) = − ∑ ( s , a , r′ , s′ ) ∈D log pθ ( s ′ , r′ | s , a ) . ( 4 ) While it is possible to place different weights on the loss for next state and reward prediction , we did not apply any special weighting and treated the reward as an additional state dimension in all of our experiments . This is straightforward to implement and does not require tuning an additional hyperparameter , which is challenging for OPE . Note that the input has |s|+ |a| dimensions . Autoregressive dynamics model . We now describe our autoregressive model . We seek to demonstrate the utility of predicting state dimensions in an autoregressive way . Therefore , rather than using a complex neural network architecture , where improvements in log-likelihood and policy evaluation are confounded by architectural differences , we opt to make simple modifications to the feedforward model described above . This allows us to isolate the source of performance improvements . The autoregressive model we use is a fully connected model that predicts the mean and log variance of a single state dimension . We augment the input space of the baseline with the previous predicted state dimensions and a one-hot encoding to indicate which dimension to predict . This is illustrated in Figure 1 . The autoregressive model therefore has 3|s| + |a| input dimensions . Hence , the autoregressive model has a small number of additional weights in the first fully connected layer , but as will be shown in our experiments , these extra parameters are not the reason for a performance gain . At training time , the autoregressive model has a similar computational cost to the fully connected model as we can mask ground truth states and use data parallelism to compute all state dimensions simultaneously . At inference , the autoregressive model requires additional forward passes , on the order of the number of state dimensions in a given environment . We use the default ordering for the state dimensions in a given environment , though it is interesting to explore different orderings in future works . The negative log-likelihood for an autoregressive model takes the form of : ` ( θ | D ) = − ∑ ( s , a , r′ , s′ ) ∈D [ log pθ ( r ′ | s , a , s′ ) + ∑n i=1 log pθ ( s ′ i | s , a , s′1 , . . . , s′i−1 ) ] , ( 5 ) where we use chain rule to factorize the joint probability of p ( s′ , r′ | s , a ) . The main advantage of the autoregressive model is that it makes no conditional independence assumption between next state dimensions . This class of models can therefore capture non-unimodal dependencies , e.g. , between different joint angles of a robot . Paduraru ( 2007 ) demonstrates this increased expressivity in the tabular setting , constructing an example on which a model assuming conditional independence fails . While the expressive power of autoregressive models have been shown in various generative models ( Parmar et al. , 2018 ; Oord et al. , 2016 ) , autoregressive dynamics models have not seen much use in Model-based RL for continuous control before this work . Algorithm 1 Model-based OPE Require : Number of rollouts n , discount factor γ , horizon length H , policy π , dynamics model p , set of initial states S0 for i = 1 , 2 , . . . n do Ri ← 0 sample initial state s0 ∼ S0 for t = 0 , 1 , 2 , . . . , H − 1 do sample from policy : at ∼ π ( · | st ) sample from the dynamics model : st+1 , rt+1 ∼ p ( · , · | st , at ) Ri ← Ri + γtrt+1 end for end for return 1n ∑n i=1Ri Model-based OPE . Once a dynamics model is trained from offline data , OPE can be performed in a direct and primitive way . We let the policy and model interact—the policy generates the next action , the model plays the role of the environment and generates the next state and reward . Due to the stochasticity in the model and the policy , we estimate the return for a policy with Monte-Carlo sampling and monitor standard error . See Algorithm 1 for pseudocode .

Summary: The paper studies offline policy evaluation (OPE) and optimization in the model-based setting. The main methodological contribution of the paper is using autoregressive models for the next state and reward prediction. The authors demonstrate that autoregressive models achieve higher likelihood compared to feedforward models on 9 environments from RL Unplugged [1] offline dataset. Given that model likelihood is only a proxy quality metric in OPE and control, they further demonstrate a positive correlation between likelihood and OPE estimates. The paper shows quantitatively that using autoregressive models results in more accurate OPE estimates than for feedforward models and model-free benchmarks. Finally, the authors apply autoregressive models for offline control and achieve higher returns than for feedforward models.