Title: COUNTERFACTUAL REGULARIZATION FOR MODEL-BASED REINFORCEMENT LEARNING

Abstract: In sequential tasks, planning-based agents have a number of advantages over model-free agents, including sample efficiency and interpretability. Recurrent action-conditional latent dynamics models trained from pixel-level observations have been shown to predict future observations conditioned on agent actions accurately enough for planning in some pixel-based control tasks. Typically, models of this type are trained to reconstruct sequences of ground-truth observations, given ground-truth actions. However, an action-conditional model can take input actions and states other than the ground truth, to generate predictions of unobserved counterfactual states. Because counterfactual state predictions are generated by differentiable networks, relationships among counterfactual states can be included in a training objective. We explore the possibilities of counterfactual regularization terms applicable during training of action-conditional sequence models. We evaluate their effect on pixel-level prediction accuracy and model-based agent performance, and we show that counterfactual regularization improves the performance of model-based agents in test-time environments that differ from training.

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. The latent state space s is represented as a convolutional feature map, and the transition function T (s) is implemented as a fully-convolutional neural network.

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 o t from an observation space O and a reward r t from a reward space R and executes an action a t from a discrete set of actions A. The agent executes actions from t = 1 until a terminal timestep t = T , generating a trajectory (o 1 , r 1 , a 1 )...(o T , r T , a T ). An optimal agent will take actions that maximize the cumulative reward T t r t over the trajectory. In pixel-based domains, o t may be a high-dimensional vector (e.g. O = R 256x256x3 ) and the underlying dynamics function p(o t , r t |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 a t 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 s t , which depends only on the previous state and action s t-1 , a t-1 . Each observation depends only on the current state. Accurately predicting future observations and rewards depends on building a highly accurate latent dynamics model. (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 network R 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 a i with each trajectory taking a random rollout for a fixed horizon T = 20. The agent executes the action a t that maximizes the estimated sum of rewards t+T i=t+1 r i , observes the next observation o t+1 and then re-plans at each step. 

2.5. COUNTERFACTUAL TRAJECTORIES:

A transition network T takes as input a latent state s t and an action a t , and produces an estimate of the successor state ŝt+1 . The network is trained to simulate a ground truth trajectory {(o 1 , a 1 ), (o 2 , a 2 )...(o N , a N )} generated by a policy π in the real environment.

The transition network is applied recurrently to generate a sequence of states:

s t+2 = T (s t+1 , a t+1 ) = T (T (s t , a t ), a t+1 )

During training, s t is an observed state and a t 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 (s t , a t ) or s t+1 = T (s t , a t ). 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 ŝ 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 L R in the objective function, the full loss becomes:

EQUATION

where D and R are the decoder and reward prediction networks, o t , r t is a ground truth observation and reward for timestep t, ŝt is a predicted latent state output by one or more recurrent applications of T , λ is a scalar hyperparameter, and L R 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 s t reached by taking actions a 1 , a 2 , ...a t-1 , there must exist some a 1 , ...a t-1 such that s t = s t and ∃ i < t, a i = a i .

Considering i = t -1 for simplicity, we can convert the above constraint into a regularization term:

min a t-1 -log ||T (s t-1 , a t-1 ) -T (s t-1 , a t-1 )|| 1

where s t = T (s t-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 a i to s t , in the sense that for any s t 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 s t . Using T (s, a 1 ...a n ) to denote recurrent application of T to simulate multiple timesteps, the term can be expanded to the general case of i < t -1.

EQUATION

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 \ {a t }. 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.

3.2. DISENTANGLEMENT REGULARIZATION:

Many practical environments contain separate and largely independent factors of variation. For example, in the game of Space Invaders, the agent controls one visible object on the screen (the player character) while enemy characters move across the screen, independently of the player. The only interaction between enemy position and player position is through projectiles fired by the player at enemies. It might be desired that a model of the game represent the position of enemies as independent of the position of the player, conditioned on projectiles.

In a non-sequential variational autoencoder setting, FactorVAE (Kim & Mnih, 2018) applies a regularization based on adversarial training using examples generated by permuting the dimensions of the latent representation of real examples. This forces the network to learn a representation such that any two latent dimensions could be swapped, and the result should still be realistic. The resulting representation maximizes the independence of latent dimensions.

Similarly but in a sequential environment, we desire a representation with the property that any two latent dimensions indices k, j of a state s t can be swapped to create an alternate s t , and the resulting sequence of counterfactual states s t+1 , s t+2 , ... would continue on with all dimensions that are independent of k and j unchanged. That is, changing the values of some dimensions of the latent representation should cause a sparse and minimal change in other latent dimensions. We represent this sparsity with an L1 regularization term:

L D = i / ∈{j,k} |s (i) t+H -s t+H (i) | (3)

where s (i) denotes the ith dimension (or channel) of the latent state vector s. Here, s t = s

(1) t , ...s

(i) t , ...s (j) t , ...s (K) t

is the state vector at time t and s t = s

(1) t , ...s

(j) t , ...s (k) t , ...s (K) t

is the vector with dimensions j and k swapped. A fixed number of time steps H is used between the counterfactual intervention and the regularization. The effect of this regularization is to promote sparse, rather than dense, relationships between latent factors. This is appropriate for environments where the underlying mechanics are discrete and deterministic.

4. EXPERIMENTS:



4.1. ACTION-CONDITIONAL STATE PREDICTION:

Our model is trained with ground-truth trajectories sampled by a random exploration policy. We train with λ = .01 and a maximum prediction horizon of H = 10 in both environments. After training, the model-predictive control (MPC) agent is applied to select the action with highest expected reward over the planning horizon at each time step.

4.2. GENERALIZATION EVALUATION ENVIRONMENT:

We train and evaluate agents in the Star Intruders game, a custom sequential environment with mechanics similar to Space Invaders, built using the StarCraft II Learning Environment (Vinyals et al., 2017) . The state space of the task includes a 256x256x3 RGB pixel matrix for visualization, and a 64x64x4 binary mask input to the network, indicating the positions of each type of object. The MSE H=3 MSE H=5 MSE H=10 MSE H=20 Avg. dynamics of the task are as follows (See Figure 4 ). The agent controls a unit at the bottom of the game map. Above the agent, randomly-positioned friendly and enemy units patrol left and right At each time step, the agent can take one of four actions: move left, move right, fire, or no-op. The fire action fires a projectile from the player's current position, which annihilates upon contact with any enemy or friendly unit.

When a projectile destroys an enemy unit, the agent receives a reward of 1, and when a projectile destroys a friendly unit, the agent receives a reward of -1. The optimal policy will fire projectiles aimed in such a way that all enemy units are destroyed but no friendly units are destroyed.

The training environment contains 8 friendly units and 8 enemy units. Each friendly unit begins at a random location in the bottom half of the map, and moves horizontally in a repeating pattern, from its starting location to the rightmost edge of the map, then to the leftmost edge, and so on. Each enemy unit patrols similarly from a random starting location in the top half of the map. Each episode continues until either all enemy units are destroyed, or a maximum number of time steps T M = 300 are elapsed. The maximum cumulative reward in any episode is 8.

4.2.1. GENERALIZATION EVALUATION:

After training, each policy is evaluated in three environments. The Training environment operates exactly as described in the previous section. The Test A environment operates in the same way as the training environment, except that instead of patrolling right-to-left, each unit will patrol either right-to-left or left-to-right with equal probability. The Test B environment operates in the same way as the training environment, except that instead of starting at random locations within the entire bottom half of the map, all friendly units are positioned within a small horizontal band close to the bottom of the map. From a subjective human perspective, the Test A and Test B environments are no more challenging than the training environment. Colors, textures, timing and input format in the test environments are identical to the training environment.

For comparison with model-free methods, we train a network with the Rainbow algorithm (Hessel et al., 2018) to maximize reward in the StarIntruders environment. The policy is trained for 1M time steps in the original training environment. The trained policy is evaluated in the original training environment, and in each variant version of the environment.

4.3. RESULTS:

Table 2 shows game score in the StarIntruders training and test environments. Without regularization, the model-free Rainbow algorithm outperforms our model-predictive control agent in the training environment. However, the model-based method generalizes to new environments more effectively.

Applying counterfactual regularization to the predictive model improves the performance of our MPC agent in both the training environment and test environments.

Table 1 shows both game score and pixel prediction error for the simpler MiniPacMan task. We note that lower MSE does not necessarily coincide with higher score, so the proposed regularizations do not simply improve pixel prediction accuracy. We conclude that counterfactual regularization improves the representation learned by our environment model.

Method:

Training Test A Test B Rainbow (Hessel et al., 2018) 7 

5. RELATED WORK:

Much previous work in model-based reinforcement learning involves planning with model-predictive control in low-dimensional state spaces using neural networks (Gal et al.) (Chua et al., 2018) or other models (Deisenroth & Rasmussen, 2011) . Other work has applied generative neural network models to high-dimensional video prediction (Chiappa et al., 2017) or state prediction in Markov reward processes (Silver et al., 2017) where transitions are not action-conditional. Some work has explored action-conditional video prediction (Oh et al., 2015) without planning. Recent work has applied model-predictive planning to high-dimensional image domains (Kaiser et al., 2019a ) (Hafner et al., 2018) .

LEARNING WITH AUXILIARY REWARDS:

Related to model regularization, some work applies auxiliary tasks during training of a model-free or model-based agent with the goal of improving performance or generalization. In (Jaderberg et al., 2016) a representation is learned using auxiliary policies based on intrinsic rewards such as maximizing changes in pixel intensity, in a model-free setting. In a model-based setting, in (Pathak et al., 2017) a representation of the environment state is trained using inverse or forward dynamics, and the prediction error of this dynamics model is subsequently used as an intrinsic reward signal to train a model-free policy network.

DISENTANGLEMENT OF FACTORS:

Factorized state representations have long been known to improve performance in reinforcement learning tasks (Degris et al., 2006) . In (Higgins et al., 2017) , the Kullback-Liebler divergence term of a variational autoencoder's objective function is reinterpreted as a regularization term promoting conditional independence of latent dimensions. In (Kim & Mnih, 2018) , independence of an autoencoder's latent dimensions is achieved via an adversarial loss.

MODEL-BASED REINFORCEMENT LEARNING:

Learning a predictive model of the environment has long been recognized as a more sample-efficient approach than model-free policy learning. Predictive models were applied to discrete grid-world tasks in Dyna-Q (Sutton, 1990), (Chentanez et al., 2005) and to low-dimensional robotic control tasks in PILCO (Deisenroth & Rasmussen, 2011) . In (de Avila Belbute-Peres et al., 2018) , a linear complementarity problem solver is used to simulate two-dimensional rigid body dynamics as a differentiable layer within a network, allowing sample-efficient prediction and high performance on control tasks. Other methods build models of environment dynamics which are then used to plan for control tasks (Lowrey et al., 2018) (Chua et al., 2018) . Many recent methods use generative deep network models trained to predict transitions in high-dimensional visual environments. Planningbased agents have been constructed using networks trained to simulate grid-based games such as Sokoban (Racanière et al., 2017) , Atari games (Kaiser et al., 2019a) , pixel-based physics tasks (Hafner et al., 2018) and visual robotic control tasks (Ebert et al., 2018) .

6. CONCLUSIONS:

The capability to generate and reason about counterfactual states is already considered valuable when deploying agents that require interpretability or explainability (Wachter et al., 2017) . We have shown how counterfactual states can also be used during the training process, in order to force desired causal dependencies or independencies that may be appropriate for a given task or environment.

We have proposed counterfactual regularization terms appropriate for two-dimensional video tasks, but loss functions based on counterfactual states provide a flexible and general way to encode desired properties into a loss function. Future work might explore new applications of counterfactual states to enforce desired causal relationships during the training process. Additional future work includes extending counterfactual regularization to new architectures and more complex physical control tasks, as well as considering more sophisticated exploration policies for generating training data, and better planning algorithms at test time.

:



Figures and Tables:

Figure1: Left: The graphical model representing the generative process. Shaded nodes are observed, and empty nodes are not observed. In this example, the state at time t = 2 has been observed. Right:The latent state space s is represented as a convolutional feature map, and the transition function T (s) is implemented as a fully-convolutional neural network.

Figure 2: Example sequence predictions in the StarIntruders task. The encoder E converts observed pixels to a latent state s. The recurrent transition function T (s, a) predicts future latent states, and a decoder network D decodes states back to pixel space.

Figure 3: During training, in addition to predicting a ground truth sequence based on actions a t , a t+1 ... our model predicts counterfactual trajectories based on alternate actions a t , a t+1 ... for which no ground truth exists. By enforcing desired relationships between alternate simulated trajectories, desired causal relationships can be imbued into the model.

Figure 4: An agent completing the training version of the StarIntruders task. Destroying an enemy generates +1 reward, and destroying a human generates -1. A regularized model-based agent destroys fewer humans than a non-regularized model-based agent or a model-free agent, even in perturbed versions of the environment.



Table content: <table><tr><td>Game Score</td></tr></table>

StarIntruders environment: Per-episode cumulative reward for training and test tasks, average of 10 episodes. The same model-predictive control (MPC) agent is applied using models trained with the standard reconstruction loss, and models trained with each proposed regularization. A model-free method, Rainbow is included for comparison.

Table content: <table><tr><td/><td>.38</td><td>4.92</td><td>5.05</td></tr><tr><td>MPC</td><td>6.50</td><td>7.30</td><td>7.00</td></tr><tr><td>MPC, Action-Control</td><td>7.60</td><td>7.70</td><td>7.60</td></tr><tr><td>MPC, Disentanglement</td><td>7.90</td><td>7.90</td><td>8.00</td></tr></table>

Peer review (iclr2020):
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). Overall the idea is well motivated - incorporating counterfactual reasoning into model based RL has potential to to improve generalization. Also, while the assumptions needed for the regularization to be correct are not always true, they do seem to hold in many cases. Lastly, the results do seem to indicate that generalization is slightly improved when using the proposed forms of regularization. My criticisms are: (1) As mentioned in the paper Action-Control assumes that at every single timestep the agent has potential to change the state. However there may be settings where the agent can always change state, but only a small component of the state. In these cases the states should be quite similar. For example a robot only moving a single object when the state consists of many objects. Also as mentioned in the paper Disentanglement will not work in stochastic environments. One concern I have is that since different environments can violate the assumptions to varying degrees, it seems like actually using the regularization and picking the correct hyperparameter to weight it will be very challenging. (2) The current results are only demonstrated in a single, custom environment. Additionally performance is shown on only 2 test tasks, and in all cases in Table 2 it is unclear how to interpret the reward. Does this performance constitute completing the task? What is the best possible cumulative reward in this case? The performance improvement seems small, but it is difficult to judge without knowing the details of the task. I think the paper would be significantly improved by (1) adding experiments in more environments, especially standard model based RL environments where the performance of many existing methods is known and (2) adding comparisons to other forms of model regularization, for example using an ensemble of models. My current rating is Weak Accept. Some other questions: - In Table 2 does MPC amount to PlaNet? - How sensitive are the current numbers to planning parameters (horizon, num samples)? - Can you provide error bars for the numbers in the tables? ______________________________________________ After author responses and closer examination of the paper I have some additional concerns about experimental details. Changing my score from *Weak Accept* to *Weak Reject*