Decision Transformer with Counterfactuals

• Decision Transformer with Counterfactuals extends standard DT by synthesizing plausible what-if trajectories to overcome the limitations of behavior cloning on offline data.
• It employs auxiliary treatment and outcome models to generate, simulate, and filter counterfactual rollouts, thereby enhancing policy stitching and transfer to new environments.
• Empirical studies show that this approach leads to notable performance gains in data-limited settings and altered dynamics, validating its robustness and practical viability.

Decision Transformer with counterfactuals refers to a class of reinforcement learning (RL) methods that augment sequence-modeling policy learners, such as the Decision Transformer (DT), with explicit counterfactual reasoning capabilities. The principal motivation is to overcome DT's intrinsic limitation of only imitating observed behaviors from offline datasets, by synthetically generating and leveraging plausible “what-if” trajectories derived from learned environment models or simulator interventions. This facilitates improved off-distribution generalization, robust policy stitching, and enhanced transfer to environments with altered or unseen dynamics. Key instantiations of this approach include the Counterfactual Reasoning Decision Transformer (CRDT) (Nguyen et al., 14 May 2025) and transfer frameworks featuring causal interventions in structured simulators (Boustati et al., 2021). The domain also intersects with causal sequence modeling for counterfactual inference in non-RL settings, as in G-Transformer (Xiong et al., 2024).

1. Foundations: Standard Decision Transformer

The Decision Transformer (DT), originally formulated by Chen et al. (2021), models the offline RL problem as trajectory sequence modeling. Given an environment dataset

$\D_{\mathrm{env}} = \Bigl\{\tau^{(i)}\Bigr\}_{i=1}^N,\quad \tau^{(i)} = (g^{(i)}_1, s^{(i)}_1, a^{(i)}_1, \dots, g^{(i)}_T, s^{(i)}_T, a^{(i)}_T),$

with $g_t = \sum_{t'=t}^T r_{t'}$ the return-to-go, DT feeds the tuple $(h_t, s_{t+1}, g_{t+1})$ through a Transformer $\M_\delta$ to learn the action distribution: $p(a_{t+1} \mid h_t, s_{t+1}, g_{t+1};\delta) = \M_\delta(h_t, s_{t+1}, g_{t+1}).$ The model is trained via standard action classification/regression losses. However, DT's behavior cloning paradigm restricts it to trajectories actually present in $\D_{\mathrm{env}}$. This severely impedes generalization in the presence of limited, sparse, or mostly suboptimal data, and precludes synthesizing high-return “stitched” trajectories from sub-trajectories—a phenomenon long observed in value-based RL but fundamentally missing in imitation-driven approaches (Nguyen et al., 14 May 2025).

2. Counterfactual Reasoning Architectures

Augmenting DT with counterfactual capabilities involves explicitly modeling the environment’s transition and reward mechanisms to propose, simulate, and filter hypothetical (“counterfactual”) trajectories that extend beyond the dataset support. The Counterfactual Reasoning Decision Transformer (CRDT) (Nguyen et al., 14 May 2025) exemplifies this approach via a tripartite pipeline:

1. Auxiliary Modeling: Train (a) a “Treatment” Transformer $\T_\theta$ to estimate action-selection probabilities, and (b) an “Outcome” Transformer $\O_\eta$ to predict subsequent states and returns.
2. Counterfactual Trajectory Generation: Roll out synthetic trajectories by sampling low-likelihood (“counterfactual”) actions from $\T_\theta$, simulate their effects on future states and rewards via $\O_\eta$, and filter the rollouts for plausibility using predicted returns and epistemic uncertainty estimates.
3. Training Augmentation: Aggregate successful counterfactual trajectories into a buffer $g_t = \sum_{t'=t}^T r_{t'}$0, and jointly retrain the base DT policy on both real and counterfactual data.

Algorithmic workflows implement these steps via interleaved stochastic sampling, filtering based on predicted improvement in returns, and rigorous uncertainty thresholds to avoid out-of-distribution dynamics. Notably, CRDT does not alter the underlying Transformer architecture of DT, instead modifying only the data distribution on which it is trained.

Complementary approaches generate counterfactual data via structural causal interventions in simulators (Boustati et al., 2021): changes in environment parameters (e.g., obstacle layouts) are effected with the do-operator, and factual and counterfactual rollouts are synthesized for policy distillation via DT.

In settings outside RL, such as longitudinal treatment-outcome modeling, the G-Transformer pairs Transformer-based sequence modeling with Monte Carlo simulation under a dynamic regime identifier $g_t = \sum_{t'=t}^T r_{t'}$1, thus enabling direct counterfactual outcome prediction in complex settings (Xiong et al., 2024).

3. Formalization and Algorithms

The formal methodology in CRDT involves:

• Treatment Model Training: For discrete actions, optimize the cross-entropy

$g_t = \sum_{t'=t}^T r_{t'}$2

For continuous actions, minimize negative log-likelihood assuming a Gaussian parameterization.

• Outcome Model Training: Predict next states and returns via mean squared error:

$g_t = \sum_{t'=t}^T r_{t'}$3

• Counterfactual Generation: At each step, select $g_t = \sum_{t'=t}^T r_{t'}$4 candidate actions sampled from the low-probability tail, accept only those predicted to improve returns or whose uncertainty is within a specified bound, and roll forward to synthesize a trajectory.
• Policy Retraining: Re-optimize the DT parameters $g_t = \sum_{t'=t}^T r_{t'}$5 with mini-batches drawn equally from $g_t = \sum_{t'=t}^T r_{t'}$6 and $g_t = \sum_{t'=t}^T r_{t'}$7:

$g_t = \sum_{t'=t}^T r_{t'}$8

• Stitching: By composing hypothetical transitions at trajectory junctions (e.g., taking high-return branches in navigation tasks), CRDT synthesizes new, higher-return trajectories absent from the dataset, thereby enhancing generalization and performance.

In the causal transfer framework (Boustati et al., 2021), the generation of counterfactual rollouts proceeds via simulation under interventions $g_t = \sum_{t'=t}^T r_{t'}$9, with trajectories weighted by the magnitude of their average treatment effect (ATE) on total return: $(h_t, s_{t+1}, g_{t+1})$0 These weights prioritize counterfactuals that meaningfully alter policy effectiveness.

The following table summarizes the three-phase CRDT pipeline (Nguyen et al., 14 May 2025):

Phase	Module / Action	Objective / Criterion
Auxiliary Model Training	$(h_t, s_{t+1}, g_{t+1})$1, $(h_t, s_{t+1}, g_{t+1})$2	Action-prob. estimation, state-return prediction
Counterfactual Generation	Counterfactual rollouts, filtering	Return improvement, uncertainty bounds
DT Policy Retraining	vanilla $(h_t, s_{t+1}, g_{t+1})$3, data mixing	Cross-entropy/MSE over real + counterfactual

4. Empirical Results and Ablations

Comprehensive empirical studies substantiate the gains of counterfactual-augmented DT methods in settings with limited data, altered dynamics, and suboptimal rollouts:

• D4RL Locomotion & Ant: CRDT achieves mean returns $(h_t, s_{t+1}, g_{t+1})$4 on locomotion (DT: $(h_t, s_{t+1}, g_{t+1})$5, EDT: $(h_t, s_{t+1}, g_{t+1})$6, REINF: $(h_t, s_{t+1}, g_{t+1})$7), reflecting a $(h_t, s_{t+1}, g_{t+1})$8 improvement versus vanilla DT (Nguyen et al., 14 May 2025). On Ant, performance improves by $(h_t, s_{t+1}, g_{t+1})$9.
• Data-limited Regimes: When training on $\M_\delta$0 of D4RL data, CRDT’s performance degrades less ($\M_\delta$1) compared to DT/REINF ($\M_\delta$2); on Maze2d umaze, DT fails ($\M_\delta$3 drop), whereas CRDT drops only $\M_\delta$4.
• Atari (1% DQN-replay): Human-normalized scores show CRDT (91.5) outperforming DT (87.9), with notable gains in Breakout ($\M_\delta$5), Qbert, and Seaquest.
• Ablations: Removing return comparison, uncertainty filtering, or using random/noisy outliers reduces the benefits, with the full CRDT pipeline consistently yielding the highest scores.
• Stitching Capabilities: In a “stitching toy” environment, DT’s “stitch rate” is $\M_\delta$6 versus CRDT’s $\M_\delta$7 (Nguyen et al., 14 May 2025).

For counterfactual-based transfer, DT trained with both factual and counterfactual trajectories (with ATE weighting) achieves near-doubling of success rates (<30% → >70%) and substantial reward increases in zero-shot evaluation on unseen environments (Boustati et al., 2021).

5. Theoretical and Practical Considerations

CRDT and related algorithms are grounded in the Rubin–Neyman potential-outcomes formalism, leveraging assumptions of consistency, overlap, and ignorability for valid counterfactual identification from offline data. Counterfactual outcomes are only as trustworthy as the underlying treatment and outcome models $\M_\delta$8 and $\M_\delta$9, with epistemic uncertainty explicitly managed to prevent distributional shift.

Practical challenges include:

• The need to train two additional large Transformer models,
• Sensitivity to hyperparameters ($p(a_{t+1} \mid h_t, s_{t+1}, g_{t+1};\delta) = \M_\delta(h_t, s_{t+1}, g_{t+1}).$0, $p(a_{t+1} \mid h_t, s_{t+1}, g_{t+1};\delta) = \M_\delta(h_t, s_{t+1}, g_{t+1}).$1, $p(a_{t+1} \mid h_t, s_{t+1}, g_{t+1};\delta) = \M_\delta(h_t, s_{t+1}, g_{t+1}).$2, $p(a_{t+1} \mid h_t, s_{t+1}, g_{t+1};\delta) = \M_\delta(h_t, s_{t+1}, g_{t+1}).$3),
• Computational cost of synthesizing sufficient counterfactual rollouts,
• Absence of finite-sample statistical error bounds for the synthesized trajectories (Nguyen et al., 14 May 2025).

A key caveat is the reliance on access to accurate simulators or generative models for outcome prediction; in transfer settings, the method assumes an available structural causal model (SCM) or high-fidelity simulator (Boustati et al., 2021). In real-world settings where only partial models are available, generative fidelity and sample quality become primary bottlenecks.

6. Comparative Context and Extensions

The methodology of decision transformer with counterfactuals is situated at the intersection of offline RL, imitation learning, and causal inference. While standard DTs operate strictly within the support of empirical data, counterfactual integration enables navigation of unobserved or altered environments and the “stitching” together of action segments—a trait previously unique to value-based RL with model-based planning.

Causal sequence modeling in non-RL applications connects closely: G-Transformer, for instance, enables counterfactual outcome prediction under time-varying dynamic treatment regimes by forward-simulating sequences under fixed regimes $p(a_{t+1} \mid h_t, s_{t+1}, g_{t+1};\delta) = \M_\delta(h_t, s_{t+1}, g_{t+1}).$4. Unlike DT, which predicts actions given desired returns, G-Transformer simulates state transitions and outcomes under hypothetical interventions, using the g-computation formula (Xiong et al., 2024).

A plausible implication is that further fusion of causal simulation and sequence modeling—such as integrating regime identifiers and causal simulation heads directly into the DT architecture—could offer richer, more robust counterfactual reasoning capabilities, especially for policy optimization and policy evaluation in complex settings.

7. Limitations and Open Problems

Aside from the aforementioned computational and modeling overhead, key open questions include:

• Scalability of counterfactual synthesis for high-dimensional and long-horizon environments,
• Formal statistical guarantees and robustness of counterfactual outcomes under misspecification,
• Extensions to settings with only partial or approximate structural models,
• End-to-end joint training of treatment, outcome, and policy models to minimize distributional mismatch or representation collapse (Nguyen et al., 14 May 2025, Boustati et al., 2021).

The limitations highlight critical frontiers in both RL and causal inference for learning robust, generalizable policies from offline data by reasoning far “beyond the known.”