Causal Influence-Aware Counterfactual Data Augmentation

We thank the reviewer for their valuable feedback. We first address the points under weaknesses and then answer the reviewer’s questions. We address weaknesses with W# and questions with Q#.

W1: Assumption about independence swapped states and goal states:

We thank the reviewer for spotting a lack of detail in the manuscript. We would like to further elaborate on this point.

We agree that in cases where demonstrations were provided to achieve desired goals, augmentations to a given state should be irrelevant to the goal that the trajectory achieves. In the Franka-Kitchen environment, we use the intersection of the uncontrollable sets along extended time windows, which excludes intermediate and final entities under influence that are potentially relevant for the demonstration. This is described in Equation 5) of the manuscript. We added a reference to section 3.3 as well.

In addition, all three experiments (Franka-Kitchen, Fetch-Push and Fetch-Pick&Lift) rely on self-supervised goals: the training data does not provide an explicit goal space, nor desired goals per trajectory. For each experiment, goals for counterfactual states are relabeled by either sampling the final state of the skill (Franka-Kitchen), a random state from any trajectory (Fetch-Push), or a future state within the trajectory (Fetch-Pick&Lift). In any of these cases, as long as all trajectories are augmented before goal sampling, no assumption on independence between desired goals and states is necessary. We added clarifications on goal sampling in the updated manuscript (see Section 5.2, A.1.1 and A.2.1).

Moreover, to the best of our knowledge, the issue of dependence between states and goals is also less of a concern when performing offline reinforcement learning with goal relabeling (as done in Fetch experiments). In this case, a reward function can be used to evaluate success (as done in HER [F1]). Thus, a given trajectory does not need to successfully achieve its goal (in other words, the dependency between states and goals does not need to hold).

Therefore, in Fetch experiments, the joint distribution over states and goals at training time only needs to assign a significant probability mass to state-goal pairs observed at test time. As goals are relabeled from augmented states, improved coverage over the joint state space at training time also results in improved coverage over the joint goal space. Finally, coverage over the space of state-goal pairs can be controlled through the relabeling strategy.

[F1] Andrychowicz, Marcin, et al. "Hindsight experience replay." Advances in neural information processing systems 30 (2017).

We thank the reviewer for their valuable feedback. We first address the points under weaknesses and then answer the reviewer’s questions. We address weaknesses with W# and questions with Q#.

W1: Distinction between CAIAC, CoDA, and CoDA-ACTION.

The algorithms differ on two aspects: (1) the extent of causal discovery and (2) the methodology used for causal discovery.

CoDA performs (1) full causal discovery, i.e., it tries to detect influence between all the entities in the causal graph, and (2) uses the attention weights of a trained transformer to infer causation. CoDA-action is a baseline we introduce. It performs (1) partial causal discovery,i.e. it only tries to detect the agent’s action-influence on the entities, and (2) uses the attention weights of a trained transformer to infer causation. Finally, CAIAC performs (1) partial causal discovery and (2) uses state-conditioned mutual information (namely CAI) to infer causation. The novelty of the method is therefore not limited to the mechanism used to infer causation, but should also include the assumption of sparsity in the causal graph outlined in Section 3.1 of the paper, which motivates partial causal discovery.

We further demonstrate these differences in the updated version of the paper. In Section A.5 (Figure 7 and 8) we visualize the computed influences for the Fetch-Push task for all the methods: CAIAC influence scores (namely CAI) , CoDA influence scores (namely the transformer’s attention matrix) and the CoDA-action scores (namely the values of the attention matrix that correspond to the row “action” and columns {“object1, object2 or agent”}). This didactic example displays how using such an implicit measure to detect influence can lead to misestimation (as in CoDA and CoDA-action), whereas CAIAC is better suited for detecting influence in this setting.

We updated the text in the updated version of the paper referring the reader to the appendix sections of interest to clarify the points above.

W2 + W3: A figure that clearly illustrates the difference between CoDA, CODA-ACTION, and CAIAC [...] + Simple, didactic toy task like the SpriteWorld task used CoDA

We hope that the explanations above helped clarify the distinctions between the different algorithms. We updated the text in the updated version of the paper referring the reader to the appendix where we include the computed influence scores mentioned in the point above. We hope that this scenario with only 3 entities (agent, object 1 and object 2) is already a good didactic example.

W4: Empirical results seem weak. In Table 1 [...] struggles in the remaining 3 tasks (Slide Cabinet, Light Switch, Hinge Cabinet).

We thank the reviewer for suggesting to further elaborate on empirical results. We provide this discussion below, and integrate it in the updated version of the manuscript (Section 5.1.2 and Appendix A.2.1).

As stated in the paper, “The results, shown in Table 1, are consistent with the challenging nature of this benchmark, as the evaluated tasks involve OOD settings in terms of states and actions”. More specifically, we found that only a subset of tasks is shown to be performed directly from the initial configuration in the training data. In particular, the training data does not contain state and action trajectories that show how to solve each of the 3 mentioned tasks from the initial robot joint configuration. Instead, these 3 tasks are only demonstrated after solving other tasks, and hence starting from other robot configurations. For this reason, even with perfect counterfactual data augmentation, these tasks remain challenging.

Specifically, out of the 1200 demonstrations in the dataset, containing different task sequences, only 3 objects are shown to be manipulated from the initial robot configuration: 60% of the trajectories solve the microwave task first, 30% show the kettle task first and 10% show the bottom right burner first. This aligns with the relative performance achieved for those tasks. For the 3 remaining tasks, namely the slide cabinet, the light and the hinge cabinet, there is no demonstration shown directly from the initial configuration, and hence the low performance.

Moreover, we remark that this issue adds up to the OOD nature of states and goals at test time, thus defining a fundamentally challenging benchmark.

W6: Empirical evidence on quality of counterfactuals (Part 2)

We provide an analysis on the created counterfactuals using CAIAC. We visually show how the created augmented samples have an increased support of the joint state space distribution in the kitchen dataset, compared to the No-Augmented baseline. In Figure 12 of the updated manuscript, we show 1000 randomly selected samples from the Franka-Kitchen dataset together with one counterfactual for each. The visualization employs t-SNE for a 2D representation of the high-dimensional state space. We observe how the augmented samples using CAIAC cover a much larger space than the observed data, suggesting that we are able to provide states with a richer amount of configurations. However, this does not tell us whether these counterfactuals are actually valid. To quantify whether an augmented datapoint is valid, we do the following: We consider time-windows $(s_t, \dots, s_{t+\tau})$ sampled from the Franka-Kitchen dataset. We create counterfactuals for this window with \method and CoDA, $(\tilde s_t, \dots,\tilde s_{t+\tau})$. We reset the simulator to $\tilde s_t$ and simulate the future with the action sequence in the original trajectory for $\tau$ steps. If the counterfactual is valid, then resulting state of the simulation should coincide with the counterfactual $\tilde s_{t+\tau}$. Unfortunately, the simulator is not perfectly deterministic (or we are unable to set the full simulator state), so measuring exact matching states is not possible. Thus, we run the simulator $K=10$ times with different env-seeds and obtain $\mathcal S' = [ \bar s^k_{t+\tau} ]_{k=1}^K$. We fit a multivariate Gaussian to the set $\mathcal S'$.

We now compute the probability of the counterfactual $\tilde s_{t+\tau}$ under the Gaussian distribution. High probability means valid counterfacturals, low means invalid. In Figure 13 of the updated manuscript we present the histogram of log probabilities. The counterfactuals by CAIAC are mostly valid, which is not the case for the CoDA baselines.

W5: [...] would like to see CAIAC evaluated on additional tasks -- tasks that show some learning progress with CAIAC and in an online learning setting [...]

Following reviewer’s suggestions, we added an additional experiment: Fetch-Pick&Lift with 4 blocks (see Figure 4 right), where there is a mismatch between the empirical joint state distribution in the training data and the joint state distribution at test time. On this task, CAIAC can create useful counterfactual samples and allows generalization beyond the joint state distribution in the training data, thus outperforming other methods by 60%. The rest of the baselines cannot create the needed counterfactuals and have significantly lower performance.

CAIAC is mainly aimed at settings in which the distribution and quantity of training data cannot be controlled, and distribution mismatch is thus a fundamental issue. While applying CAIAC to the online setting remains possible, the continual collection of on-policy data can already alleviate issues such as spurious correlation in the replay buffer or insufficient coverage of the joint state space. While data augmentation can also be helpful in online learning, we believe that centering our experiments on the offline setting can better isolate the issues the method is tackling, and constitutes a more meaningful benchmark.

W6: Empirical evidence on quality of counterfactuals

We appreciate the reviewer’s nice suggestion since we believe it provides additional empirical support to this work. First of all, we want to refer the reader to the Section A.4 of the updated version of the paper, where we show the attention weights of the trained transformer (see Figures 7 and 8) to provide evidence of incorrect causal estimation by CoDA and CoDA action. The weights are even unable to recover influence along the diagonal, i.e., that an entity state at time $t+1$ depends on the entity state at time $t$. These visual results give evidence for claim 1).

Comment 1: I found Figure 6 to be quite helpful in understanding the CAI scores.

We appreciate the comment. We added a reference to the figure in the main paper (Section 5.1.2) as well to an additional analogous figure for the Fetch-Push task (see Figure 7 in the Appendix) which also has an added reference in the main paper (see Section 5.2).

Comment 2: When describing the local causal structure in the chosen benchmark tasks, it may be beneficial to concretely describe the structure. In particular, the agent's actions only affect an object if the agent is in contact with the object.

We would like to clarify that through our work we adopt a unified definition of influence: the agent has influence over an entity in a specific state configuration if it can modify the state of the entity at next time step through its actions (see equation 3 for a formal definition). Concretely, in the case of robotic manipulation environments physical contact is not necessary as long as the agent can change the object pose, even if indirectly, in a single simulation step. We clarified this in the updated version of the manuscript (see Appendix A4).

Comment 3: Related literature

We thank the reviewer for pointing us out to these papers. We will add and discuss the 2 mentioned papers in an extended related work section. After checking these papers, we would like to mention that [2] was not published yet by the time we submitted our manuscript and will be cited as concurrent work. Furthermore [1] is an interesting follow-up of CoDA but assumes access to the true causal graph, and hence does not address the problem we are tackling.

Q1: Additional online RL experiment

Please see answer to weakness W5 above.

Q1:High data regime performance:

This is a good point. We have updated the results on the Fetch-Push experiment (Figure 4, left) and added a new experiment on a Fetch-Pick&Lift environment (Figure 4, right). The updated results stem from a more systematic hyperparameter tuning for each method (see reviewer LtoA (Part 5, Q5) and our answer there). While the results follow a similar trend to the original one, it notably demonstrates significantly higher performance for CAIAC compared to other methods. In order to further analyze the obtained results, we would like to consider three conditions separately. We also added these clarifications in the updated version of the manuscript (see Section 5.2).

1. When data is scarce (see Fetch-Push, low-data regimes), as long as the marginal state distributions at training and test time match for all the entities, CAIAC is able to cover the full support of the state joint distribution by creating unseen configurations of the entities.
2. When data is abundant AND it adequately covers the full joint state space (as in Fetch-Push, high-data regime), no data augmentation is needed by definition and hence CAIAC and the No-Augmented baseline perform similarly. For CoDA and CoDA action we see a light drop in performance possibly due to the creation of unfeasible counterfactuals which harm performance.
3. When data is abundant AND there is yet a mismatch between the joint state distribution at training and test time, CAIAC remains very effective. We showcase this in a new experiment on Fetch-Pick&Lift (see Figure 4 right). We use abundant data, namely 2M samples/40k episodes. However, the state joint distribution in the training data does not match the state joint distribution at test time. CAIAC reaches almost 100% whereas all baselines perform purely. CoDA cannot create helpful counterfactuals, even with an optimized detection threshold.

Q2: Computational demands:

Please see answer to W3 point above.

We thank the reviewer for their valuable feedback. We first address the points under weaknesses and then answer the reviewer’s questions. We address weaknesses with W# and questions with Q#.

W1: Novelty: To the best of our knowledge, this is the first proposed method for data augmentation that uses an explicit measure of influence between entities in the environment to create counterfactuals. While CoDA also suggests creating counterfactuals by swapping local independent subgraphs from different transitions, it either relies on hardcoded dependencies, or uses a heuristic upon the attention weights of a transformer world model to infer the local factorization. As we show, such a heuristic is generally not accurate, and leads to the creation of unfeasible augmented samples that hurt performance. In the updated version of the paper, we provide evidence by showing a comparison between the different measures of influence on the Franka Kitchen and Fetch-Push tasks (Section A.4, Figures 6, 7 and 8), and an analysis of the accuracy of influence prediction for CAI and CoDA in Section A.5 (Figures 9,10 and 11) for the Fetch-Push task.

W2: Performance difference/ gap in analysis: To rephrase, the reviewer is asking to motivate the difference in performance between CoDA and CAIAC. To analyze this, we designed the baseline CoDA-action, which aims at discovering the same part of the causal structure as CAIAC, but using the aforementioned heuristic of transformer attention weights. The fact that CoDA-action performs comparable to CoDA suggests that the assumption of sparsity in the causal graph is not sufficient, but the more accurate measure of causal influence is crucial.

W3: Computational demands: CAIAC relies on computing the CAI measure for data augmentation. In turn, CAI can be evaluated for all entities at once, through $k$ forward passes for the $k$ counterfactual actions, which are performed in a batch-wise fashion. $k$ is a constant factor, and does not scale with the number of entities. Methods relying on a transformer world model, like CoDA and CoDA-action only need one forward pass (which internally has quadratic cost in the number of entities due to cross attention). However, CoDA also needs to compute the connected components from the adjacency matrix, which has a quadratic cost. For relatively few entities, as is common in the robotic manipulation environments, the computational overhead is relatively small. For the high data regime Fetch-Push environment we timed how long it takes for each algorithm to compute influence on all 2M datapoints:

The algorithms were benchmarked on a 12-core Intel i7 CPU. We would like to mention that counterfactuals could be generated in parallel to the learning algorithm and hence not significantly impact runtime of the algorithm. Furthermore, in our offline setting, counterfactuals can be fully precomputed. This information was added in the updated version of the manuscript (see Appendix A6).

W4: Control condition: swapping connected components: To rephrase, the reviewer is suggesting to create counterfactuals by swapping the controllable subgraphs instead of the uncontrollable ones. Both approaches are correct and would lead to feasible counterfactuals. However, the coverage of the state space is greater using the current approach. In the case that given a specific environment configuration only one task is demonstrated (such as in the kitchen), swapping the controllable subgraphs would lead to counterfactuals that are likely already part of the dataset.