Counterfactual Rollouts for Offline Contrastive Goal-Conditioned Reinforcement Learning

Tingrui Huang; Ai Bo; Junzhe Zhang

cs.LGOct 10, 2026

Counterfactual Rollouts for Offline Contrastive Goal-Conditioned Reinforcement Learning

Authors: Tingrui Huang, Ai Bo, Junzhe Zhang

Organizations: Eindhoven University of Technology, Eindhoven, The Netherlands · Department of Electrical Engineering and Computer Science, Syracuse University, Syracuse, NY, USA

Abstract

Contrastive reinforcement learning (CRL) learns goal-conditioned values by associating state--action pairs with the states subsequently reached, assuming that no unobserved variable jointly affects the recorded action and its outcome. When offline demonstrations come from a demonstrator with information hidden from the learner, this assumption fails and the critic absorbs spurious action--outcome associations. We formalize this setting with confounded Markov decision processes and introduce counterfactual transitions that separate what the data identify from what they do not; minimizing over all compatible transitions yields a worst-case occupancy that lower-bounds the true one. Causal Contrastive Reinforcement Learning (CauCRL) learns a pessimistic counterfactual transition from positive and negative offline demonstrations and trains the critic on counterfactual rollouts, substantially improving success rates over standard CRL on confounded goal-reaching tasks.

Figures & tables

Panel of Figure 1 — Figure 1: Overview of CauCRL. A pessimistic counterfactual transition $\mathcal{T}_{\eta}$ is learned from positive demonstrations $\mathcal{D}^{+}$ and negative demonstrations $\mathcal{D}^{-}$ , confounded by $U$ . Its rollout from $(s,a)$ yields a counterfactual future $\tilde{s}^{+}$ that replaces the observed future $s^{+}$ as the positive sample, contrasted against a negative sample $s^{-}\sim P^{-}(s)$ .

Panel of Figure 2 — Figure 2: Windy corridor (left) and success rate over training (right): under confounded demonstrations, all baselines stay below the expert success rate of $1.0$ .

Panel of Figure 3 — Figure 3: Lipschitz continuity of the counterfactual transition. Along the diagonal $a=a^{\prime}$ (thick line), $\mathcal{T}(s,a,a,\cdot)=P(\cdot\mid s,a)$ is identified (filled densities over $s^{\prime}$ ). An off-diagonal entry $\mathcal{T}(s,a,a^{\prime},\cdot)$ is not identified (dashed densities). The shortest transport to the diagonal is L-shaped (arrow): it moves the evaluated action from $a$ and the demonstrator’s action from $a^{\prime}$ to a common value $x$ , with total length $d_{\mathcal{A}}(a,a^{\prime})$ .

Panel of Figure 4 — Figure 4: Continuous-control environments: (a) PointMaze with a stochastic swamp region (yellow) and (b) AntMaze with falling-rock hazards (red).

Figure 5: Success rates over training on Windy Corridor (left), PointMaze (middle), and AntMaze (right). Curves show means over five random seeds, with shaded regions indicating standard errors.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

	Windy Corridor	PointMaze	AntMaze
State dimension	3	8	29
Action dimension	3 (discrete)	2	8
Goal dimension	2	8	2
Episode horizon	400	50	800
Batch size	64	256	1024
Actor hidden layers	$[128]$	$[256,256]$	$[1024,1024]$

Appendix

Table 1: Actor–critic architectures and training hyperparameters.

	Windy Corridor	PointMaze	AntMaze
Conditioning context	state (one-hot, 34-D)	state + goal (16-D)	state (29-D)
Action dimension	4 (discrete)	2	8
Mixture components $K$	– (categorical)	5	5
Hidden layers	$[128,128]$	$[128,128]$	$[512,512]$
Activation	ReLU	SiLU	ReLU
Covariance	–	diagonal	diagonal

Appendix

Table 2: Nominal behavioral-policy architectures and training settings.

	Windy Corridor	PointMaze	AntMaze
Prediction target	next state (75 local offsets)	planar displacement (2-D)	full-state displacement (29-D)
Hidden layers	$[128,128]$	$[128,128]$ + generator	$[512,512]$
Activation	ReLU	ReLU	ReLU
Lipschitz control	locality restriction, $r=2$	target-action spectral projection, $L_{a}\leq 1$	all-layer spectral normalization, coef. $2.0$
Power iterations	–	–	10
Predictive form	categorical (tabulated)	stochastic, $M=8$	deterministic (Dirac)

Appendix

Table 3: Counterfactual transition-model architectures and training settings.

Hyperparameter	Value
Rollout start	route-dependent split point
Rollout horizon	90
Rollouts per refresh	2,000
Refresh interval	every training round
Generated-data fraction	0.5
State feasibility	state bounds and maze free space

Appendix

Table 4: AntMaze counterfactual rollout and replay hyperparameters.

Hyperparameter	Value
Rollout start	recorded $(s,a)$ , uniform over $\mathcal{D}^{+}$
Rollout horizon	400
Rollouts per refresh	10,000
Refresh interval	every training round (5k updates)
Training rounds	20
Generated-data fraction	1.0 (critic); actor BC uses $\mathcal{D}^{+}$

Appendix

Table 5: Windy Corridor counterfactual rollout and replay hyperparameters.

Benchmark	Method	Per-seed
Windy	CauCRL	0.797	0.984	0.919	0.895	1.000
	CRL	0.513	0.513	0.513	0.513	0.513
	GCIQL	0.513	0.513	0.513	0.513	0.513
	GCBC	0.586	0.562	0.635	0.574	0.574
	HIQL	0.513	0.513	0.513	0.513	0.513
PointMaze	CauCRL	1.000	1.000	1.000	1.000	1.000

Appendix

Table 6: Per-seed final success rates (mean over the final

10\%

of evaluation checkpoints) on all three benchmarks.

Explore similar work

Aug 27, 2026

Arrive and Survive: Scaling Safe Goal-Conditioned Policy Learning from One-Bit Failure Signals

Contrastive reinforcement learning (CRL) scales effectively in goal-conditioned tasks by casting policy learning into a self-supervised contrastive objective. However, in a failure-terminated Markov decision process, established CRL considers pre-failure future goals only when constructing positive samples, without accounting for the probability mass removed by failure termination. Our theoretical analysis shows that this omission induces a systematic overestimation bias in goal-reaching values. Consequently, near-failure trajectories provide disproportionately strong supervision of success despite retaining little future occupancy. Unsafe actions can thereby be reinforced through catastrophic failure bootstrapping, leading to failed policy learning and unsustainable goal-reaching behaviours. To address this problem, we introduce two minimal yet strong corrections: mass-weighted InfoNCE corrects the overweighting of short surviving futures in critic learning, and a log-survival-mass score restores the missing survival mass in policy optimization. The resulting method, Safe Contrastive Reinforcement Learning (Safe-CRL), requires only the one-bit signal provided by failure termination to scale safe goal-conditioned policy learning. Across twelve failure-prone robot navigation and locomotion tasks, Safe-CRL consistently improves survival and substantially outperforms the Scaling-CRL baseline in goal-reaching performance. Additionally, deep Safe-CRL policies exhibit complex failure-avoidance behaviours. This study completes the CRL theory under failure termination and provides a scalable safe RL framework. The code is available via https://github.com/RomainLITUD/safe-crl.

Guopeng Li, Yiyang Duan, Yiru Jiao +1

Southeast University Nanjing, China · Lab, Yinwang Intel. Tech. Co. Ltd. Shanghai, China

Goal-Conditioned RL Safe RL Contrastive RL Robotic RL

Ask PDF

Jun 23, 2026

An Introduction to Causal Reinforcement Learning

Causal inference provides a set of principles and tools that allow one to combine data and knowledge about an environment to reason with questions of counterfactual nature, i.e., what would have happened had reality been different, even when no data of this unrealized reality is currently available. Reinforcement learning provides methods to learn a policy that optimizes a specific measure (e.g., reward, regret) when the agent is deployed in an environment and pursues an exploratory, trial-and-error approach. These two disciplines have evolved independently and with virtually no interaction between them. We note that they operate over different aspects of the same building block, counterfactual relations, which makes them umbilically connected. Based on these observations, novel learning opportunities arise when this connection is explicitly acknowledged and mathematized. To realize this potential, we note that any environment where the RL agent is deployed can be decomposed as a collection of autonomous mechanisms with different causal invariances, parsimoniously modeled as a structural causal model; any standard RL setting implicitly encodes such a model. This formalization allows us to put under a unifying treatment different modes of learning, including online, off-policy, and causal calculus learning, which appear unrelated in the literature. However, these modalities are not exhaustive: we introduce several natural and pervasive classes of learning settings that entail novel dimensions of analysis. Specifically, we introduce and discuss through causal lenses generalized policy learning, where to intervene, imitation learning, and counterfactual learning. These tasks lead to a broader view of counterfactual learning and suggest great potential for studying causal inference and reinforcement learning side by side, which we call causal reinforcement learning (CRL).

Elias Bareinboim, Junzhe Zhang, Sanghack Lee

Department of Computer Science Columbia University New York, USA · Graduate School of Data Science Seoul National University Seoul, Republic of Korea

Reinforcement Learning Causal Inference Causal RL Counterfactual Learning

Ask PDF

Jun 8, 2026

Counterfactual Transport Flows for Offline Conservative Trajectory Refinement

Offline reinforcement learning (RL) offers a path to policy improvement from logged data alone, using historical returns or other measurable outcomes as world feedback. A key difficulty is improving observed behavior without extrapolating beyond what the offline data supports. We propose counterfactual transport flows, a source-conditioned trajectory refinement framework for offline decision-making guided by world feedback. Given a low-feedback candidate trajectory, we construct local preference pairs from offline data by retrieving nearby trajectories in latent trajectory space with higher task-specific feedback, and use them as weak supervision for conservative refinement. The framework learns instance-specific refinement directions: at inference time, a refinement strength parameter controls how far the candidate trajectory is transported, enabling a trade-off between preserving the original behavior and applying stronger improvement. Experiments on D4RL benchmarks, including AntMaze and MuJoCo tasks, show that our method improves behavior from historical returns as world feedback, while providing interpretable trajectory-level refinement paths.

Lena Krieger, Xuan Zhao, Zhuo Cao +3

1IAS-8, Forschungszentrum J ulich, Germany · 2LMU Munich, Munich Center for Machine Learning (MCML), Germany · Department of Computer Science, Aarhus University, Denmark.

Trajectory Optimization Offline RL

Ask PDF