Counterfactual Regularization for Model-Based Reinforcement Learning

1 INTRODUCTION . Recent advancements in the use of variational inference and generative neural network architectures have made it possible to build accurate transition models for high-dimensional sequential decisionmaking environments . Networks have been trained to accurately predict future states conditioned on actions in large Markov decision process state spaces including Atari , VisDoom , and pixel-based robotic control tasks ( Kurutach et al. , 2018 ) ( Ha & Schmidhuber , 2018 ) ( Hafner et al. , 2018 ) . Using such a transition model , an agent can plan sequences of future actions to maximize an expected reward , using model-predictive control or other search algorithms . These approaches fall into the paradigm of model-based reinforcement learning ( RL ) , distinct from the standard model-free approach of training a policy or value network . Model-based approaches have better sample efficiency than model-free methods , which usually require a large amount of training data to perform well . There are distinct differences between traditional model-based RL , in which a transition between states are defined by a Markov transition matrix , and model-based approaches in which the state representation is learned , in which the state is represented as a continuous vector . This is made more unclear with environments that start with a scene described as an image , where it is well-known typical transition models that maximize a variational bound on the likelihood of future observations , similar to a variational autoencoder ( Kingma & Welling , 2013 ) tend to generate blurry samples and exhibit posterior collapse ( van den Oord et al. , 2017 ) , a phenomenon showing that the learned states may not be robust enough . We hypothesize that a representation in which states are factored and well-separated will improve the performance of model-based agents . In order to achieve this , we propose a new family of regularizations applicable to action-conditional sequence prediction models , which we call counterfactual regularization . These regularizations are loss terms based on counterfactual states , which can be generated by feeding non-ground-truth actions or states into a transition model . Specifically , we propose regularization terms to achieve two goals . First , the set of possible states should be distinct in such a way that any change in input actions should result in a change in latent state . Second , that the dimensions of the learned state should be disentangled into separate factors in such a way that a change to one factor should effect only a sparse set of other factors . In order to test these regularizations , we introduce the StarIntruders game , a sequential environment with test-time dynamics that vary from training dynamics to test the generalization performance of trained agents . We propose a convolutional recurrent latent-state architecture well-suited for modeling two-dimensional video environments and trained using counterfactual regularization . Our experiments show that a simple planning agent using an action-conditional model trained with the proposed counterfactual regularization outperforms both non-regularized models , and a model-free agent , in test environments that differ from training . 2 LATENT DYNAMICS MODELS . 2.1 PRELIMINARIES . We consider a partially-observable Markov decision process ( POMDP ) in which an agent observes at each time step t an observation ot from an observation space O and a reward rt from a reward space R and executes an action at from a discrete set of actions A . The agent executes actions from t = 1 until a terminal timestep t = T , generating a trajectory ( o1 , r1 , a1 ) ... ( oT , rT , aT ) . An optimal agent will take actions that maximize the cumulative reward ∑T t rt over the trajectory . In pixel-based domains , ot may be a high-dimensional vector ( e.g . O = R256x256x3 ) and the underlying dynamics function p ( ot , rt|o < t , r < t , a < t ) may not be known . We are concerned with the problem of learning from experience a parameterized function pθ that estimates the distribution of future observations and rewards conditioned on actions . We model the environment using a neural network similar to the deterministic state space model from ( Hafner et al. , 2018 ) , trained on trajectories collected by an exploration policy . Using such a function ( a model of the environment ) , planning algorithms such as model-predictive control can be applied to compute an optimal action at at each timestep . Unlike model-free reinforcement learning agents , a model-based agent does not learn a policy network or value function estimate ; instead , actions are chosen by searching among possible future trajectories . Previous work has explored the efficiency benefits of model-based methods , which can achieve similar performance to model-free methods using far fewer training examples ( Hafner et al. , 2018 ) ( Kaiser et al. , 2019b ) . Our work instead explores the capability of model-based agents to generalize to changes in the environment . The graphical model in Fig . 1 illustrates our assumptions about the generative process . We model all dependencies on previous timesteps through an unobserved latent state variable st , which depends only on the previous state and action st−1 , at−1 . Each observation depends only on the current state . Accurately predicting future observations and rewards depends on building a highly accurate latent dynamics model . 2.2 MODEL ARCHITECTURE . Our latent dynamics model consists of the following components : an encoder E : O → S , an action-conditional transition function T : S ×A→ S , a decoder G : S → O , and a reward decoder R : S → RK , as shown in Figure 1 . Each component is a convolutional deep neural network with strided and transposed convolutions , and LeakyReLU nonlinearities . Spectral Normalization ( Miyato et al. , 2018 ) is applied to each convolutional layer of T to assist in gradient propagation . For environments in which state estimation requires observation of multiple sequential states , the input layer of E is resized to take as input a stack of K sequential frames . Where not otherwise noted , K = 3 . The transition network T converts each integer action input to a one-hot categorical representation , broadcast across the spatial input to the network . The reward networkR predicts a positive or negative reward at each timestep , for each spatial location , for each reward type . To improve training stability , reward estimates are discretized in the following way : each spatial location is classified as causing -1 , 0 , or +1 reward , and the output reward is the sum of all spatial locations . 2.3 TRAINING . The model is trained to reconstruct sequences of observations gathered from the true environment by an exploration policy . Given a sequence of ground truth observations , the model must reconstruct future observations and rewards ( see Figure 2 ) . Similar to ( Hafner et al. , 2018 ) , our model consists of separate encoder/decoder networks which convert between observations and a learned latent representation , and a transition network which operates recurrently in the latent space ( See Figure 1 ) . This structure allows the latent representation to keep track of parts of the state not visible in current observations . 2.4 MODEL-PREDICTIVE CONTROL . Given a predictive model of future observations and rewards , we apply a simple model-predictive control algorithm to select actions at each time step during evaluation . The planning process produces a batch of B = 64 trajectories for each possible action ai with each trajectory taking a random rollout for a fixed horizon T = 20 . The agent executes the action at that maximizes the estimated sum of rewards ∑t+T i=t+1 ri , observes the next observation ot+1 and then re-plans at each step . 2.5 COUNTERFACTUAL TRAJECTORIES . A transition network T takes as input a latent state st and an action at , and produces an estimate of the successor state ŝt+1 . The network is trained to simulate a ground truth trajectory { ( o1 , a1 ) , ( o2 , a2 ) ... ( oN , aN ) } generated by a policy π in the real environment . The transition network is applied recurrently to generate a sequence of states : st+2 = T ( st+1 , at+1 ) = T ( T ( st , at ) , at+1 ) During training , st is an observed state and at is a known ground-truth action , selected by the training exploration policy . However , an alternate action a′t or an alternate state s ′ t can be input to the transition function , thereby generating an alternate or counterfactual state : s′t+1 = T ( st , a ′ t ) or s′t+1 = T ( s ′ t , at ) . Repeated application of the transition function then generates a counterfactual trajectory : s′t+2 = T ( s ′ t+1 , a ′ t+1 ) For many tasks , especially in physical environments , it may not be possible to repeat a previously observed trajectory , and so the “ true ” value of a counterfactual state may never be observable . However , the estimate ŝ′t+1 is computable , and because T is a differentiable parameterized function , the gradient of a loss function involving generated counterfactual states can be computed with respect to network parameters . Relationships among estimates of real and counterfactual states provide a language by which statements about causality can be converted into differentiable loss terms suitable for inclusion in a neural network training process . Many useful desiderata can be posed in the language of counterfactuals , and converted to loss terms along with standard pixel-reconstruction and reward losses used in predictive models . By including a counterfactual loss term LR in the objective function , the full loss becomes : L = ∑ t ||D ( st ) − ot||22 + ||R ( st ) − rt||22 + λLR ( 1 ) where D and R are the decoder and reward prediction networks , ot , rt is a ground truth observation and reward for timestep t , ŝt is a predicted latent state output by one or more recurrent applications of T , λ is a scalar hyperparameter , and LR is one of the counterfactual regularization terms described below . 3 CAUSAL REGULARIZATION THROUGH COUNTERFACTUALS . By generating counterfactual trajectories during training ( see Figure 3 ) , we can augment the training objective with new terms . We explore a number of possible types of counterfactual regularization , given below . 3.1 ACTION-CONTROL REGULARIZATION . When exploring a complex environment with an untrained policy , it may be the case that the reward function can be accurately modeled as a deterministic function of time , unrelated to the agent ’ s actions . For instance , in a timed game ending with a reward of -1 or +1 indicating win or loss , a poorly-initialized agent might learn that its actions have no effect on the environment , and may model the reward mechanism as a timer that outputs 0 or -1 . This failure case could be considered analogous to the psychological phenomenon of learned helplessness , the perception of independence between actions and outcomes ( Langer , 1975 ) . To address this problem , we might wish to limit our learning system to consider only models with the property that every action taken by the agent has a causal effect on the state . Using the language of counterfactuals , we can express this property precisely : for any state st reached by taking actions a1 , a2 , ... at−1 , there must exist some a′1 , ... a ′ t−1 such that st 6= s′t and ∃ i < t , ai 6= a′i . Considering i = t− 1 for simplicity , we can convert the above constraint into a regularization term : min a′t−1 − log ||T ( st−1 , at−1 ) − T ( st−1 , a′t−1 ) ||1 where s′t = T ( st−1 , a ′ t−1 ) is a counterfactual state generated by the alternate a ′ i . This regularization ensures that the learned environment model must include a causal link from ai to st , in the sense that for any st there must exist a counterfactual action a′i which would have resulted in a different state s ′ t. In a more general case , the counterfactual a′i might be multiple time steps removed from st . Using T ( s , a1 ... an ) to denote recurrent application of T to simulate multiple timesteps , the term can be expanded to the general case of i < t− 1 . LR = min a′i − log ||T ( si , ai ... at−1 ) − T ( si , a′i ... a′t−1 ) || ( 2 ) For finite discrete action spaces , a′i can be minimized by exhaustive search , but for simple environments it suffices to generate a′i by sampling uniformly from A \ { at } . This loss term has the effect of biasing the transition model by forcing the latent state to change as a result of every agent action , which is an appropriate assumption for many control tasks .

Summary: The paper presents regularization techniques for model based reinforcement learning which attempt to build counterfactual reasoning into the model. In particular, they present auxiliary loss terms which can be used in "what if" scenarios where the actual state is unknown. Given certain assumptions, they show that this added regularization can improve generalization to unseen problem settings. Specifically they propose two forms of regularization: (1) enforcing that for different actions the predicted next state should be different (action-control) and (2) enforcing that when certain parts of the low dimensional state are perturbed, over a model rollout the perturbation should only affect the perturbed parts of the state, essentially encouraging the latent space features to be independent (disentanglement).