Source-linked AI summary
Counterfactual Rollouts for Offline Contrastive Goal-Conditioned Reinforcement Learning
Tingrui Huang, Ai Bo, Junzhe Zhang
TL;DR
When offline demonstrations are confounded by information hidden from the learner, standard contrastive RL can learn spurious action–outcome associations. CauCRL learns pessimistic counterfactual transitions and trains critics on their rollouts; experiments report improvement across discrete and continuous-control tasks.
Problem
When unobserved confounders jointly influence demonstrated actions and their outcomes, contrastive RL's goal-reaching probabilities may be non-identifiable from the available data.
Source: §2 · p. 3Method
CauCRL learns a pessimistic counterfactual transition from confounded demonstrations and uses its counterfactual rollouts as positive samples for contrastive critic learning.
Source: §1 · p. 2, §3.2 · p. 7Results
Takeaways & Limitations
The worst-case counterfactual occupancy lower-bounds true occupancy, providing a pessimistic proxy for policy improvement in the real environment.
Source: §3.2 · p. 7Learning counterfactual transitions assumes continuity in the evaluated action and the demonstrator's action, corresponding to smooth treatment response and selection.
Source: §3.1 · p. 5
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Offline goal-conditioned RL can be confounded when demonstrators use information unavailable to the learner, undermining the action–outcome associations used by contrastive RL. CauCRL addresses this with pessimistic counterfactual rollouts for critic training.
1 INTRODUCTION
- Goal-conditioned RL targets goals specified at test time, while a learned policy can be reused across multiple tasks.The introduction frames this setting as one response to sparse or absent reward signals.Source: §1 · p. 1
- Unobserved variables jointly influencing recorded actions and outcomes can make goal-reaching probabilities non-identifiable from offline data.Contrastive RL uses future states from the same trajectory as positives, so confounding can invalidate its action–outcome associations.Source: §1 · p. 1
- In the windy corridor, standard contrastive RL overvalues short-route moves because favorable wind selects when the expert moves forward, although unfavorable wind can push the agent into lava.All baselines remain below the expert success rate of 1.0 in Figure 2.Source: §1 · pp. 1–2, Figure 2 · p. 2
- Despite prior causal-RL work across several settings, a systematic approach to goal-conditioned RL under unobserved confounding remains missing.The unresolved question is how to learn the goal-reaching occupancy when confounded demonstrations supply contrastive positives.Source: §1 · p. 2
- CauCRL trains contrastive critics on pessimistic counterfactual rollouts consistent with confounded demonstrations.The paper introduces counterfactual transitions and a worst-case occupancy that lower-bounds the true goal-reaching probability, and evaluates the method on discrete navigation and high-dimensional continuous control.Source: §1 · p. 2
2 GOAL-CONDITIONED RL WITH UNOBSERVED CONFOUNDING
Offline goal-conditioned RL cannot generally recover policy values from demonstrations when an unobserved variable affects both the demonstrator’s actions and state transitions. The section uses counterfactual transitions to represent what the data identify and derives a worst-case occupancy that lower-bounds true goal-reaching probability.
2 GOAL-CONDITIONED RL WITH UNOBSERVED CONFOUNDING
- A confounded goal-conditioned MDP models latent variables as affecting both environment transitions and the demonstrator’s behavioral policy.The learner replaces the demonstrator’s behavior mechanism with its own policy when evaluating interventional outcomes.Source: §2 · pp. 2–3
- Because the demonstrator’s action and next state share an unobserved cause, offline CRL recovers observational occupancy rather than the policy’s Q-function, biasing its critic and improved policy.Positive demonstrations supply observational continuations, while negative demonstrations provide a state distribution used for contrastive negatives.Source: §2 · p. 3
2.1 Counterfactual Modeling
- Counterfactual transitions ask what next state would follow action a, conditional on the demonstrator having actually chosen a′, retaining information about the confounder.When a = a′, consistency identifies the transition from observational data; off-diagonal action pairs remain free parameters among compatible transitions.Source: §2.1 · p. 4
- Minimizing occupancy over compatible counterfactual transitions yields a worst-case measure that lower-bounds the true probability of reaching the goal.The compatible set contains the true transition, while the observational data cannot distinguish it from the other compatible transitions.Source: §2.1 · pp. 4–5
3 CAUSAL CONTRASTIVE REINFORCEMENT LEARNING
CauCRL approximates a worst-case counterfactual construction with learned models, retaining observed dynamics for the factual component and modeling unidentified outcomes pessimistically. It trains contrastive critics on counterfactual rollouts and improves the policy with those critics.
3.1 Learning Counterfactual Transitions
- In continuous domains, identifying transitions only where the evaluated and demonstrator actions match does not constrain nearby actions, motivating additional smoothness constraints.In discrete domains, off-diagonal entries are unrestricted and the worst case follows Manski bounds.Source: §3.1 · p. 5
- Lipschitz continuity restricts unidentified transitions to a Wasserstein ball around identified diagonal transitions, with the allowed distance growing with action separation.The figure depicts the shortest L-shaped path to the identified diagonal, with total length dA(a, a′).Source: §3.1 · p. 5, Figure 3 · p. 5
- The transition model fits identified diagonal entries to positive demonstrations and steers unidentified action pairs toward states from negative demonstrations.When negative demonstrations concentrate on absorbing failures, this pessimistic choice approximates the worst case of the Manski bounds.Source: §3.1 · pp. 5–6
3.2 Counterfactual Rollouts
- Counterfactual rollouts average over the nominal demonstrator policy while following learner-selected actions, resampling the demonstrator action even at the first step.The procedure uses rollout states as positive goals for contrastive learning.Source: §3.2 · p. 6
- The optimal rollout-trained critic recovers the policy Q-function in the induced MDP up to a negative-goal marginal factor, preserving its action ranking.Truncating rollouts at horizon H changes the positive-goal distribution by at most γ^(H+1) in total variation.Source: §3.2 · pp. 6–7
- With a Bayes-optimal critic, policy improvement is exact in the induced MDP; learned pessimistic transitions approximate a lower bound on true occupancy as a proxy for real-environment improvement.With a δ-accurate critic and λBC = 0, the tabular guarantee allows for critic error.Source: §3.2 · p. 7
4 EXPERIMENTS
The experiments evaluate CauCRL on three goal-reaching tasks where demonstrators observe information hidden from learners. CauCRL achieves strong route and success outcomes across the reported environments.
4 EXPERIMENTS
- Across all three tasks, demonstrators observe latent information that affects both their actions and subsequent transitions, but learners and evaluated methods do not.The tasks span discrete Windy Corridor, continuous planar PointMaze, and eight-dimensional continuous-control AntMaze.Source: §4 · p. 7
- Reported outcomes are mean goal-reaching success and, where applicable, robust-route rates, averaged over five seeds with standard errors over final evaluation checkpoints.The final 10% of evaluation checkpoints are used; per-seed results are in Appendix C.5.Source: §4 · p. 7
- CauCRL achieves 0.919 ± 0.036 mean success and 1.00 robust-route rate on Windy Corridor, plus 1.00 success and detour rates on PointMaze.Fig. 5 shows the success-rate curves for these tasks.Source: §4 · p. 8, Figure 5 · p. 8
- 0.815 ± 0.030 success on AntMaze for CauCRL exceeds CRL’s 0.629±0.015.The comparison is reported in Fig. 5 (right).Source: §4 · p. 8, Figure 5 · p. 8
5 CONCLUSION
The conclusion presents CauCRL as a counterfactual approach to offline goal-conditioned RL under unobserved confounding. It reports improved robustness in discrete and continuous-control tasks and identifies broader settings for future evaluation.
5 CONCLUSION
- CauCRL separates identified factual transitions from unidentified counterfactual outcomes and uses a worst-case goal-reaching occupancy that lower-bounds the true occupancy.The paper reports improved robustness over confounding-unaware baselines across discrete and continuous-control tasks.Source: §5 · p. 8
AI Use Statement
- The authors used generative AI tools across research, experimentation, coding, literature review, and manuscript drafting, and state that outputs were reviewed and verified.They retain responsibility for the paper’s final decisions and content.Source: AI Use Statement · p. 9
A PROOFS
The appendix proves three claims underlying CauCRL: a worst-case occupancy lower bound, the optimal critic for counterfactual rollouts, and policy improvement in the induced MDP. The proofs use fixed-point and Bellman-operator arguments.
A PROOFS
- The shared proof strategy characterizes quantities as unique fixed points of γ-contractive Bellman operators, then verifies identities and inequalities at the operator level.The contraction argument gives a unique bounded fixed point.Source: A PROOFS · p. 11
Proof of Proposition 1
- Minimizing counterfactual occupancy over compatible transitions lower-bounds true goal-reaching occupancy because the true transition is compatible and minimization can only decrease occupancy.The proof first establishes that the counterfactual Bellman equation is well posed for every compatible transition.Source: A PROOFS · pp. 11–12
Proof of Proposition 2
- For untruncated CtfRollout, positive goals follow the policy’s goal-reaching occupancy in the induced MDP, yielding the Bayes-optimal critic for the counterfactual NCE objective.The proof identifies that critic by computing the objective’s optimum.Source: A PROOFS · pp. 11, 13
Proof of Proposition 3
- With an approximately optimal critic, greedy policy updates improve the policy in the MDP induced by the counterfactual transition.The proof adapts an approximate policy-improvement argument and propagates the greedy inequality through the Bellman operator.Source: A PROOFS · pp. 11, 13
B IMPLEMENTATION DETAILS
The appendix specifies environment-specific nominal policies, counterfactual transition models, and rollout/replay construction for CauCRL. Generated continuations supply contrastive future goals while recorded demonstrations remain available for behavioral cloning.
B IMPLEMENTATION DETAILS
- CauCRL’s implementation details supplement the algorithmic description with environment-specific parameterizations and refer readers to Appendix C.4 for architectures and optimization hyperparameters.Source: B IMPLEMENTATION DETAILS · p. 14, §B.2 · p. 15
B.1 Nominal Behavioral Policy
- Continuous-control nominal policies use five-component censored Gaussian mixtures, while Windy Corridor uses a categorical policy; PointMaze conditioning excludes privileged hazard and demonstrator metadata.The nominal-policy fitting procedure differs by environment.Source: §B.1 · p. 14
B.2 Counterfactual Transition Model
- Transition models vary by environment: Windy Corridor uses a locally restricted categorical kernel, PointMaze predicts stochastic planar displacement, and AntMaze predicts deterministic full-state displacement.Windy Corridor’s kernel support is restricted to candidate successors near both factual anchors.Source: §B.2 · p. 15
B.3 Counterfactual Rollout and Replay Construction
- Generated continuations replace or extend recorded trajectories, and future-goal relabeling stays within each generated or recorded trajectory.Recorded actions continue to provide behavioral-cloning supervision; AntMaze retains a recorded prefix through the route decision.Source: §B.3 · pp. 15–16
B.4 Contrastive Critic and Actor Training
CauCRL uses the corresponding CRL actor–critic parameterization while training its contrastive critic and goal-conditioned actor with environment-specific choices.
B.4 Contrastive Critic and Actor Training
- CauCRL and the corresponding CRL baseline use the same actor–critic parameterization within each environment.Source: §B.4 · p. 16
- Critic training uses in-batch negatives, with one critic in PointMaze and two independent critics in AntMaze.AntMaze actor actions are evaluated using the minimum of the two critic scores.Source: §B.4 · p. 16
- The actor is a tanh-Gaussian goal-conditioned policy, regularized by behavioral cloning on recorded offline actions.PointMaze uses a separately balanced BC minibatch; AntMaze uses the standard training minibatch.Source: §B.4 · p. 16
C.1 Windy Corridor
Windy Corridor contrasts a short, hazardous route with a longer safe route while hiding wind from the learner. Demonstrators observe wind and predominantly use the short route.
C.1 Windy Corridor
- The 15 × 15 corridor offers a roughly 13-step short route above lava and a roughly 39-step safe route.Entering lava terminates the episode as a failure.Source: §C.1 · p. 16
- Wind is resampled each step and affects forward motion; south wind pushes agents into lava from six positions on the short route.The safe route has no such positions.Source: §C.1 · p. 16
- The learner observes its cell and heading but not the wind, which the demonstrator sees before acting.Source: §C.1 · p. 16
- Demonstrators choose the short route in 95% of episodes and wait at exposed positions under south wind; all 1,000 demonstrations reach the goal.The dataset has 14,028 transitions, including 952 short-route and 48 safe-route episodes, and no lava failures.Source: §C.1 · p. 16
- Evaluation uses 300 seeded episodes with deterministic actions under natural wind, measuring goal success and safe-route use.The safe-route measure is the fraction entering row y = 13.Source: §C.1 · pp. 16–17
C.2 PointMaze
PointMaze compares a hazardous shortcut with a lower detour in a setting where demonstrators observe hazard states that the learner does not. Its offline dataset includes several demonstration types and both successful and failed episodes.
C.2 PointMaze
- PointMaze has a short route crossing three hazard cells and a lower safe detour, with a 50-step episode horizon.Source: §C.2 · p. 17
- Hazard cells activate independently with probability 0.30 each step, and entering an active cell causes absorbing failure.Actions are noisy two-dimensional velocities, and each step uses 10 simulation substeps with per-axis wall collision.Source: §C.2 · p. 17
- The learner receives four recent planar positions as an 8-dimensional state but never observes hazard indicators.The goal uses the corresponding stacked representation.Source: §C.2 · p. 17
- Demonstrators plan the safe route with hazards blocked and use observed hazard indicators to avoid predicted failures on the shortcut.The dataset contains 6,600 episodes and 330,000 transitions, with 4,514 successes, 1,451 absorbed failures, and 635 timeouts.Source: §C.2 · p. 17
- Evaluation uses 300 seeded episodes with deterministic actions and reports goal success and safe-detour rate.An episode counts as a detour if it enters the lower route before any hazard cell.Source: §C.2 · p. 17
C.3 AntMaze
AntMaze contrasts a direct route exposed to stochastic rockfalls with a longer perimeter detour, while hiding rockfall conditions from the learner. Demonstrators observe schedules, and an additional failure pool is used to fit the counterfactual transition model.
C.3 AntMaze
- The ring maze offers a direct route through two rockfall zones and a longer perimeter route that avoids them.Episodes have an 800-step horizon.Source: §C.3 · p. 17
- The learner observes a 29-dimensional Ant state and uses 8-dimensional joint torques, but rockfall states and schedules are hidden.The goal is represented by its planar (x, y) position.Source: §C.3 · p. 17
- Demonstrators choose routes independently of rockfall activations and attempt direct crossings only when predicted intervals avoid rockfall periods.They use a 50-step crossing estimate and an 8-step safety margin.Source: §C.3 · p. 17
- The 1,000-episode demonstration dataset contains equal shares of direct and detour routes, 994 successes, six timeouts, and no rock-contact failures.A separate pool of 300 episodes retains cases where overriding a Zone 1 wait with a go decision causes rock-contact failure; it is excluded from actor–critic replay.Source: §C.3 · pp. 17–18
- Table 1 summarizes actor–critic architectures and training hyperparameters for all three environments.Within each environment, CRL and CauCRL use the same actor and critic architectures.Source: Table 1 · p. 18
- Evaluation uses 300 seeded deterministic-policy episodes at the same rockfall activation probability as data collection, pactive = 0.40.Success is reaching the goal within the 800-step horizon.Source: §C.3 · p. 18
C.4 Network Architectures and Training Hyperparameters
CauCRL and CRL use the same actor and critic architectures within each environment, while nominal-policy and counterfactual-transition configurations vary across environments. The appendix also specifies baseline training and evaluation choices.
C.4 Network Architectures and Training Hyperparameters
- Within each environment, CRL and CauCRL use the same actor and critic architectures.The configurations for all three environments are summarized in Table 1.Source: §C.4 · p. 18
- Windy Corridor uses one-hot state–action and goal features for the critic, and actor minibatches sample distinct recorded state–action pairs uniformly.This sampling choice prevents frequent short-route decisions from outweighing rare safe-route decisions; its contrastive temperature is 0.07 and λ = 0.5.Source: §C.4 · p. 18
- PointMaze and Windy Corridor nominal policies are trained separately and frozen, whereas AntMaze optimizes its nominal-policy head jointly with the counterfactual transition model.PointMaze uses 4,321 training and 479 validation episodes; Windy Corridor uses a random 90% episode split of D+.Source: §C.4 · p. 18
- PointMaze selects the earliest validation minimum for its counterfactual-transition model, while AntMaze and Windy Corridor use the final iterate.The Windy Corridor factual next-state network and counterfactual kernel share settings, with background weight ϵ = 10^-3.Source: §C.4 · p. 18
- HIQL uses three 512-unit hidden layers, with a 10-step PointMaze subgoal horizon and a 25-step AntMaze horizon.The 5-step PointMaze setting produced nearly degenerate subgoal advantages, so the horizon was set to 10.Source: §C.4 · pp. 18–19
- Windy Corridor baselines train for 100,000 updates; HIQL uses a 3-step subgoal horizon, while GCIQL uses expectile 0.9 and an AWR actor.GCIQL and HIQL use binary 1/0 relabeled-goal rewards; failures have zero bootstrap masks, while time-limit truncations continue bootstrapping.Source: §C.4 · p. 19
C.5 Result Details
C.5 provides per-seed final-checkpoint success rates for all three benchmarks, underlying the means reported in Section 4 and plotted in Figures 2 and 5. The related-work section positions CauCRL as complementary to advances in contrastive objectives because its change is to transition and future-outcome supervision.
C.5 Result Details
- Table 6 gives each method’s per-seed success rate averaged over the final 10% of evaluation checkpoints across all three benchmarks.These are the per-seed values behind the Section 4 means and Figures 2 and 5.Source: §C.5 · p. 19, Table 6 · p. 21
D RELATED WORK
- Recent CRL work explores stable architectures, temporal-difference objectives, symmetric InfoNCE with distance-based critics, and quasimetric representations.The cited works include Stable CRL (Zheng et al., 2024a), TD InfoNCE (Zheng et al., 2024b), JaxGCRL (Bortkiewicz et al., 2025), and quasimetric representations (Myers et al., 2025).Source: D RELATED WORK · p. 19
- CauCRL is complementary to these developments because it changes transition and future-outcome supervision rather than requiring a particular contrastive objective.InfoNCE methods can use counterfactually augmented future states as positive pairs, while TD-based methods can use the robust construction for one-step successors.Source: D RELATED WORK · p. 19