CausalDyna: Improving Generalization of Dyna-style Reinforcement Learning via Counterfactual-Based Data Augmentation
Authors: Deyao Zhu
Li Erran Li
Mohamed Elhoseiny
Venue: ICLR 2022 Submitted (2022)

Abstract: Deep reinforcement learning agents trained in real-world environments with a limited diversity of object properties to learn manipulation tasks tend to suffer overfitting and fail to generalize to unseen testing environments. To improve the agents' ability to generalize to object properties rarely seen or unseen, we propose a data-efficient reinforcement learning algorithm, CausalDyna, that exploits structural causal models (SCMs) to model the state dynamics. The learned SCM enables us to counterfactually reason what would have happened had the object had a different property value. This can help remedy limitations of real-world environments or avoid risky exploration of robots (e.g., heavy objects may damage the robot). We evaluate our algorithm in the CausalWorld robotic-manipulation environment. When augmented with counterfactual data, our CausalDyna outperforms state-of-the-art model-based algorithm, MBPO and model-free algorithm, SAC in both sample efficiency by up to 17% and generalization by up to 30%. Code will be made publicly available.

Under review as a conference paper at ICLR 2022
CAUSALDYNA: IMPROVING GENERALIZATION OF DYNA-STYLE REINFORCEMENT LEARNING VIA COUNTERFACTUAL- -BASED DATA AUGMENTATION
Anonymous authors Paper under double-blind review
ABSTRACT
Deep reinforcement learning agents trained in real-world environments with a lim- ited diversity of object properties to learn manipulation tasks tend to suffer overfit- ting and fail generalize to unseen testing environments. To improve the agents' ability to generalize to object properties rarely seen or unseen, we propose a data- efficient reinforcement learning algorithm. CausalDyna, that exploits structural causal models (SCMs) to model the state dynamics. The learned SCM enables us to counterfactually reason what would have happened had the object had a differ- ent property value. This can help remedy limitations of real- wol-w environments or avoid risky exploration of robots (e.g., heavy objects may damage the robot). We evaluate algorithm in the CausalWorid robotic- pobotic-manipulation environment. When augmented with counterfactual data, our CausalDyna outperforms state-of- the-art model- based algorithm, MBPO and model-free algorithm, SAC in both sample efficiency by up to 17% and generalization by up to 30%. Code will be made publicly available.
INTRODUCTION
Classical model-free reinforcement learning approaches require a massive amount data collected in the environment work, which slows down its success in tasks where data collection is time- consuming or costly, like robot manipulation. Model-based reinforcement learning (MBRL) meth- ods alleviate this issue by maintaining world model that simulates the real environment. The world model can serve as a surrogate of the real environment for the agent to interact with to reduce the amount of the required conecconu interaction in the real environment. MBRL methods (Kael- bling et al. 1996; Wang et al.) 2219)] Janner et al.) (019) learn from model rollouts of previously observed states. Recently, CTRL (Lu et al.) (020) takes structural causal model (SCM) approach that can generate samples counterfactually had a different action had been taken for a state previously observed. However, these methods are limited for robotic manipulation tasks since the environment is often the key limiting factor. In this paper, perform counterfactual reasoning properties. For example, when the task manipulates objects with different masses, the real environ- ment may not have uniform distribution of object masses. Furthermore, avoid damaging the robot, certain exploration of the gripper torque may be limited during training.
the object
To this end, propose Dyna-style MBRL method, CausalDyna in robotics that improves the policy performance by counterfactual reasoning of physics properties of objects and enriching the diversity of the generated rollouts. We leverage the structural causal model (SCM) to model the state dynamics. CausalDyna can be applied to generate episodes with unseen or rarely seen objects improve the sample efficiency and generalization of the policy.
Our contributions are summarized as follows.
We introduce a novel Dyna-style causal reinforcement learning algorithm, dubbed as CausalDyna that learns from counterfactually generated episodes with intervened object property values.
We compare with state-of-the- the-art model based reinforcement learning algorithm, MBPO and model free algorithm, SAC the World environment. Experimental results

---

Under review as a conference paper at ICLR 2022
Real Environment
Real Environment
Original Object
Counterfactual Property
Modified
Object to Interact
Object
state
reward
reward
action
action
World Model
Robot
World Model
Robot
Normal Generation
Counterfactual Property Generation
Figure 1: In classical Dyna-style methods, the world model generates episodes starting from a real environment state. Then, our robot can practice in the world model and learn how manipulate the original object. To improve the generalization of the learned policy, we further modified the object property in the state. So the robot has the chance to play with objects with more diverse properties.
show that Dyna outperforms MBPO and SAC on sample efficiency by up to 17% and generalization by up to 30% when manipulating objects with unseen rarely properties.
RELATED WORK
Causal Inference in Reinforcement Learning There an increasing interest causal inference in the field of reinforcement learning. Counterfactullly-Guuee Policy Search (CF-GPS) (Buesing (ct al. 2018) assumes that the real transition, observation, and reward functions are all known. They show that any partially observable Markov decision process (POMDP) can be represented as struc- tural causal model (SCM). Therefore, counterfactual inference can be applied to improve the off- policy evaluation and policy-guided search. Coun Terfactual Reinforcement Learning (CTRL) (Lu et al. 2020) leverages bidirectional conditional GAN to model the environment dynamic for data augmentation. The model takes noise vector input besides the state and action model the ran- domness of the environment. Before generating counterfactual data given alternative actions, they first infer the value of this noise vector. Then, the inferred noise is used to generate predictions with new actions. Causal Partial Models (CPM) (Rezende et al., 2020) studies the causal incorrectness of world models that don't 't condition on the full observation. To fix this issue, CPM introduces backdoor variable that helps the rollout of the model to be causally correct. We propose an SCM framework to model the physics properties of objects across the temporal dimension. addition, we show that generating episodes with counterfactual object properties helps improve the generalization of the learned policy.
Model-Based Reinforcement Learning Model-based Reinforcement Learning (MBRL) ap- proaches have shown a potential to improve the sample efficiency by a large margin compared to classical model-free approaches (Kaelbling et al. 1996 Wang et al. 2019). Autoencoder-bant based
algorithms like World Models (Ha & Schmidhuber 2018) and Dreamer (Hafner et al. 22019 2020)
the world model to better represent the visual observation and faster the policy training. Policy Search with Backpropagation algorithms like PILCO (Deisenroth & Rasmussen] 2001 Deisenroth
al. 2013 Kamthe & Deisenroth 2018) and GPS (Levine Koltun Levine Abbeel 2014]
Montgomery & Levine (2016) train the policy by maximizing the simulated return of the policy in world model. Because the world model differentiable, the policy can be directly trained by gradient descent. Shooting algorithms like PETS-RS (Chua et al., (018)) and MB-MF (Nagabandi jtt 2018)) alleviate the receding horizon problem in model predictive control (MPC). Recent works include Ross & Bagnell (2012), MOPO, (Yu et al.) 2020) and Morel (Kidambi et al. 2020) show that MBRL can work well in the offline RL. setting. Unlike the traditional MBRL that approximates the local transition function, L (Zhang et al. 2021 builds the world model
graph of states for better reasoning ability. Dyna-style algorithms (Sutton 1990; 1991a b) use the learned world model to roll out simulated episodes to reduce the demand for real data for policy training. As
2

---

Under review as a conference paper at ICLR 2022
1
Qt-2
Gc-1
The structure causal model of a robot environment. The time-invariant property is modeled node across the temporal dimension that affects all the causal mechanisms. 8-m,t and ar denotes the time-variant state and the action at the step Ł, respectively.
Figure 2:
recent development of Dyna-style algorithms, ME-TRPO (Kurutach et al.] 2018 uses an ensemble of world models to catch the epistemic uncertainty; MB-MPO (Clavera al. 2018) viewed each model the ensemble as a task and meta-learn a policy that adapts quickly to handle the model-bias issue; MBPO (Janerr al. 2019) rolls out short episodes branched from real data to improve the generation quality. Our method follows the Dyna-style framework and targets designing and using causal world model to generate better and more diverse rollouts in robotic environments.
BACKGROUND
3.1
STRUCTURAT. CAUSAL MODEL.
Structural Causal Model (SCM) is widely used framework to describe the causal mechanism of system. Let's denote X {x1, ...,
..., X, as the set of N variables system. Knowing their causal relationships allows us to build directed acyclic causal graph to describe this system. Each node represents a variable, which is directly caused by parent nodes. In this way, a node x,, can be modeled as the following function:
X (i(Paata(Xa)) un.)
(1)
Here, Paoos(X.) denotes the observed parent nodes of X,, u, is a noise that represents the effect of omitted factors. This function is also called causal mechanism. SCM is the set of these causal mechanisms that describes the whole system. SCM defines joint distribution of the variables p(x1, XN) following the causal Markov assumption: given its direct causes, each variable x,, independent of other indirect causal variables.
3.2 DYNA-STYLE MODEL-BASHE REINFORCEMENT LEARNING
Dyna-style model-based reinforcement learning uses the world model to roll out simulated episodes, which can be viewed data augmentation. The training of Dyna-style MBRL is composed of three steps: First, the agent interacts with the real environment and collects real data to train the world model. Then, this world model is used interact and collect simulated data. After that, the agent can be trained together with the real and the simulated data using classical model-free reinforcement learning algorithms. These three steps are executed repeatedly until the training converges. In case we apply Dyna-Style algorithm on RL algorithms with experience-replay buffers and would like to collect whole simulated episodes, as the world model is trained to only approximate the transition of the environment p(s;-1 1 az), we need an initial state to start the simulated episodes. A usual way to solve or randomly sampled state from the collected real episodes as the start point of the simulated episodes. As the real episode already contains the future of 8, under the original action sequence {a:, l...) executed in this episode, generating new simulated episodes starting from under different action sequences can be viewed as answering a counterfactual "what if" question: What would happen if the agent behave differently this time instead of doing {a,, at+1,....? The world model gives the agent a chance to figure out the answer without interacting in the real environment, and helps the agent learn faster.
simulator of the real environment for the agent to
using the first state
3

---

Under review as a conference paper at ICLR 2022
Algorithm 1: Counterfactual Property Generation
Data: Rollout length K, Real experience buffer D,, Policy Pa. World model pww-
Counterfactual property space M, Empty episode buffer IB
Result: TB
Sample state [s-mi m] from the real experience buffer ID,, (append(ss)
2 Sample counterfactual property value mCF from M, set = [s_mi mcF
for K steps do
~ Pr a 8). s' pw (s'ssssss.0) ss a) B.append(a, 8'), 4-
end
4 METHOD
4.1
STRUCTURE CAUSAL MODEL OF A ROBOT ENVIRONMENT
Let's consider environment where robot needs to manipulate object. We can describe this environment using different states. Many of these states are changing over time, including object position and the end-effect position. It is important to model them as they directly contain the dynamic information of the environment. Some other states are time-invariant, like the object mass or the floor friction coefficient. Although their values are fixed, they determine the environment dynamics and affect how other time-variant states change. Let's denote the total state at step f as 8t- m m.] is the concatenation of the time-variant state 8-m,t I at step t and the object time- 8, m invariant property m. The motor torque to execute at step is denoted as a,. As shown in F can build structural causal model (SCM) to describe this environment. The time-invariant property m is modeled as a fixed node across the temporal dimension, which affects all the causal mechanisms.
4.2 COUNTERFACTUAL PROPERTY GENERATION
Policy generalization ability is essential as the testing environment of the policy is not always the same as the training environment. For example, when learning lift an object, the robot might only interact with objects whose masses are in a suitable range. Lifting frequently too-heavy object might reduce its service life, and most reinforcement learning algorithms need large amount of interaction data to work. However, knowing how to lift a heavy object is still desirable when deploying the robot. A typical Dyna-style method generates simulated rollouts branching from a starting state seen in previous real episodes. If the world model takes physics properties input, is possible to go a step further and intervene in these properties. For example, we could modify the mass of an object in the world model to make it heavier. So the agent can learn to manipulate them in the world model as much as we want without harming its service life. Inspired by this, design simple generation strategy to enrich the simulated rollouts by modifying the original object's property to improve the policy generalization. Concretely, instead of taking starting state 8; (s m] sampled from real episodes as it is like most of the Dyna-style methods, we replace the object property m by desired counterfactual value mcF sampled from a predefined counterfactual property space M before rolling out the simulated episodes. We name this type of episodes generation as counterfactual property generation, illustrate it in Fig and show the process
A
4.3 TRAINING PROCEDURE
The training of our model follows the Dyna-style model-based reinforcement learning framework. The world model an additional imperfect substitute for the real environment for the policy to interact with. The policy is still trained using the traditional model-free reinforcement learning approach, but the data for training is a mixture the data from the real environment data and that maintain two replay buffers. The real
from the world model. During the training procedure, experience replay buffer D,, stores the interaction data from the real environment. The world model is trained using the real experience replay buffer only. The simulated episodes from the world model
4

---

Under review as a conference paper at ICLR 2022
Algorithm 2: Training Procedure
Data: Policy Pr- World model Pw Empty real experience replay buffer D,, Empty episode
buffer B, Rollout length K, Counterfactual property space M, Counterfactual generation
ratio a
Result: Trained Policy Psr
Prefill Di, by executing the untrained policy pa in the environment
while Not Converge do
Split DD into a training set Dr.rain and holdout set hootooout randomly
Train the world model PwM M on Dy,tainn until converge on hoooldoo..
Empty the simulated experience buffer DD.
Generate %% N simulated episodes with K steps by counterfactual property generation as Alg I to D3
Generate (1 a%) yNy simulated episodes with K steps with original property to D,
for E steps do
Collect step of data the real environment; add it to D, Update policy parameters via SAC on the combination of D, and DD, for G steps
11
end
12 end
are stored in the simulated experience buffer D., which is used to train the policy net and the real experience replay buffer DD,. The whole training procedure is shown in Alg The policy is trained via soft actor-critic (SAC) (Haarnoja et al.) 2018)) using the data from both the real experience buffer Dr and the simulated buffer ID,, As we generate the simulated episodes with counterfactual property and we following the Dyna-style MBRL framework, we name our model CausalDyna.
World Model Training Each time the world model is trained, the real experience replay buffer D, is split into a training set, and a holdout set randomly. The world model trained predict the next state $t-1 by maximizing the log-like loglliiel given the current state s; and the action in the training set until converging measured by the holdout set.
Augment Data Collection We adopt the generation strategy of model-based policy optimization (MBPO) (Janner et al.) 2019) to roll out the world model. The simulated episodes start from a real state randomly sampled from the real experience replay buffer W,, and are rolled out for K steps. We generate two types of simulated episodes: a% of the rollouts generated with counterfac- tual property generation, where we intervene the object property as described in Alg to generate episodes with different objects. The remaining (1 - a%) episodes generated using the original property. Each time N, simulated episodes are generated in total. Note that each time we collect the simulated episodes, all the previous data in the simulated experience buffer D, is discarded as the world model generated them few training steps before and are not 'fresh' anymore.
EXPERIMENTS
5.1 BENCHMARK
We evaluate our method CausalDyna on a recently proposed robotic benchmark CarlalWorld (Ahmed et al. 2020). CausalWorld is designed for causal structure and transfer learning in a robotic manipulation environment. The robot in Corldl a 3-finger gripper. Each finger has three joints. The mission of the robot is to move objects to specified target locations. The observations of the WorlalWorld we use includes the time stamp i, the robot state 8,, the object state 8,, the time- invariant property m, and the goal information 8,. The robot state is consists of 9 joint positions, 9 joint velocities, and the Cartesian coordinates of the three end-effectors (fingertips). The object state so contains the Cartesian coordinate, the velocity, the quaternion orientation, and the object's angular velocity. The property m includes the object mass and the friction coefficient. The goal information 8y contains the target location and orientation of the object.
5

---

Under review as a conference paper at ICLR 2022
Evaluated Models We evaluate three approaches in experiments: Model-Based Policy Op- timization (MBPO) (Janner et al., 2019), one of the state-of-the-art Dyna style methods with high sample efficiency, Soft Actor-Critic (SAC) (Haarnoja et al. 2018). widely-used model-free ap- proach, and our method CausalDyna.
Task Settings and Performance Metries We define three settings to evaluate our method: Picking Mass, Pushing Mass, and Pushing Friction. In Picking Mass, the robot needs to pick up an object to a target location in the air. The object mass is different over different episodes. In contrast, the target locations in Pushing Mass and Pushing Friction are the ground. The object mass and the floor friction in Pushing Mass and Pushing Friction are different over different episodes, respectively. We use the default reward signals of CausalWorld signals to encourage the robot to get close to the object and move it toward the target. The reward weighted sum over the reduction of the distance between the end effectors and the object and the distance between the object and the target. We evaluated our approach and competing methods using fractional success rate (FSR), which defined as the overlapping ratio between the object and the target. We compute the FSR of given episode as the average FSR over the last 20 steps. To quantify the sample efficiency in our benchmark, we propose a metric named Area-Under-the Curve Ratio (AUCRatio). Given learning curve FSR number of the environment steps collected already, AUCRatio until step Nstep is computed as 2q As 0 S FSR S 1, a policy with AUCRatio 1 means it can perform the task perfectly without training.
train our method. The reward provides rich
(Nstep))))) where Natey denotes the
Natep
AUCRatio
fiearn (Natep)
(2)
Notep Detep==
5.2 EXPERIMENTS WITH
PROPERTY
An intelligent robot might encounter various objects when deploying, If the robot need to manipulate an object unseen during training, its performance might be reduced. This can be viewed as an out- of-distribution problem: how to generalize well to the object not in the training distribution? The counterfactual property generation approach has the potential to increase the performance on objects with unseen property values if we roll out simulated episodes with object property that is out of the training range. To verify our assumption, we create an experiment to study whether our method helps improve the agent performance training. In detail, in our Picking Mass and Pushing Mass setting, the robot is trained with objects of which the mass is uniformly distributed from 0.015kg 0.045kg. But during the testing stage, the robot is asked to interact with heavier objects up to 0.1. In Pushing Friction setting, the friction coefficient is from 0.3 to 0.6 during training. And the robot is deployed also handle friction from 0.6 to 0.8.
objects whose property value is not encountered during
As target the performance of the objects with unseen property value during training, we our method here to imagine these objects. In detail, when the counterfactual property generation is applied, replace the original property value with counterfactual value uniformly sampled from the unseen test range. In this way, our agent can practice manipulating these unseen objects in the world model in advance.
Hyperparameters The length of the simulated episodes K is 10. bootstrap ensemble of world models is used following Kurutach et al. (2018) for both MBPO and our method. The ensemble size For each generation step, we randomly pick one model from the ensemble to predict the next state. When training the policy, 20% of the training data are from the real experience replay buffer. The remaining are from the simulated episodes. In our CausalDyna, 20% of the simulated episodes are generated by counterfactual property generation (a in Alg 2}. We use Adam (Kingmal & Ba 2014) as the training optimizer for all experiments. All the models we evaluated are trained for 1.2 million steps in Picking Mass and 600 thousand steps in Pushing Mass and Pushing Friction. Each model in this experiment has 5 training cases. The model architecture and the remaining hyperparameters can be found in Appx and Appx|B
and

---

Under review as a conference paper at ICLR 2022
0.8
Fractional Success Rate
0.8
9.6
0.7
0.4
0.4
Fractional
0.5
CausaiDyns: 0.47
0.2
Causalyyne 0.87
Causalyyns
MBPG: 6.37
NBPP: 0,70
3.4
NBPG: 0.53
0.0
SAC: 0.30
SAC: 0.25
SAC: 0.48
0.02
0.04
9.06
0.08
0.30
0.02
0.04
0.06
0.30
0.3
0.4
0.5
0.7
0.8
(a) Picking Mass
(b) Pushing Mass
(c) Pushing Friction
Figure 3: Experimental results of counterfactual property generation in the out- -itto---iitributiinn ex- periment. The vertical black line shows the boundary between the and unseen property values during training. The left part is the seen region. Counterfactually generating the simulated episodes with unseen property value helps alleviate the performance drop when evaluating unseen property during training. Numbers in the legend denote the average performance in the unseen value range. Each curve contains 5 training cases.
1.0
1.0
1.0
NBPO:
Practtoral SLc00s
Enscttonal Success
2.6
8.8
0.6
0.4
0.4
0.0
0.8
0.0
140
(a) Picking Mass
(b) Pushing Mass
(c) Pushing Friction
Figure 4: The learning curve of the evaluated models on the training property range in the out- of-distribution experiments. CausalDyna converges as fast as MBPO, although it generates 20% less simulated episodes in the training property range. Numbers in the legend denote the average AUCRatio.
Performance The experimental results are shown in Fig The vertical black line denotes the boundary between the seen and unseen values during training. The left part the seen region. In Picking Mass and Pushing Mass, the performance of all the methods declines when the object mass out the training range. Moreover, the performance reduction is more significant when the tested object mass is farther away from the training range. Our CausalDyna alleviates this performance reduction the unseen range by a large margin compared to MBPO. In Picking Mass, CausalDyna improves the unseen FSR by 24% from 0.37 0.47. For Picking Mass it is 27% from 0.7 to 0.87. This indicates that hallucinating episodes with unseen objects during training helps improve the generalization ability of the policy. In Pushing Friction, CausalDyna achieves similar performance as MBPO since the unseen range performance reduction here is not obvious. As SAC is less sample efficiency than both model-based methods, SAC cannot achieve compatible results given the same training data as MBPO and CausalDyna. Note that in Picking Mass, although our Causal performs better than MBPO the outto--dist range, the absolute performance not high when the object is too heavy (like 0. 0kg). This might be caused by the reduced performance of world model when counterfactually generating episodes with unseen objects. A better world model design that can better understand the physics and reason the future more causally might help alleviate this issue when combined with our method. We leave this for future research.
Sample Efficiency We show the learning curve of MBPO, CausalDyna, and SAC this experi- ment in FigHt Although augment 20% fewer simulated episodes in the original property range

---

Under review as a conference paper at ICLR 2022
0.56
0.55
0.94
Fractional Success Ratt
0.50
0.92
0.7
0.6
9.88
Fractoonal
0.00
0.00
9.84
Model
CausalDyna: 0.73
GasselDyna:
CousalDyna:
0.96
0.92
MBPO: 0.92
SAC
SAC 8
SAC 0:87
0.002
0.81 Moss
0.05
0.092
0.01
0.05
0.3
0.55 Friction
0.8
Mass
(a) Picking Mass
(b) Pushing Mass
(c) Pushing Friction
Figure 5: Experimental results of counterfactual property generation in the unbalanced distribution experiment. When counterfactually generating episodes where the object is less encountered during training, CausalDyna helps improve the policy performance and tail values. For each property, the median value occurs 90% of the time in the environment, and the rest values share the remaining 10% equally. Numbers in the legend denote the average performance over the tail values. Each model has training cases.
both the objects with head values
compared to MBPO, CausalDyna converges as fast as MBPO in the original training range. Results indicate that method improves the -of--ii-iittitttionn performance without sacrificing the sam- ple efficiency. The model-free SAC training is much slower than MBPO and CausalDyna, as SAC doesn't 't have simulated data to train on.
5.3 EXPERIMENTS WITH UNBALANCED TRAINING DISTRIBUTION
In real environments like warehouses, the numbers of different wares unequal, and sorting robot might manipulate some objects less frequently. This can be described as an unbalanced train- ing distribution. If the training distribution is heavily unbalanced and some objects are significantly less encountered than others during training, counterfactually generating episodes with such objects might help improve the policy performance on them. We create simple heavily unbalanced training distribution consisting of head property value and two tail property values to verify this assump- tion. The object property in 90% of the training episodes equals the head value. The two tail values share the remaining 10%, each value obtains 5%. Concretely, in Picking Mass and Pushing Mass, we have three different objects with mass values 0.002kg, 0.01kg, and 0.05kg, respectively. 90% of the time, the robot sees and manipulates the object with the median mass value of 0.01kg. The robot plays with the heavy 0.05kg object and the light 0.002kg object equally in the remaining time. For Pushing Friction, the three friction coefficients are 0.3, 0.55, and 0.8 that occur in 5%, 90%, and 5% of the time, respectively. In the testing stage, models need to perform well on all three property values.
As the objects with tail values occur less frequently in the training stage, CansalDyna in this experi- ment imagines what would happen if the given head object is the tail. Concretely, when CausalDyna generating simulated episodes, the property value of original objects are counterfactually modified one of the tail property values randomly. Therefore, the agent interact with the tail objects more in the world model to improve the tail performance.
Hyperparameter In CausalDyna, 2/3 of the simulated episodes are generated by counterfactual property generation (a in Alg(). All the models all the 3 settings trained for 600 thousand steps. Each model in this experiment has 6 training cases. The remaining hyperparameters are the same in the previous experiment.
Performance As shown in Fig 5 the performance on the head property value (0.0lkg for mass and 0.55 for friction) is better than the tail property values for all the methods in all the 3 settings. How- ever. CausalDyna improves the performance the tail property and shows the smallest performance difference between the head and the tail among the three models. For example, the performance gap between the head and the tail of CausalDyna in Picking Mass is about 0.1, much smaller than MBPO
8

---

Under review as a conference paper at ICLR 2022
0.5
Cyusa MBPO:
Frectonall Succese Rate
0.8
Frectondl Success Ratte
0.
0.6
0.5
0.8
Frocional
0.3
0.2
Cassa Dyna: 50%
Causa Dyne: 59%
MaPO: 18%6
MBPO: 42%
0.0
0.8
SAC: 31%
0.0
SAC: 40%
Siaps
1:5
145
Steps
105
(a) Picking Mass
(b) Pushing Mass
(c) Pushing Friction
Figure 6: Average policy performance at different environment steps. Our method CausalDyna, which counterfactually generating episodes where the object is less frequently encountered during training, reduces the required amount of environment steps and shows the best sample efficiency in the unbalanced training distribution experiment. Numbers in the legend denote the average AUCRa- tio. Each model has 6 training cases.
(0.2-0.3), and the tail performance is increased by 30% from 0.56 to 0.73. Besides, we notice that CausalDyna improves the policy performance on both objects that are less frequently seen during training and the head objects compared to MBPO. This might be because learning how to behave well in the tail cases helps the model better understand the environment dynamics and improves overall performance. In addition, the performance variance in Pushing Mass and Pushing Friction of CausalDyna is much lower than the other two methods, which suggests that the performance of CausalDyna is more consistent than other methods. With the same amount of training data as MBPO and CansalDyna, the model-free SAC's performance worse than the model-based MBPÒ and CausalDyna, which is the same as the out-of- experiment.
Sample Efficiency The learning curves of the evaluated models are shown in Fig 6 The fractional success rate is uniformly averaged over all the property values. CausalDyna shows better sample efficiency and converges faster. In all three settings, CausalDyna requires about 100k fewer envi- ronment steps to converge compared to MBPO and increase the sample effiency by about 17%. This might be because CausalDyna has more simulated episodes with the tail property values to train the agent, which helps the agent understand the task better and adapt to all the property values faster.
6 CONCLUSION AND FUTURE WORK
In this paper, we focus on improving the generalization ability of model-basee reinforcement learn- ing in robotic environments. We propose a novel Dyna-style causal reinforcement learning algorithm named CausalDyna that rollouts episodes with intervened object properties. CausalDyna leverages the diversity of the simulated episodes augmented by the world model and improves the generaliza- tion of the policy when manipulating objects with property unseen or rarely seen during training. Experiments show that our method helps the robot generalize to objects with unseen property val- better. In addition, when the training distribution is unbalanced, our method requires fewer environment steps to converge and performs better with rarely seen objects.
To our knowledge, we are the first to propose counterfactual reasoning on environment properties to improve the generalization of reinforcement learning. We believe this is a promising direction to solve many complex reinforcement learning tasks where the policy generalization ability is essential. When combined with model predictive control and counterfactual reasoning on actions, it is possible to further improve sample efficiency and generalization of RL algorithms. One limitation of our method is that the quality of our counterfactual episodes depends on how well our world model understands the enviromment. We plan to design a better world model that takes prior knowledge like simple physics laws into account. Finally, we have assumed that the properties in environment are fully observable in our current work. We plan to investigate causal models with latent variables representing unobserved properties of the environment.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Ossama Ahmed, Frederik Träuble, Anirudh Goyal, Alexander Neitz, Yoshua Bengio, Bernhard Schölkopf, Manuel Wüthrich, and Stefan Bauer. Causalworld: A robotic manipulation bench- mark for causal structure and transfer learning. arXiv preprint arXiv:2010.04296 2020.
Lars Buesing, Theophane Weber, Yori Zwols, Sebastien Racaniere, Arthur Guez, Jean-Baptiste Lespiau, and Nicolas Heess. Woulda, coulda, shoulda: Counterfactually-guie policy search. arXiv preprint arXiv:/811.0277 2018.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learn- ing in a handful of trials using probabilistic dynamics models. arXiv preprinl arXiv: 1805. 12114, 2018.
Ignasi Clavera, Jonas Rothfuss, John Schulman, Yasuhiro Fujita, Tamim Asfour, and Pieter Abbeel. Model-based reinforcement learning via meta-policy optimization. In Conference on Robot Learning, pp. 617-629. PMLR, 2018.
Mare Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on machine learning (ICML-11), pp. 465-472. Citeseer, 2011.
Mare Peter Deisenroth, Dieter Fox, and Carl Edward Rasmussen. Gaussian processes for data- efficient learning in robotics and control. IEEE transactions on pattern analysis and machine intelligence, 37(2):408-423 2013.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint 1rXiv:1803.10122, 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Confer- ence on Machine Learning, pp. 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: behaviors by latent imagination. In International Conference on Learning Representations, 2019.
Learning
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv:2010.021933 2020.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. arXiv preprint 1rXiv:1906.08253, 2019.
Leslie Pack Kaelbling, Michael Littman, and Andrew W Moore. Reinforcement learning: A survey, Journal of artificial intelligence research, 4:237-285, 1996.
Sanket Kamthe and Mare Deisenroth. Data-efficient reinforcement learning with probabilistic model predictive control. In International conference artificial intelligence and statistics, pp. 1701- 1710. PMLR, 2018.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. arXiv preprint trXiv:2005.0595/ 2020.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXhv:1412.6980, 2014.
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. arXiv preprint arXiv:/802. 10592, 2018.
Sergey Levine and Pieter Abbeel. Learning neural network policies with guided policy search under unknown dynamics. In NIPS, volume 27, pp. 1071-1079. Citeseer, 2014.
Sergey Levine and Vladlen Koltun. Guided policy search. International conference on machine learning, pp. 1-9. PMLR, 2013.
10

---

Under review as a conference paper at ICLR 2022
Chaochao Lu, Biwei Huang, Ke Wang, José Miguel Hemández-Lobate Kun Zhang, and Bernhard Schölkopf. Sample-efficient reinforcement learning via counterfactual based data augmentation. arXiv preprint arXhv:2012.09092, 2020.
William H Montgomery and Sergey Levine. Guided policy search via approximate mirror descent. Advances in Neural Information Processing Systems, 29:4008-4016, 2016.
Anusha Nagabandi, Gregory Kahn, Ronald S Fearing, and Sergey Levine. Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning. In 2018 IEEE Interna- tional Conference on Robotics and Automation (ICRA), pp. 7559-7566. IEEE, 2018.
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In leml, 2010.
Prajit Ramachandran, Barret Zoph, and Quoc V Le. Searching for activation functions. arXiv preprint arXiv:1110.0544 2017.
Danilo Rezende, Ivo Danihelka, George Papamakarios, Nan Rosemary Ke, Ray Jiang, Theophane Weber, Karol Gregor, Hamza Merzic, Fabio Viola, Jane Wang, et al. Causally correct partial models for reinforcement learning. arXiv preprint arXiv:2002.02886, 2020.
Stephane Ross and Andrew Bagnell. Agnostic system identification for model-based reinforcement learning. arXiv preprint arXiv: 1203. 2012.
Richard Sutton. Integrated architectures for learning, planning, and reacting based on approxi- mating dynamic programming. In Machine learning proceedings 1990, pp. 216-224. Elsevier, 1990.
Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4):160-163, 1991a.
Richard S Sutton. Planning by incremental dynamic programming. In Machine Learning Proceed- ings 1991, pp. 353-357. Elsevier, 1991b.
Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang. Pieter Abbeel, and Jimmy Ba. Benchmarking model-based reinforce- ment learning. arXiv preprint arXiv; 1907000977 2019.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. arXiv preprint 20X5::2005.13239 2020.
Lunjun Zhang, Ge Yang, and Bradly C Stadie. World model graph: Learning latent landmarks for planning. In International Conference on Machine Learning, pp. 12611 -12620. PMLR, 2021.

---

Under review as a conference paper at ICLR 2022
A MODEL ARCHITECTURE
Here we list the architecture of the world model, the policy actor net and the policy critic net we use all experiments for all methods. All the models are built using linear-layers. The world model Swish activation function (Ramachandran et al. 2017) and the policy uses ReLU (Nair & Hinton 2010
Table 1: Model Architecture
Modules
Hidden Layers
Neurons Per Layer
World Model
3
200
Policy Actor
2
256
Policy Critic
2
256
B HYPERPARAMETER
The size of the real experience replay buffer D, is 100k for MBPO and our method Cyna in all three settings. For SAC, it is 1M as we notice SAC with 100k-size replay buffer cannot be trained well. For world model training, The replay buffer D, is split randomly into training set Dr.rain with 80% of the data and a holdout set ID, holdoutt containing the remaining data. We train the model once for every 250 real environment steps until converge evaluated on the holdout set. The learning rate is 3e-4. Batch size is 256. For the policy training, the policy net is updated for iterations per real environment step. The batch size is 256, and the learning rate is set to le-4.
C QUALITATIVE RESULTS
Here we demonstrate episodes from CausalDyna and MBPO in the Pushing Mass setting in unbal- anced training distribution experiments with the heavy tail object in Fig and Fig Both models are trained for 600k environment steps. The object to manipulate is in blue color. Target location is shown as the green shade. Each column corresponds to episode. CausalDyna generalizes the heavy tail object well and pick it to the location successfully shown in Fig 7 while MBPO fails to lift the object up episodes shown in Fig
12

---

Under review as a conference paper at ICLR 2022
Figure 7: Cyns with the heavy tail object. Pushing Mass, Unbalanced Training Distribution.
13

---

Under review as a conference paper at ICLR 2022
Figure 8: MBPO with the heavy tail object. Pushing Mass, Unbalanced Training Distribution.
14


Reviews:
[{"review_id": "pGRBa7WTQvS", "reviewer": "Reviewer_CWLF", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- connects data augmentation to counterfactual property generation\n- clearly written\n- What is novel about applying counterfactual data augmentation to DYNA, as opposed to standard data augmentation techniques in other areas of machine learning (e.g. computer vision) is that the goal of counterfactual data augmentation should be to learn a model that is _equivariant_ to data augmentation, whereas standard data augmentation techniques mostly focus on making the model _invariant_ to data augmentation. Therefore, this is a novel problem setting worth studying.\n\nWeaknesses\n- The authors predefine the counterfactual property space $M$, which avoids the hard problem of learning the counterfactual property space from experience. But this is a very restrictive assumption: it's not merely that the authors assume that the mass properties are observable, the authors assume the entire range of masses is observable as well. Because the authors assume access to this counterfactual property space, their empirical experiments amount to comparing a method with additional data augmentation to methods that do not have this privileged data augmentation, which seems to not be an interesting question to ask, since it is well known that data augmentation, especially if you have knowledge of what the augmentations should be, would help performance.\n-  What would be a more interesting question to ask is, given that we know data augmentation would help, how can we supervise the world model to be equivariant to different data augmentations? However, the paper does not discuss how they supervise the world model to make the correct prediction for different values of $m_{cf} \\in M$. In particular, if $m_{cf}$ is intended to take on values of out-of-distribution $m$ from the perspective of the policy, how can we guarantee that the model will actually learn to be equivariant to the data augmentation to generate useful data for training the policy? This seems to be the most crucial point, because the paper seems to assume that the world model will make the correct prediction, but it is not discussed how the paper supervises the world model to do so.\n- Figure 5 does seem to suggest that the world model does seem to predict slightly better with the proposed data augmentation, but why would it do so without supervision on the correct prediction for different values of $m_{cf} \\in M$? This would be an interesting question to investigate, but this investigation and analysis is missing from the paper.\n- On the other hand, Figure 3 and 4 seem to suggest that the proposed data augmentation does not help at all compared to MBPO (the error bars overlap enormously). This result seems to refute the author's hypothesis that counterfactual data augmentation is crucial, especially since in Figure 6 MBPO does not perform that much worse than CausalDyna.\n", "summary_of_the_paper": "This paper considers the problem of generalization of a dynamics model to different versions of the same environment, indexed by parameter m. For example, this parameter $m$ might be the mass of an object to be picked up. The problem that the authors tackle is to generalize the agent's policy to environments with either different parameters $m$ or a different distribution of parameters $m$. The solution the authors propose is to randomly sample different values of $m$ during training of the model, as opposed to using the $m$'s observed from the agent's rollouts. Empirically we observe that the proposed method performs similar to MBPO in generalization to out-of-distribution $m$ (Section 5.2) and that generalizes better on heavier masses than MBPO in the unbalanced distribution setting (Section 5.3).", "summary_of_the_review": "The problem setting is good, but the main conceptual weakness is a lack of explanation for how the model is supervised to be equivariant to different data augmentations. Furthermore, empirical results seem to suggest that the proposed data augmentation does not have much of an effect on performance, which seems to undermine the motivation for this paper.\n\nEDIT (AFTER REBUTTAL)\nThank you to the authors for your response. Based on the authors' response, it appears that the majority concerns I had raised will be addressed in the next iteration of the paper. Therefore, I maintain my original score and look forward to the new and improved version.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "nTvG8d2nHB", "reviewer": "Reviewer_cJzu", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n- The idea of the paper is neat and intuitive. Improving generalization ability via causal learning has been a growing research area. This paper provides a solid demonstration of its promise in the context of model-based RL.\n- The paper is very well written. The experimental analysis is rich and sound.\n\nWeaknesses\n- The claimed connection between the SCM and the proposed dynamic model seems vague. From what I've understood, the proposed method is essentially pretty much the same as the standard dynamic model with state inputs, and then performs interventions over a particular variable. It is not clear how SCM comes into play.\n- The assumption made in the paper is very strong. It's known that learning an SCM (or a dynamic model that allows for counterfactual intervention) is a very challenging task by itself. Yet, the paper assumes it can be easily solved. In particular, it assumes\n    - full access to the high-level variables (e.g., `object mass` and `friction coefficient`) in the true generative process\n    - the learned dynamic model can generalize to unseen settings within a `predefined counterfactual property space`\\\n  These assumptions do not look realistic in the real-world.", "summary_of_the_paper": "The paper presents a method to improve the generalization of model-based RL by means of interventional data augmentation. The key idea is to intervene the value of a particular variable (e.g., object property) in the learned dynamic model for episode simulations. Experimental results show that it improves (i) the generalization ablity in the OoD scenarios with respect to the intervened variable, (ii) sample efficiency in the presence of unbalanced training distribution.", "summary_of_the_review": "The paper well demonstrates the benefits of counterfactual data augmentation for model-based RL. However, the technical contribution seems limited and involves very strong assumptions. I, therefore, consider it below the acceptance threshold.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "kIM-uGcOOzX", "reviewer": "Reviewer_VSvt", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Enhancing the diversity of training data (here, time-invariant properties of the state) to improve the robustness and extrapolation capacities of an RL agent is a good idea but is standard practice known as domain randomization.  In the second application, countering data imbalance, this amounts to conventional data up-/downsampling.\n\nThis paper motivates data augmentation / domain randomization and data up-/downsampling in terms of causality theory.  However, this perspective provides no additional insight; the structural causal model it introduces does not appear to be used by the method at all.  Moreover, the presentation does not cleanly separate counterfactual reasoning from intervention.  Altering state is a form of intervention; calling it counterfactual because such state was not seen in real data dilutes the notion of counterfactual reasoning (is all simulated training counterfactual?).\n\nMethodologically, the greatest weakness of the method, acknowledged by the authors, is that there is no way to train the model on altered data.  Thus, the performance of the policy on these altered data hinges on the extent to which the model, trained without such data, happens to make accurate predictions.  In other words, if the trained model does not extrapolate, the entire method is ineffective.\n\nThe method is compared to two representative baselines, but the most important baseline is absent: an ablated CausalDyna that does not alter any state.\n\nSome Details:\n\n- Figure 1 is not referred to in the text.\n\n- Section 3.2 misrepresents Dyna by separating model training on real data, collecting simulated experience, and training the value function / policy into three distinct steps.  In the original formulation, there is no experience-collection phase, and model training and value function / policy training on real and simulated data proceed in a parallel or interleaved fashion.  The separate data-collection phase only makes sense with RL algorithms that use experience-replay buffers.\n\n- Why is the FSR averaged over the last 20 steps? Are the trained policies unable (or not trained) to bring the object to rest?\n\n- In Eqn. 2 the factors of 1 in the numerator and denominator are superfluous.\n\n- \"5 training cases\": Only later it becomes implicitly clear that this is intended to mean 5 independent runs from which means and (presumably) standard deviations are computed.", "summary_of_the_paper": "This paper addresses the problem of out-of-distribution states encountered in reinforcement learning, and proposes a data augmentation method for Dyna-style algorithms that amounts to domain randomization. The central idea, altering elements of the state vector while training on the model, is motivated from a causality perspective.\n", "summary_of_the_review": "- The paper claims to use structural causal models but this is the case only in a trivial sense at best.\n\n- The causality perspective provides no novel insight.\n\n- The experimental setup pales in comparison to other uses of domain randomization.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "eyjVlt28sM0", "reviewer": "Reviewer_RJye", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "In general, the paper is well presented, and I can easily grasp the main idea. The major concern lies in the novelty, I think building an environment with a counterfactual technique is not new, which has been proposed before. This paper did not provide enough increments on the previous methods. In addition, I would like to ask how to handle the error produced by the SCM, maybe theoretical analysis can be provided to demonstrate the bound of the agent when being trained on such error-involved SCM. ", "summary_of_the_paper": "This paper aims to design a counterfactual reinforcement learning model to improve the training efficiency of the agent. To achieve this goal, SCM is leveraged to simulate the environment and generate new counterfactual trajectories. To demonstrate the effectiveness of the proposed model, the authors have conducted many experiments.", "summary_of_the_review": "Not quite ready for top conferences like ICLR", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper describes a new method to improve the generalization of model-based RL by means of interventional data augmentation. The key idea is to intervene the value of a particular variable (e.g., object property) in the learned dynamic model for episode simulations. Experimental results show that it improves (i) the generalization ablity in the OoD scenarios with respect to the intervened variable, (ii) sample efficiency in the presence of unbalanced training distribution.

Strengths:

- connects data augmentation to counterfactual property generation
- clearly written
- novel about applying counterfactual data augmentation to DYNA, as opposed to standard data augmentation techniques in other areas of machine learning
- The paper well demonstrates the benefits of counterfactual data augmentation for model-based RL

Weaknesses:

- a lack of explanation for how the model is supervised to be equivariant to different data augmentations.
- empirical results seem to suggest that the proposed data augmentation does not have much of an effect on performance
- The claimed connection between the SCM and the proposed dynamic model seems vague
- The technical contribution seems limited and involves very strong assumptions.
- The structural causal model it introduces does not appear to be used by the method at all.
- the presentation does not cleanly separate counterfactual reasoning from intervention
- he greatest weakness of the method, acknowledged by the authors, is that there is no way to train the model on altered data. Thus, the performance of the policy on these altered data hinges on the extent to which the model, trained without such data, happens to make accurate predictions

All the reviewers voted for rejection. I recommend the authors to use the reviewrs' comments to improve the paper and resubmit to another venue.
Decision: Reject