Complementing a Policy with a Different Observation Space

Gokul Swamy 1 Sanjiban Choudhury 2 J. Andrew Bagnell 3 1 Zhiwei Steven Wu 1

Abstract

Williams, 1992), and when searching over a set of deterministic policies (Bagnell et al., 2003).

We consider the problem of improving upon ablack-box policy which operates on a different
observation space than the learner. Such problems occur when augmenting an existing handengineered system with a new machine learningmodel or in a shared autonomy / human-AI complementarity context. We prove that following
the naive policy gradient can lead to a decreasein performance because of incorrect grounding ina different observation space. Then, if we haveaccess to both sets of observation at train time, we
derive a method for correctly estimating a policy
gradient via an application of backdoor adjustment. If we don’t, we prove that under certain
assumptions, we can use the proxy correction tocorrectly estimate a direction of improvement.

An oft-underemphasized requirement for the above policy
improvement theorems to hold is that the proposed change
to the policy uses the same observation space as the policy
it is improving over. In practice, this often doesn’t hold,as ML practitioners may be attempting to improve a policythat comes from an existing, effectively black-box system
to them, or even to provide action recommendations toexpert humans who operate with a fundamentally different
observation space.
For example, consider a medical assistance system to help
doctors at a walk-in clinic. In response to observing inflamed tonsils, a doctor might prescribe an antibiotic. The
automated system, which operates purely based on numerical information, does not have access to this feature
but does observe that patients always get better after being prescribed amoxicillin. It would then over-estimate
Qdoctor (o, amoxicillin), leading to it suggesting uneccesaryprescriptions of a strong medication that could have unfortunate side-effects for patients with issues other than strep
throat. More generally, naively computing Qπ (o, a) where
π uses a different observation space than o can lead to estimates of cumulative reward that will not match what thelearner would actually receive at test time.

1. Introduction
Classical results tell us that in a Partially Observed Markov
Decision Process (POMDP), finding high-value policies
that map observations to actions is computationally hard
(Papadimitriou & Tsitsiklis, 1987). Furthermore, a naive
application of MDP techniques (e.g. value-iteration, policyiteration, or Q-learning) empirically leads to unstable and often decaying performance over iterations (Singh et al., 1994;
Littman, 1994). However, folk wisdom suggests that if one
computes (or approximates well) Qπ (o, a) – the long-term
expected return of an observation-action pair – and thenmakes “small adjustments” to the policy to select higher
value actions for a given observation, the policy can be reliably improved. Such small changes can then be repeatedto achieve a notion of local optima. This folk wisdom finds
precision for stochastic mixtures of policies (Jaakkola et al.,
1994; Kakade & Langford, 2002; Daumé et al., 2009), parameterized stochastic policies (Baxter & Bartlett, 2001;

In this work, we study the perhaps initially surprising failure of these policy improvement algorithms. We identify
their fundamental issue: when observations spaces differ, Qπ (o, a) becomes causally confounded, rendering thestandard Monte-Carlo estimate inconsistent. We then study
two regimes- - first where a richer observation space is available when we train a policy, but not when we execute thatpolicy ; second, a more difficult setting where at both train
and test we have access only to a history of an observation
space more limited than the initial policy we wish to improve over. Our technical tools leverage recent advancesin proxy variable correction techniques (Miao et al., 2018;
Tchetgen et al., 2020) and approximate back-door corrections (Pearl, 2009; Xu & Gretton, 2022) to provide correct
policy gradients in confounded settings.

*

Equal contribution 1 Carnegie Mellon University 2 Cornell University 3 Aurora Innovation. Correspondence to: Gokul Swamy<gswamy@cmu.edu>.

More explicitly, our work makes the following three contributions.

Proceedings of the Interactive Learning with Implicit Human Feedback Workshop at ICML 2023., Honolulu, Hawaii, USA. PMLR
202, 2023. Copyright 2023 by the author(s).

1Complementing a Policy with a Different Observation Space

for small enough ϵ.

1. We formalize the problem of improving upon a policy
with a different observation space. We prove that correct
estimation of Qπ (o, a) values is sufficient for policy improvement. We then give a family of examples under which
the standard, Monte-Carlo estimate for Qπ (o, a) will notlead to policy improvement.

2.1. Policy Improvement in POMDPs
The first complexity we have to deal with is that our policies
depend on observations rather than states. Following anargument similar to that of Jaakkola et al. (1994), we prove
the following theorem.

2. We consider the setting when both sets of observations
are available at training time. We show how we can use
Pearl’s backdoor adjustment (Pearl, 2009) to de-confoundour Q-value estimates. We derive a procedure efficiently
approximating the backdoor adjustment using the output ofa simple classifier.

Theorem 2.1. Define πϵ = (1 − ϵ)πb + (ϵ)π, where πboperates on O2 and π operates on O1 . Then, we have that
J(πϵ ) − J(πb ) =ϵT Eo∼ρπϵ [Ea∼π(a|o) [Es∼ρπb (s|o) [Aπb (s, a)]]] + O(ϵ2 T 2 ).
(3)

3. We consider the setting in which the learner neverobserves the features the baseline policy used to makedecisions. Under some assumptions, we show how we canuse the proxy correction (Miao et al., 2018; Tchetgen et al.,2020; Bennett & Kallus, 2021) to correct estimate Q-values
in the learner’s observation space. We derive two algorithms
for doing so, one of a generative-modeling flavor and oneof a game-theoretic flavor.

DefineAπb (o, a) = Es∼ρπb (s|o) [Aπb (s, a) − V πb (s)],We now turn our attention to formalizing the problem weconsider.

and Qπb (o, a), V πb (o) analogously. The above result tells
us that as long as we are able to estimate Qπb (o, a) values
correctly, putting more probability mass on actions with positive advantages with a step of size ϵ < O( T1 ) is sufficientto guarantee improvement.

2. Formalism
Consider a POMDP (Kaelbling et al., 1998) with two different observation spaces, O1 and O2 . We can represent our
setup via the following structural causal model (SCM).

For comparison, consider the (vanilla) policy gradient for
MDPs.

st

∇θ J(πθ ) = Eξ∼πθ

at

#
πθ

∇θ log(πθ (a|s))Q (s, a) .

(5)
One usually subtracts a baseline to reduce variance. 1 The
optimal control variate is the value function (Sutton & Barto,2018), which gives us the following expression

Qt

o2t

∇θ J(πθ ) = Eξ∼πθ

Figure 1. The causal model we consider in this paper. The green
arrow denotes the dependence in the training data, the orange arrowdenotes the dependence of the policy we’re learning.

J(πθ ) = Eξ∼πθ

" T
X
t

o1t

Let

(4)

" T
X

∇θ log(πθ (a|s))A (s, a) .
(6)

Similar to our above result, the theory of policy gradients(Kakade & Langford, 2002; Agarwal et al., 2020) tells usthat mixing in a small enough portion of the greedily optimal
policy is sufficient to guarantee improvement. Thus, themain difference comes in the definition of the advantage
function: for MDPs, advantages depend on states while forour setup, advantages depend on the observation space ofthe policy we’re mixing in. We therefore turn our attention
to estimating Q-values / advantages.

(1)tdenote the performance of a policy. We assume we see data
generated by some behavior policy πb : O1 → ∆(A) thatconsists of tuples of (o2t , at , Qt ) and would like to find a direction in policy-space that improves J(πb ). More formally,
we’re looking for a policy π : O2 → ∆(A) such that
J((1 − ϵ)πb + (ϵ)π) ≥ J(πb )

#
πθ

t

#
r(st , at )

" T
X

1

One usually pre-conditions with the Fisher information matrix
to follow the natural policy gradient (Kakade, 2001), but we ignore
those details for simplicity.

(2)2Complementing a Policy with a Different Observation Space
Greedy Switch Correctp(error ) = e

figure corresponds to the learner under-estimating the valueof chemotherapy when few people need it and doctors rarely
spuriously prescribe it.

Q̂(chemo ) > Q̂(!chemo )0.8

Q(chemo ) < Q(!chemo )0.6

If such a problem can happen even in a single-step problem,
it can clearly happen with multiple steps of interaction. We
therefore turn our attention to using techniques from causal
inference to de-bias our Q-value estimates.

0.4
0.2
0.2

0.4

0.6

0.8

p(cancer ) = α

2.3. Estimating Q-values when both Observations are
Available

Figure 2. We see that when the base rate of cancer is low and thedoctor makes few mistakes, we under-estimate the value of chemo
and would end up providing poor care to people. This is indicatedin the figure by the red dots.

We first consider the relatively simple setting when both setsof observations are provided to the learner in the behavioral
data. Formally, we want to compute
Q(o2t , a) = E[

2.2. An Example of why the Naive Q-value Estimate
Failsr(sτ , aτ )|do(at ), o2t ].

(9)

τ =t

We illustrate the issue with the naive Q-value estimator witha simple, one-step example.

Using Qt as a shorthand for the sum of rewards over thehorizon, we can apply the rules of do-calculus (equivalently,
the backdoor adjustment formula (Pearl, 2009)) to computethe above causal effect.
X
E[Qt |do(at ), o2t ] =P (o1t |do(at ), o2t )E[Qt |do(at ), o1t , o2t ]Problem 2.2 (Cancer Treatment). Consider the following
one-step (contextual bandit) problem. A doctor observes
whether a patient has cancer (O1 = {cancer, !cancer}) anddecides whether to instruct them to begin chemotherapy.
The rewards are as follows:o1t=

Q(cancer, chemo) = 1, Q(!cancer, chemo) = 0, (7)X

P (o1t |do(at ), o2t )E[Qt |at , o1t , o2t ]o1t

Q(cancer, !chemo) = 0, Q(!cancer, !chemo) = 0.5=
We assume α fraction of the population has cancer. The
doctor makes a mistake (i.e. picking the action with thelower Q value) with probability ϵ. We want to learn a policy
that, without any observations, improves upon the doctor(O2 = ∅).

X

P (o1t |o2t )E[Qt |at , o1t , o2t ],o1twhere the second equality follows from (Qt ⊥⊥ a|o1t , o2t )Gat
and Rule 2 and the third equality follows from (o1t ⊥⊥a|o2t )Gat and Rule 3. Lastly, simple algebra tells us that

To solve this problem, we compute both Q(∅, chemo) andQ(∅, !chemo) by averaging the rewards received when apatient was / wasn’t given chemotherapy, regardless of theirunderlying cancer status. More formally,Q̂(∅, a) = Eπb [Q(S, a)]TXP (o1t , o2t )
1 2
E[Q
|a
,
o
,
o
]
t t t
t
P (o1 )P (o2 ) 1t 2 t= Eo1t r(ot , ot )E[Qt |at , o1t , o2t ] , (10)E[Qt |do(at ), o2t ] = Eo1t

where we use r(o1t , o2t ) to denote the density ratio between
the joint and the product of marginals.

(8)

We pick the action with the higher Q value always (i.e.
π(∅) = maxa Q(∅, a)). We visualize the region of (α, ϵ) in
which we pick !chemo in yellow in the below figure. Given
access to the base rate of cancer (α), we can also compute
the ground-truth Q values, which tell us when giving nopatients chemotherapy would be preferable to giving allof them chemotherapy. Intuitively, this is true when few
patients have cancer. We visualize this in teal in the figure.

The above expression gives us the following procedure forpolicy improvement.
1. Estimate E[Qt |at , o1t , o2t ] via regression.
2. Fit a density ratio estimate for r(o1t , o2t ).
3. Fit Q(o2t , at ) via a secondary regression with targets
determined via the previous two components.

Mixing in any amount of the greedy optimal policy is harmful when Q̂ and Q rank the two actions in the opposite order,
which correspond to the red dots in the above figure. Intuitively, the cluster of red dots in the bottom left of the

4. Compute the greedily optimal policy and plug into
your favorite policy improvement method to control
the step size.
3Complementing a Policy with a Different Observation Space

o2t

X

Z

U

A

W

Y

at−1

at+1

o1t , stat

Qt

st−1 , st

o2t−1

at

o2t

Qt

Figure 3. Causal models that satisfy the independence assumptions required to apply the proxy correction.

We can perform the seconds step via classification between paired and unpaired samples of o1t , o2t (Hyvarinen& Morioka, 2016).

The second is via generative modeling. First, we learn amodel of P (W |a, z, x). Then, we minimize the following
MSE

2.4. Estimating Q-values with only Learner
Observations

min EA,Z,X [(E[Y |a, z, x]−E[g(a, x, W )|a, z, x])2 ]. (15)g

Importantly, by drawing separate samples of W for both
evaluations of g inside the square, we avoid the so-called“double-sample” problem that can lead to inconsistent estimates of a causal effect (Baird, 1995).

We first describe the general technique we use before specializing to our case of POMDPs. The argument we present
below is a slight generalization of that in Xu et al. (2021).

A third family of techniques, described by Xu et al. (2021),
alternates between solving two regression problems, eachof which corresponds to one of the stages of the above
generative modeling approach. More formally, we assume
that we can write g and the conditional feature mean as

2.4.1. T HE P ROXY C ORRECTION
We are attempting to determine causal effect E[Y |do(a), x].
Unfortunately, Y, A, X all have a common parent: unobserved confounder U . However, we have access to two“proxies”, W, Z, for the unobserved U that we can use to
estimate the effect. Formally, we require the following independence conditions (Miao et al., 2018):
C1: W ⊥⊥ (A, Z)|U, X

g(a, x, w) = uT (ΨA(2) (a) ⊗ ΨX(2) (x) ⊗ ΨW (w))
EW |a,z,x [ΨW (w)] = V (ΨA(1) (a) ⊗ ΨZ (z) ⊗ ΨX(1) (x)),(11)

where all Ψ are learned neural networks and u and V area vector and matrix, respectively. Observe that because
we take a tensor product ⊗ between the output of neuralnetwork features which can be arbitrarily high dimensional,
An example of a graph that satisfies these conditions is
the above separation into functions of each random variabledrawn in Figure 3, left. Consider the following set of condiis without meaningful restriction on functions representable.
tional moment restrictions (CMR) on g:
We then alternate between minimizing two loss functions
E[Y |a, z, x] = E[g(a, x, W )|a, z, x], ∀(a, z, x) ∈ (A×Z×X ). over subsets of the parameters. We refer interested readers
to Xu et al. (2021)’s excellent paper for more details.
(13)
We prove the following theorem.
2.4.2. A PPLYING THE P ROXY C ORRECTION TO
Theorem 2.3. Define
POMDP S
C2: Y ⊥
⊥ Z|A, U, X

(12)h(a, x) = EW |x [g(a, x, W )].

(14)

Due to space concerns, we only write out the full form of
the game-solving variant.

If g satisfies Equation 13, then E[Y |do(a), x] = h(a, x).

Approach 1. Consider the SCM in Figure 3, center.
Swamy et al. (2022) provide two methods for learning a gthat satisfies the CMR in one of two ways. The first is via
game-solving:

Observe that we satisfy the requisite independence assumptions:C1: at−1 ⊥⊥ (at , at+1 )|st , o2t , o1t(16)

min max E[f (A, X, Z)(Y −g(A, X, W ))−0.5f 2 (A, X, Z)].
gf

4

C2: Qt ⊥⊥ at−1 |at , st , o2t , o1t(17)Complementing a Policy with a Different Observation Space

Putting it all together, we suggest the following method forpolicy improvement
1. Solve gamemin max E[f (at , o2t , at−1 )(Qt −g(at , o2t , at+1 ))−0.5f 2 (at , o2t , at−1 )]gf

2. Fit generative model of P (at+1 |ot ).
3. Fit a model of Qπb (o2t , at ) = Eat+1 |ot [g(at , ot , at+1 )]via regression.
4. Compute the greedily optimal policy and plug into apolicy improvement method to control the step size.
Note that the generative model we have to fit in the second
step is just a policy.
Approach 2. An alternative approach which would allow us
to estimate the same quantity would be to use the past andcurrent observations as proxies (Bennett & Kallus, 2021).
We visualize this approach graphically in Figure 3, right.
Note that because we directly use the observation as one ofthe proxies, we don’t need to perform a secondary averaging.
We can solve the following game to recover Qt :min max E[f (at , ot−1 )(Qt − g(at , ot )) − 0.5f 2 (at , ot−1 )]
g

f

(18)
After solving for g, we can directly plug it into an advantage
estimation procedure (Schulman et al., 2018), allowing us
to compute causally correct policy gradients.

5Complementing a Policy with a Different Observation Space

References

Miao, W., Geng, Z., and Tchetgen Tchetgen, E. J. Identifying causal effects with proxy variables of an unmeasured
confounder. Biometrika, 105(4):987–993, 2018.

Agarwal, A., Jiang, N., Kakade, S. M., and Sun, W. Reinforcement learning: Theory and algorithms. CS Dept.,
UW Seattle, Seattle, WA, USA, Tech. Rep, pp. 10–4, 2019.

Papadimitriou, C. H. and Tsitsiklis, J. N. The complexity ofmarkov decision processes. Mathematics of operations
research, 12(3):441–450, 1987.

Agarwal, A., Kakade, S. M., Lee, J. D., and Mahajan, G. Optimality and approximation with policy gradient methods
in markov decision processes. In Conference on Learning
Theory, pp. 64–66. PMLR, 2020.

Pearl, J. Causal inference in statistics: An overview. 2009.
Schulman, J., Moritz, P., Levine, S., Jordan, M., and Abbeel,
P. High-dimensional continuous control using generalized
advantage estimation, 2018.

Bagnell, J., Kakade, S. M., Schneider, J., and Ng, A. Policy
search by dynamic programming. Advances in neural
information processing systems, 16, 2003.
Baird, L. Residual algorithms: Reinforcement learning withfunction approximation. In Machine Learning Proceedings 1995, pp. 30–37. Elsevier, 1995.

Singh, S. P., Jaakkola, T., and Jordan, M. I. Learning without state-estimation in partially observable markovian
decision processes. In Machine Learning Proceedings
1994, pp. 284–292. Elsevier, 1994.

Baxter, J. and Bartlett, P. L. Infinite-horizon policy-gradientestimation. Journal of Artificial Intelligence Research,
15:319–350, 2001.

Sutton, R. S. and Barto, A. G. Reinforcement learning: An
introduction. MIT press, 2018.
Swamy, G., Choudhury, S., Bagnell, J. A., and Wu,
Z. S. Game-theoretic algorithms for conditional moment
matching, 2022.

Bennett, A. and Kallus, N. Proximal reinforcement
learning: Efficient off-policy evaluation in partiallyobserved markov decision processes. arXiv preprint
arXiv:2110.15332, 2021.

Tchetgen, E. J. T., Ying, A., Cui, Y., Shi, X., and Miao,
W. An introduction to proximal causal learning. arXiv
preprint arXiv:2009.10982, 2020.

Daumé, H., Langford, J., and Marcu, D. Search-based
structured prediction. Machine learning, 75:297–325,
2009.

Williams, R. J. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Reinforcement learning, pp. 5–32, 1992.

Hyvarinen, A. and Morioka, H. Unsupervised feature extraction by time-contrastive learning and nonlinear ica.
Advances in neural information processing systems, 29,
2016.

Xu, L. and Gretton, A. A neural mean embedding approachfor back-door and front-door adjustment, 2022.
Xu, L., Kanagawa, H., and Gretton, A. Deep proxy causal
learning and its application to confounded bandit policy
evaluation. Advances in Neural Information Processing
Systems, 34:26264–26275, 2021.

Jaakkola, T., Singh, S., and Jordan, M. Reinforcement learning algorithm for partially observable markov decision
problems. Advances in neural information processing
systems, 7, 1994.
Kaelbling, L. P., Littman, M. L., and Cassandra, A. R. Planning and acting in partially observable stochastic domains.
Artificial intelligence, 101(1-2):99–134, 1998.
Kakade, S. and Langford, J. Approximately optimal approximate reinforcement learning. In Proceedings of the
Nineteenth International Conference on Machine Learning, pp. 267–274, 2002.
Kakade, S. M. A natural policy gradient. Advances in neural
information processing systems, 14, 2001.
Littman, M. L. Memoryless policies: Theoretical limitations and practical results. In From Animals to Animats
3: Proceedings of the third international conference on
simulation of adaptive behavior, volume 3, pp. 238. MIT
Press Cambridge, MA, USA, 1994.
6

Complementing a Policy with a Different Observation Space

A. Proofs
A.1. Proof of Theorem 2.1
Proof. Via the performance difference lemma (Kakade & Langford, 2002),
J(πϵ ) − J(πb ) = T Ea,o∼ρπϵ [Es∼ρϵ (s|o) [Aπb (s, a)]].

(19)

Fix o ∈ O2 . ConsiderEa∼ρπϵ (o) [Es∼ρϵ (s|o) [Aπb (s, a)]] =X(πb (a|o) + ∆(a|o))

a∈A

=

X

∆(a|o)

X

X

X
a∈A

(ρπb (s|o) + ∆(s|o))(Qπb (s, a) − V πb (s))

s∈Sπ(a|o)a∈A=ϵ(ρπb (s|o) + ∆(s|o))(Qπb (s, a) − V πb (s))

s∈S

a∈A

=ϵ

X

X

(ρπb (s|o) + ∆(s|o))(Qπb (s, a) − V πb (s))

s∈Sπ(a|o)

X(ρπb (s|o))(Qπb (s, a) − V πb (s)) + O(ϵ2 T 2 ).

s∈SThe last line comes from the fact thatdT V (ρπb (s|o), ρπϵ (s|o)) ≤ ϵT,(20)

as proved in Agarwal et al. (2019).
Taking the average over observations tells us that
J(πϵ ) − J(πb ) = ϵT Eo∼ρπϵ [Ea∼π(a|o) [Es∼ρπb (s|o) [Aπb (s, a)]]] + O(ϵ2 T 2 ).

(21)

A.2. Proof of Theorem 2.3
Proof. Via backdoor adjustment (Pearl, 2009),
E[g(a, x, W )|a, z, x] = EU |a,z,x [E[g(a, x, W )|a, z, x, u]] = EU |a,z,x [E[g(a, x, W )|x, u]],(22)

where the second equality comes from assumption C1. Next, we note thatE[Y |a, z, x] = EU |a,z,x [E[Y |a, z, x, u]] = EU |a,z,x [E[Y |a, x, u]],(23)

where the second equality comes from assumption C2. Now, if we satisfy the CMR uniformly, we know that
EU |a,z,x [E[g(a, x, W )|x, u] − E[Y |a, x, u]] = 0.

(24)

If we assume ”completeness” of our confounder (Miao et al., 2018; Tchetgen et al., 2020), this implies that the above
condition implies the following pointwise guarantee:
E[g(a, x, W )|x, u] = E[Y |a, x, u].

(25)This guarantee allows us to plug the LHS into the backdoor adjustment formula for the causal effect. More formally,E[Y |do(a), x] = EU |x [E[Y |a, u, x]] = EU |x [E[g(a, x, W )|x, u]] = EW |x [E[g(a, x, W )],where the first equality comes from the backdoor adjustment formula and the last from the law of iterated expectation.

7(26)