Counterfactual Reasoning Decision Transformer

• The paper introduces CRDT, which expands Decision Transformer training data using counterfactual trajectories to overcome limitations of suboptimal and sparse data.
• It employs auxiliary transformer models for treatment and outcome estimation to generate alternative actions and predict resulting states, enhancing trajectory stitching and generalization.
• Empirical results demonstrate that CRDT significantly outperforms standard DT on benchmarks like D4RL and Atari by achieving up to a 16.1% performance gain in challenging environments.

Counterfactual Reasoning Decision Transformer (CRDT) is a framework for offline reinforcement learning that augments a standard Decision Transformer (DT) with counterfactual reasoning so that the agent can generate and utilize counterfactual experiences, reason beyond observed data, and improve decision-making in unseen scenarios. It is explicitly motivated by the fact that DT requires high-quality, comprehensive data to perform optimally, whereas many real-world offline datasets are limited, suboptimal, or contain scarce optimal behaviours; under those conditions, CRDT is designed to expand behavioural diversity and improve performance without architectural modifications to the base DT (Nguyen et al., 14 May 2025).

1. Problem setting and motivation

CRDT addresses a failure mode of sequence-modeling approaches to offline RL that arises when the training corpus does not adequately cover high-return behaviour. In the formulation used for CRDT, DT treats trajectory modeling as a supervised sequence prediction task, but this dependence on logged data quality produces three specific difficulties: overfitting to suboptimal data, weak trajectory stitching, and limited generalization to unseen or altered environment conditions (Nguyen et al., 14 May 2025).

The overfitting issue is described in explicitly causal terms. If better trajectories are underrepresented, DT will mostly imitate, by likelihood maximization, the poorer ones. The stitching issue is also central: DTs are weak at “stitching” together optimal segments from suboptimal trajectories since they only reproduce seen trajectory fragments. A further consequence is limited generalization, because the policy remains anchored to historically observed behaviors and therefore struggles in out-of-distribution settings or under altered dynamics (Nguyen et al., 14 May 2025).

CRDT is proposed as a remedy that uses hypothetical “what-if” alternatives for actions and outcomes. Rather than modifying the DT backbone itself, it expands the training signal by constructing plausible counterfactual experiences and then reusing them during DT training. This makes CRDT a data-augmentation and regularization framework built around counterfactual reasoning rather than a new sequence-modeling architecture in the narrow sense (Nguyen et al., 14 May 2025).

2. Formal construction of the framework

CRDT retains the standard Decision Transformer as the acting model and adds two transformer-based auxiliary models: a treatment model $\mathcal{T}$ and an outcome model $\mathcal{O}$ (Nguyen et al., 14 May 2025). The treatment model estimates the conditional probability distribution of actions given trajectory history, with “treatment assignment” taken in the causal-inference sense, while the outcome model predicts the resulting state and return-to-go if a given action is taken.

For the treatment model, the discrete-action formulation is

$$p(a_t \mid h_{t-1}, s_t, g_t; \theta) = \mathcal{T}(h_{t-1}, s_t, g_t; \theta),$$

where $h_{t-1}$ is past trajectory, $s_t$ is current state, and $g_t$ is current return-to-go. For continuous action spaces, the model assumes

$$a_t \sim \mathcal{N}(\mu_t, \sigma_t^2), \qquad (\mu_t, \sigma_t^2) = \mathcal{T}(h_{t-1}, s_t, g_t; \theta).$$

The corresponding losses are

$$\mathcal{L}_{\mathcal{T}(\theta)} = -\frac{1}{N} \sum_{i=1}^{N} a_{t}^{*(i)} \log p(a_{t}^{(i)} \mid h_{t-1}^{(i)}, s_t^{(i)}, g_t^{(i)}; \theta)$$

for discrete action spaces, and

$$\mathcal{L}_{\mathcal{T}(\theta)} = \frac{1}{N} \sum_{i=1}^{N} \left( \frac{(a_t^{*(i)} - \mu_t^{(i)})^2}{2(\sigma_t^{(i)})^2} + \frac{1}{2}\log\!\big(2\pi(\sigma_t^{(i)})^2\big) \right)$$

for continuous action spaces (Nguyen et al., 14 May 2025).

The outcome model operates on

$h_t = (g_1, s_1, a_1, \dots, g_t, s_t, a_t)$

and predicts

$\mathcal{O}$0

with mean-squared-error objective

$\mathcal{O}$1

These auxiliary models supply the counterfactual mechanism, while the final policy remains a standard DT (Nguyen et al., 14 May 2025).

3. Counterfactual trajectory generation and filtering

The core operational step in CRDT is the generation of counterfactual trajectories. At every step in a trajectory, $\mathcal{O}$2 is used to estimate the action distribution given the current history, and candidate actions are selected specifically from low-probability regions so that the model explores alternatives that were not preferred by the behavioural policy in the data (Nguyen et al., 14 May 2025).

For discrete actions, CRDT samples $\mathcal{O}$3 actions classified as counterfactual when their selection probability under $\mathcal{O}$4 is below a threshold $\mathcal{O}$5. For continuous actions, the construction uses the bound on the max of Gaussian random variables: $\mathcal{O}$6 and defines candidate actions by

$\mathcal{O}$7

where $\mathcal{O}$8, $\mathcal{O}$9 is the standard normal inverse CDF, and $$p(a_t \mid h_{t-1}, s_t, g_t; \theta) = \mathcal{T}(h_{t-1}, s_t, g_t; \theta),$$0 is a step size (Nguyen et al., 14 May 2025).

Candidate counterfactual trajectories are then subject to two filters. The first is a high return filter: only trajectories where the counterfactual return-to-go prediction $$p(a_t \mid h_{t-1}, s_t, g_t; \theta) = \mathcal{T}(h_{t-1}, s_t, g_t; \theta),$$1 is lower, meaning higher return, than the original $$p(a_t \mid h_{t-1}, s_t, g_t; \theta) = \mathcal{T}(h_{t-1}, s_t, g_t; \theta),$$2 are retained. The second is an uncertainty filter based on a dropout-derived uncertainty estimate from $$p(a_t \mid h_{t-1}, s_t, g_t; \theta) = \mathcal{T}(h_{t-1}, s_t, g_t; \theta),$$3. A trajectory is discarded if accumulated prediction variance exceeds a threshold $$p(a_t \mid h_{t-1}, s_t, g_t; \theta) = \mathcal{T}(h_{t-1}, s_t, g_t; \theta),$$4: $$p(a_t \mid h_{t-1}, s_t, g_t; \theta) = \mathcal{T}(h_{t-1}, s_t, g_t; \theta),$$5 Here the variance is taken as the maximum over state dimensions. The process continues iteratively, substituting counterfactual actions and propagating predicted outcomes as new inputs, until the end of the trajectory or until filtering blocks the path. Successful trajectories are stored in $$p(a_t \mid h_{t-1}, s_t, g_t; \theta) = \mathcal{T}(h_{t-1}, s_t, g_t; \theta),$$6 (Nguyen et al., 14 May 2025).

The final training stage mixes original and counterfactual experience: $$p(a_t \mid h_{t-1}, s_t, g_t; \theta) = \mathcal{T}(h_{t-1}, s_t, g_t; \theta),$$7 This realizes the central design principle of CRDT: improved decision-making is sought not by changing the DT architecture, but by strictly expanding the trajectory dataset with plausible, high-return, low-uncertainty counterfactuals (Nguyen et al., 14 May 2025).

4. Empirical behavior, benchmarks, and observed effects

CRDT is evaluated on D4RL benchmarks, including Locomotion, Ant, and Maze2d; on Atari games Breakout, Qbert, Pong, and Seaquest; and on a toy environment used to analyze stitching and data scarcity. The comparisons include Decision Transformer, Reinformer, Elastic Decision Transformer, Behavior Cloning, CQL, IQL, MOReL, and MOPO, under full-data training, a small-data regime using only $$p(a_t \mid h_{t-1}, s_t, g_t; \theta) = \mathcal{T}(h_{t-1}, s_t, g_t; \theta),$$8 of the dataset, modified-environment generalization, and stitching-focused toy experiments (Nguyen et al., 14 May 2025).

On continuous-control D4RL tasks, CRDT outperforms DT and the other sequential modeling approaches in both overall reward and in environments with few or underrepresented optimal trajectories. The reported gains include a $$p(a_t \mid h_{t-1}, s_t, g_t; \theta) = \mathcal{T}(h_{t-1}, s_t, g_t; \theta),$$9 average gain over DT on Locomotion, a $h_{t-1}$0 gain over DT on Ant, and a largest single improvement of $h_{t-1}$1 in walker2d-med-rep. On Atari, CRDT improves normalized scores on 3 out of 4 games, with a reported $h_{t-1}$2 on Breakout (Nguyen et al., 14 May 2025).

The small-data experiments are used to highlight robustness to limited offline data. With only $h_{t-1}$3 of the dataset, CRDT shows a $h_{t-1}$4 degradation, compared with $h_{t-1}$5 for DT and REINF. Under altered reward or dynamics at test time, CRDT generalizes better and outperforms DT and REINF in all but one of four tested modified scenarios (Nguyen et al., 14 May 2025).

The most distinctive qualitative effect is stitching. In the toy experiments with highly imbalanced data favoring bad paths, DT achieves only $h_{t-1}$6 success, whereas CRDT achieves $h_{t-1}$7. In Maze2d and Ant, which are described as requiring stitching for success, CRDT attains much higher returns and success rates than DT and REINF. The paper attributes this to hypothetical counterfactual reasoning that allows the model to imagine, and therefore practice, transitions unseen in any single real trajectory (Nguyen et al., 14 May 2025).

Ablation results further constrain the interpretation. Simple data augmentation by adding noise or random alternate actions helps less than proper counterfactual reasoning as formalized in CRDT. Both the high-return filter and the uncertainty filter contribute, ablating them marginally reduces benefits, and broader action search together with more counterfactual samples improves performance until saturation. Visualizations also show broader action distributions in state-action space where training data is sparse (Nguyen et al., 14 May 2025).

5. Causal assumptions and conceptual interpretation

CRDT is explicitly rooted in the potential outcomes theory of causal inference and assumes consistency, sequential ignorability, and sequential overlap (Nguyen et al., 14 May 2025). In that sense, the treatment model and outcome model are not merely auxiliary predictors: they provide the machinery for estimating potential outcomes under alternative actions and for generating action-conditioned hypothetical rollouts from an observational offline dataset.

At the same time, CRDT does not rely on explicit structural causal models or graph learning. Instead, it learns the necessary counterfactual mappings directly using transformer sequence models. This choice differentiates it from approaches that represent interventions via explicit structural equations, while preserving a causal interpretation of generated alternatives through the treatment-assignment and potential-outcome vocabulary (Nguyen et al., 14 May 2025).

The filtering stage also clarifies the intended semantics of the counterfactuals. The high-return filter is aligned with DT’s tendency to improve with higher return data, whereas the uncertainty filter is meant to avoid unfamiliar or out-of-distribution states. A plausible implication is that CRDT treats counterfactual reasoning as constrained imagination: the generated trajectories are not arbitrary perturbations, but alternatives that are simultaneously low-probability under the behaviour policy, predicted to be beneficial, and screened for familiarity (Nguyen et al., 14 May 2025).

A common misconception is to treat CRDT as a replacement for the DT architecture itself. The framework is instead described as operating “without architectural modifications” and as a wrapper data-augmentation framework around the backbone DT. The substantive intervention occurs in the data regime and training signal rather than in the sequence-modeling core (Nguyen et al., 14 May 2025).

6. Relation to adjacent counterfactual transformer methods

CRDT belongs to a broader literature on counterfactual reasoning for sequential decision problems, but it is not interchangeable with adjacent transformer-based causal models. A related offline-RL study trained a Decision Transformer on factual and counterfactual simulations generated by structural interventions $h_{t-1}$8 on the environment, using trajectories collected by running an unchanged source policy in both source and counterfactual environments; it also considered weighting trajectories according to Average Treatment Effect (ATE) to improve transfer to new environments (Boustati et al., 2021). That work shares the use of counterfactual trajectories for DT training, but its mechanism is framed around transfer under changes in environment dynamics rather than the treatment-model/outcome-model construction used in CRDT.

A separate line of work addresses explicit counterfactual outcome estimation rather than offline-RL policy learning. The “Causal Transformer” estimates counterfactual outcomes over time from observational longitudinal data by combining three transformer subnetworks for time-varying covariates, previous treatments, and previous outcomes, together with a counterfactual domain confusion loss that learns adversarial balanced representations predictive of the next outcome but non-predictive of the current treatment assignment (Melnychuk et al., 2022). In that formulation, the target is counterfactual outcome estimation under the Neyman-Rubin causal model with time-varying confounding, not return-conditioned policy optimization in an MDP.

The same distinction appears in spatial-temporal causal inference. A transformer-based framework for spatial-temporal counterfactual outcomes estimation models the intensity function of an outcome point process, uses CNN-based propensity estimation for high-dimensional treatments, and gives an IPW-based estimator that is consistent and asymptotically normal under stated assumptions (Li et al., 26 Jun 2025). Likewise, the Counterfactual Gaussian Process models continuous-time outcome trajectories under sequences of future actions by jointly modeling event times, action choice, and outcomes within a marked point process, with an adjusted likelihood designed to undo confounding from the action policy in observational traces (Schulam et al., 2017).

Taken together, these works indicate that “counterfactual reasoning” in transformer-adjacent sequence modeling spans at least two distinct regimes. One regime, exemplified by CRDT and related DT transfer methods, uses counterfactual trajectories to improve policy learning from offline data (Nguyen et al., 14 May 2025, Boustati et al., 2021). The other regime, exemplified by the Causal Transformer, spatial-temporal transformer estimators, and CGP, focuses on explicit estimation of counterfactual outcomes under hypothetical interventions in observational longitudinal data (Melnychuk et al., 2022, Li et al., 26 Jun 2025, Schulam et al., 2017). This suggests that CRDT is best understood as an offline-RL framework for counterfactual data generation and reuse, rather than as a direct estimator of longitudinal causal effects.