Counterfactually-Guided Policy Search (CF-GPS)

• Counterfactually-Guided Policy Search (CF-GPS) is a reinforcement learning framework using structural causal models to perform policy evaluation and improvement in partially observable environments.
• It leverages counterfactual inference on real logged data to generate model-based rollouts, thereby reducing bias and variance compared to traditional synthetic rollouts or importance sampling methods.
• Empirical results on tasks like PO-Sokoban demonstrate that CF-GPS achieves near-zero policy evaluation error when fully conditioned on observations, outperforming conventional model-based planning algorithms.

Counterfactually-Guided Policy Search (CF-GPS) is a framework for reinforcement learning in partially-observable Markov decision processes (POMDPs) that leverages structural causal models (SCMs) to perform policy evaluation and search via counterfactual inference. Key to the method is conditioning model-based rollouts on real, logged data to reduce bias and variance, contrasting with de novo synthetic rollouts or importance sampling approaches. CF-GPS generalizes prior methods such as Guided Policy Search and can be related to value gradient approaches through the lens of counterfactual structural inference (Buesing et al., 2018).

1. Formal Foundations: POMDPs and Structural Causal Models

CF-GPS operates within the episodic POMDP setting. Let $\mathcal{S}$ be the hidden state space, $\mathcal{A}$ the action space, and $\mathcal{O}$ the observation space. The agent trajectory is $\tau = (S_1, O_1, A_1, ..., S_T, O_T)$ with the stochastic policy $\pi(a_t|h_t)$ over the full action-observation history $h_t$. Transition and observation kernels are $P_{S_{t+1}|S_t,A_t}$ and $P_{O_t|S_t, A_{t-1}}$ respectively, and the undiscounted return is $G(\tau) = \sum_{t=1}^T R_t$, where $R_t$ is included in $\mathcal{A}$0.

The entire generative process of the POMDP plus policy $\mathcal{A}$1 is encoded as a structural causal model $\mathcal{A}$2, with scenario (noise) variables $\mathcal{A}$3 encoding randomness in initial state, transitions, observations, and policy actions. Deterministic mechanisms $\mathcal{A}$4 map scenario variables to process variables; for instance, $\mathcal{A}$5 and $\mathcal{A}$6.

Interventions—substituting the policy mechanism $\mathcal{A}$7 by another policy $\mathcal{A}$8—correspond to the standard SCM operation $\mathcal{A}$9, yielding the SCM $\mathcal{O}$0 and corresponding trajectory distribution $\mathcal{O}$1.

2. Counterfactual Evaluation in SCMs

Given observed variables $\mathcal{O}$2, an intervention $\mathcal{O}$3 (e.g., $\mathcal{O}$4), and query variables $\mathcal{O}$5, the counterfactual distribution $\mathcal{O}$6 is obtained by:

1. Inferring scenario variables $\mathcal{O}$7 from $\mathcal{O}$8.
2. Fixing $\mathcal{O}$9 to the sampled value, applying the intervention, and forward-simulating the new mechanism.
3. Returning the queried variables.

Algorithmically, this involves: sampling $\tau = (S_1, O_1, A_1, ..., S_T, O_T)$0, replacing $\tau = (S_1, O_1, A_1, ..., S_T, O_T)$1 with a delta function at $\tau = (S_1, O_1, A_1, ..., S_T, O_T)$2, altering mechanisms per $\tau = (S_1, O_1, A_1, ..., S_T, O_T)$3, and simulating $\tau = (S_1, O_1, A_1, ..., S_T, O_T)$4. Averaging over real data $\tau = (S_1, O_1, A_1, ..., S_T, O_T)$5 ensures unbiased estimation: $\tau = (S_1, O_1, A_1, ..., S_T, O_T)$6.

For policy evaluation (CF-PE), with off-policy dataset $\tau = (S_1, O_1, A_1, ..., S_T, O_T)$7, the counterfactual value estimator for candidate $\tau = (S_1, O_1, A_1, ..., S_T, O_T)$8 is

$\tau = (S_1, O_1, A_1, ..., S_T, O_T)$9

Under a correct model, $\pi(a_t|h_t)$0.

3. The CF-GPS Algorithm and Policy Improvement

CF-GPS alternates between logging real experience and model-based planning, with simulated rollouts "grounded" in inferred scenarios from observed episodes. Each logged episode $\pi(a_t|h_t)$1 is used to derive a posterior over scenario variables $\pi(a_t|h_t)$2. Multiple counterfactual rollouts are generated for each $\pi(a_t|h_t)$3 under a planner or candidate policy $\pi(a_t|h_t)$4.

Policy improvement then proceeds using a return-weighted regression objective, with rollouts from the counterfactual distribution,

$\pi(a_t|h_t)$5

Here, each trajectory $\pi(a_t|h_t)$6 is simulated conditioned on the inferred $\pi(a_t|h_t)$7 ("posterior grounding") instead of starting from the prior, thereby reducing model mismatch.

CF-GPS Algorithmic Outline

$P_{O_t|S_t, A_{t-1}}$4

CF-GPS can incorporate a "planner" $\pi(a_t|h_t)$8 that uses an expert policy or current policy, allowing for annealed mixtures.

4. Methodological Comparisons and Theoretical Insights

CF-GPS situates itself among multiple strands of reinforcement learning:

• Vanilla model-based RL (MB-PE, MB-PS): These approaches draw scenario variables from the model prior $\pi(a_t|h_t)$9, leading to biased results under model mismatch between $h_t$0 and the true environment.
• Importance-Sampling off-policy methods: Reweight logged returns according to trajectory likelihood ratios $h_t$1. While unbiased, these methods suffer from high variance if $h_t$2 and $h_t$3 differ significantly or lack mutual support.

CF-GPS achieves bias reduction by sampling from the posterior $h_t$4, grounding rollouts in real experience. Variance is reduced by generating counterfactual trajectories under fixed scenario noise, randomizing only action noise if needed, and dispensing with trajectory-level importance weighting.

Relations to Other RL Methods

• Guided Policy Search (GPS): GPS, as in Levine & Koltun (2014), is a special case of CF-GPS for fully-observed MDPs with linear-Gaussian SCMs and LQR planners. Here, linear dynamics parameters inferred from real data act as counterfactual scenario inference.
• Stochastic Value Gradient (SVG): SVG interprets policy and model computation as deterministic functions of shared noise $h_t$5, making inference over $h_t$6 and freezing $h_t$7 equivalent to counterfactual trajectory analysis and gradient estimation.

5. Empirical Evaluation and Benchmarking

CF-GPS performance is evaluated on the PO-Sokoban grid-world benchmark:

• Environment: $h_t$8 grid, 3 boxes, 3 targets, strong partial observability (each tile flipped to empty with $h_t$9 probability in observation), episode length $P_{S_{t+1}|S_t,A_t}$0, structured rewards.
• Scenario variables: The initial level layout $P_{S_{t+1}|S_t,A_t}$1, inferred for each episode.

Off-policy Evaluation (CF-PE)

• Model $P_{S_{t+1}|S_t,A_t}$2 is learned only over initial state distribution $P_{S_{t+1}|S_t,A_t}$3; transitions/observations are assumed known.
• Three policies $P_{S_{t+1}|S_t,A_t}$4 are trained under a uniform behavior $P_{S_{t+1}|S_t,A_t}$5.
• Value estimates using MB-PE ($P_{S_{t+1}|S_t,A_t}$6) incur high error ($P_{S_{t+1}|S_t,A_t}$7). As more transitions are conditioned on ($P_{S_{t+1}|S_t,A_t}$8), error drops to near zero. Importance-sampling approaches perform poorly (error $P_{S_{t+1}|S_t,A_t}$9).

Policy Search (CF-GPS)

• Logged data are collected by 64 actors under the behavior policy $P_{O_t|S_t, A_{t-1}}$0.
• Scenario inference $P_{O_t|S_t, A_{t-1}}$1 utilizes DRAW.
• The planner $P_{O_t|S_t, A_{t-1}}$2 is an annealed mixture of the current policy and an expert.
• Each scenario supports 10 counterfactual rollouts, used for return-weighted regression policy updates.

Key Outcomes

• Counterfactual policy evaluation error approaches zero with full smoothing, versus large errors for purely prior-based simulations.
• CF-GPS learning curves surpass both vanilla model-based planning and GPS-like baselines that do not fully condition on observations.
• Given the same real data budget, CF-GPS outperforms model-free IMPALA.

6. Context, Limitations, and Generalizations

CF-GPS formalizes model-based policy learning as counterfactual inference with SCMs, centering the role of posterior scenario inference in reducing bias otherwise endemic to model-simulator approaches. The methodology assumes correct or well-calibrated models for scenario variable inference; model errors can propagate into policy search steps. The generalization to fully-observed MDPs and incorporation of planners or experts in $P_{O_t|S_t, A_{t-1}}$3 broadens applicability, subsuming Guided Policy Search and linking to Stochastic Value Gradient methods ("reparameterization-based algorithms can be interpreted as counterfactual methods"). This suggests potential for further theoretical unification and practical extensions in RL.

• Buesing, L., Weber, T., Zwols, Y., Racanière, S., Guez, A., Lespiau, J.-B., & Heess, N. (2019). "Woulda, Coulda, Shoulda: Counterfactually-Guided Policy Search" (Buesing et al., 2018).
• Pearl, J. (2009). Causality. Cambridge University Press.
• Levine, S., & Koltun, V. (2014). "Guided policy search." ICML.
• Heess, N., et al. (2015). "Learning continuous control policies by stochastic value gradients." NIPS.