Beyond the Known: Decision Making with Counterfactual Reasoning Decision Transformer

TL;DR

CRDT addresses the challenge of suboptimal or limited offline data for Decision Transformers by introducing counterfactual reasoning. It learns separate Treatment and Outcome models to generate plausible counterfactual experiences, filters them by predicted quality and uncertainty, and augments DT training with these samples, enabling improved generalization and stitching without changing the DT architecture. Across continuous and discrete domains, CRDT yields consistent gains and demonstrates stitching capabilities, particularly in data-scarce or modified-dynamics scenarios. This approach leverages the potential outcomes framework to reason about unobserved actions and their consequences, offering a scalable way to extend offline RL with counterfactual data.

Abstract

Decision Transformers (DT) play a crucial role in modern reinforcement learning, leveraging offline datasets to achieve impressive results across various domains. However, DT requires high-quality, comprehensive data to perform optimally. In real-world applications, the lack of training data and the scarcity of optimal behaviours make training on offline datasets challenging, as suboptimal data can hinder performance. To address this, we propose the Counterfactual Reasoning Decision Transformer (CRDT), a novel framework inspired by counterfactual reasoning. CRDT enhances DT ability to reason beyond known data by generating and utilizing counterfactual experiences, enabling improved decision-making in unseen scenarios. Experiments across Atari and D4RL benchmarks, including scenarios with limited data and altered dynamics, demonstrate that CRDT outperforms conventional DT approaches. Additionally, reasoning counterfactually allows the DT agent to obtain stitching abilities, combining suboptimal trajectories, without architectural modifications. These results highlight the potential of counterfactual reasoning to enhance reinforcement learning agents' performance and generalization capabilities.

Figures (8)

• Figure 1: (a): A toy environment where the goal of the agent is to move from the green circle position to the red circle position given that data is biased toward moving from bottom-left to top-right (green trajectory/diagonal line) over top-left to bottom-right (blue trajectory/ diagonal line). When using traditional DT, the agent will most likely follow the green trajectory and fail to reach the goal. (b): The empirical result of the counterfactual reasoning process following CRDT on the toy environment, with the green and blue trajectories forming an intersection. At the intersection, notice that turning right yields a higher potential outcome/return, CRDT generates counterfactual experience accordingly. As shown by the bold yellow, blue, and green dots, none of the counterfactual experiences followed the green trajectory after the crossing point; they all show a clear right turn. Training DT with these counterfactual experiences improved the overall performance (refer to Sect. \ref{['toy_env']} for performance results). (c) Top: The CRDT framework follows three steps: first, learning to reason counterfactually with the CRDT agent; second, perform counterfactual reasoning to generate counterfactual experiences; and third, use these experiences to improve decision-making. Bottom: A single step in the iterative counterfactual reasoning process of a trajectory. The outcomes of one-step reasoning are the counterfactual action $\hat{a}_t$, the next state $\hat{s}_{t+1}$ and returns-to-go $\hat{g}_{t+1}$ will replace the original values $a_t, s_{t+1}, g_{t+1}$ and the generated data will be used in next iteration.
• Figure 2: Performance comparison on limited subset of $D_\text{env}$ (more results are provided in Appendix \ref{['continuous_action_space_full']} Fig. \ref{['fig:smaller_dataset_full']}). The results are over 5 seeds. For each seed, evaluation is conducted over 100 episodes. The X-axis represents the percentage of the dataset used in the experiment.
• Figure 3: Given two trajectories $\left( s^{a}_{t-1}, s_t, s^{a}_{t+1} \right), \left( s^{b}_{t-1}, s_t, s^{b}_{t+1} \right)$. We want our agent to be able to start from state $s^{b}_{t-1}$, however, can reach state $s^{a}_{t+1}$
• Figure F.1: Performance comparison on limited subset of $D_\text{env}$ dataset. We report the results over 5 seeds. For each seed, evaluation is conducted over 100 episodes. The X-axis represents the percentage of the dataset used in the experiment.
• Figure F.2: The impact of varying the number of counterfactual experiences $n_e$ in $D_\text{CRDT}$ on the performance. The agent is trained using the 10% walker2d-medium-replay dataset. The terms CRDT 4000, 8000, and 16000 refer to configurations of CRDT with $n_e$ set to 4000, 8000, and 16000 samples, respectively. We report the results over 5 seeds. For each seed, evaluation is conducted over 100 episodes.