# Categorical-measure-theoretic approach to optimal policies tending to seek power

The paper [Optimal Policies Tend to Seek Power](https://arxiv.org/abs/1912.01683) tries to justify the claim in the title with a mathematical argument, using a clever formal definition of "power". I like the general approach, but I think that parts of the formalisation do not correspond well to the intuitive understanding of the concepts involved. This wouldn't be an issue if the goal of the paper was to prove a mathematical theorem. Here, however, the goal is to be able to deduce something about the behaviour of future AI systems. Therefore, the interface between "math" and "the world" is much more important. This post attempts to take another route to formalisation, which I think is closer to the intuition, albeit more difficult to work with mathematically.

# Intro

The power of a state $s$, written $POWER_\mu(s)$, is defined as (roughly) the average optimal value of that state, over some chosen distribution of rewards $\mu$. The paper quite convincingly argues for it:
> Power is the ability to achieve a range of goals. In an MDP, the optimal value function $V^*_R(s, \gamma)$ captures the agent’s ability to “achieve the goal” $R$. Therefore, average optimal value captures the agent’s ability to achieve a range of goals.

The problem with this approach is that it's very sensitive to the choice of the distribution of rewards $\mu$. In particular, one might take a Dirac delta on a reward giving 1 in any chosen state, and 0 everywhere else, which would trivially assign maximal power to that particular state. To combat this difficulty, authors introduce some technical devices allowing them to say "for most distributions, the power of (one, obviously better, state) is greater than the power of (another, obviously worse, state)".

To me, "$X$ happens for most distributions of rewards" should mean that there is some reasonable, uninformative, broad prior $\mathcal{D}$ over distributions of rewards, and that $\mathbb{P}_\mathcal{D}(X) > \frac{1}{2}$. The approach chosen in the paper not only doesn't correspond to this intuitive understanding, but seems to me to be (in certain sense) proving the thesis by assuming the conclusion.

More concretely, the approach taken in the paper is as follows. Given a reward $R$ and a permutation $\sigma$ over states of the MDP, one can obtain another reward $(\sigma R)(s) = R(\sigma(s))$ by permuting the states' rewards. This also works if we have not one reward $R$, but a distribution over rewards, by pushing it forward by $\sigma$. Therefore, each such permutation $\sigma$ induces an orbit $S_\sigma$ in the space of (distributions of) rewards. The authors then say that the power of one state $s_1$ is greater than power of another state $s_2$ *for most distributions*, if, for each permutation $\sigma$, for the majority of distributions of rewards $\mu$ in the orbit $S_\sigma$, we have the inequality $POWER_\mu(s_1) > POWER_\mu(s_2)$. 

This is arguably quite far from how anyone would interpret the phrase "for most distributions" in real life. What we care about is the global behaviour in the whole space, not the the local behaviour in orbits. (Moreover, the permutations considered here are arbitrary in a sense of not respecting the structure of the MDP - but that is another issue).

The main problem of making the probabilistic meaning of "most" proposed above work here is that it is mathematically difficult to work in the infinite-dimensional space of distributions over rewards. There are many pathologies, such as the fact that there is no counterpart to the (Lebesgue-) uniform measure (after scouring the internet, I found that authors make this argument in a response to the review of the paper). 

I still think it is possible to make progress on this. It should be possible to use something like Dirichlet/Pitman-Yor/Gaussian processes, or even something more geometric such as Wasserstein distance, based on the [metric structure](https://arxiv.org/pdf/2006.13900.pdf) of the rewards in an MDP. However, regardless of the concrete distribution, the first step is to prove that the event $POWER_\mu(s_1) \leq POWER_\mu(s_2)$ is measurable in the space of distributions of rewards (that is, in $\mu$).

Thus, this post contains:
 -  the proof of this fact, and
 -  some (unfruitful) exploration of two simple cases of concrete distributions over distributions of rewards

The main technical device I use is the Giry monad, which gives a characterisation of the canonical $\sigma$-algebra on the space of probability distributions. I include the basic definition below, but those unfamiliar with the topic might enjoy [this tutorial](https://www.youtube.com/watch?v=3Da88Tgw_rM) - which assumes only basic familiarity with measure theory and category theory.

# Technical content

## What the original paper did

To refresh the memory, the original paper proceeds as follows. We assign "power" to each state of a MDP (under a fixed distribution $\mu$ over reward functions), to be:
  $$
  POWER_\mu(s) = \frac{1 - \gamma}{\gamma} \cdot \mathbb{E}_{R \sim \mu}[V^*_R(s) - R(s)]
$$
For a set $\mathcal{D}$ of distributions over rewards, we define the inequality between powers of states $s_1, s_2$ to be:
$$
POWER_\mu(s_1) \geq_{most(\mathcal{D})} POWER_\mu(s_2)
$$
if and only if, for all $B_\mu = orb(\mu, \Sigma_\mathcal{S}) = \{\sigma \cdot \mu: \sigma \in \Sigma_S\}$, for the majority of $\nu \in B$ (which is a finite set) we have:
$$
POWER_\nu(s_1) \geq POWER_\nu(s_2)
$$

(The paper considers only rewards of a form $R: S \to \mathbb{R}$ and only deterministic policies $\pi \in S \to A$.)

## Giry monad

First, we note that if $S$ and $A$ are finite (countable case is the same), then the space of reward functions $S \to \mathbb{R}$ (in general, $S \times A \times S \to \mathbb{R}$) is equal to $\mathbb{R}^n$ for some $n \in \mathbb{N}$ - therefore, it is a standard Borel space.

On the category of Borel spaces, we have the Giry monad $G$, which assigns, to each measurable space $X$, the space of the probability measures on $GX$. Therefore, $G$ points out a canonical $\sigma$-algebra on the $\Delta(\mathbb{R}^n)$, which we denote by $\Sigma$.

We wiil describe $\Sigma$. First, we say that for a fixed $U \in \Sigma_{\mathbb{R}^n}$, the evaluation function is given by:
$$
\begin{align*}
ev_U &: \Delta(\mathbb{R}^n) \to \mathbb{R} \\
ev_U(\mu) &= \mu(U)
\end{align*}
$$
Then, $\Sigma$ It is generated by all $ev_U$, that is, by sets all sets of the form $ev^{-1}_U(V)$, for any measuable set $V \in \Sigma_\mathbb{R}$.

**Lemma 1**: For any measurable function $f: \mathbb{R}^n \to \mathbb{R}$, we can take the evaluation with respect to this function:
$$
ev_f(\mu) = \int f(x) d\mu(x)
$$
and then $\Sigma$ is equivalently generated by all $ev^{-1}_f(V)$.

## Measurability of the $POWER$ inequality

Let us fix some states $s_1, s_2$. We can ask whether the set
$$
A = \{\mu \in \mathcal{D} : POWER_D(s_1) \geq POWER_D(s_2)\}
$$
is measurable in $\Sigma$ (since ultimately, we want to say something about it's measure).

Unrolling the definitions gives us
$$
\begin{align}
POWER_\mu(s_1) &\geq POWER_\mu(s_2) \iff \\
\frac{1 - \gamma}{\gamma}\mathbb{E}_{R \sim \mu}[V^*_R(s_1) - R(s_1)] &\geq \frac{1 - \gamma}{\gamma}\mathbb{E}_{R \sim \mu}[V^*_R(s_1) - R(s_1)] \iff \\
\mathbb{E}_{R \sim \mu}[(V^*_R(s_1) - V^*_R(s_2)) - (R(s_1) - R(s_2))] &\geq 0
\end{align}
$$
Let's now denote:
$$
\begin{align*}
f &: \mathbb{R}^n \to \mathbb{R} \\
f(R) &= V^*_R(s_1) - V^*_R(s_2) - (R(s_1) - R(s_2)) \\
\\
F &: \Delta(\mathbb{R}^n) \to \mathbb{R} \\
F(\mu) &= \mathbb{E}_{x \sim \mu}[f(x)] = \int f(x) d\mu(x)
\end{align*}
$$
Rewriting the definition of $A$ gives us:
$$
A = \{\mu \in \Delta(\mathbb{R}^n) : \mathbb{E}_{x \sim \mu}[f(x)] \geq 0 \} = \{\mu \in \Delta(\mathbb{R}^n) : F(\mu) \geq 0 \} = F^{-1}([0, \infty))
$$
Now, we only have to argue that $f$ is measurable - then, by **Lemma 1**, we will know that $A$ is measurable. It is a sum of four terms - it is enough to show that $g(R) = R(s_i)$ and $h(R) = V^*_R(s_i)$ are measurable.

The first function $g(R) = R(s_1)$ is just a projection to the $s_i$-th coordinate, therefore measurable.

The second function
$$
h(R) = \max_{\pi \in \Pi} f^{\pi, s_i}(\gamma)\cdot R
$$
is just a maximum over a finite set $\Pi = S \to A$ of linear functionals (on a finite-dimensional space), therefore measurable as well - where $f^{\pi, s_i}(\gamma)$ is the visit distribution
$$
f^{\pi, s_i}(\gamma) = \mathbb{E}_{\{s_t\} \sim \pi}\left[\sum_{t=0}^\infty \gamma^t e_{s_t}\right]
$$
(doesn't really matter what it is, the important point is that it is independent of $R$).

## Probabilities on $\mathcal{D}$

Ok, so now we can think about what is the probability mass of $A = A(s_1, s_2)$ under different probability distributions over $\Delta(\mathbb{R}^n)$. To make our lives easier, we'll now switch to $\Delta(I^n)$, $I = [0, 1]$. This doesn't change anything: we can rescale reward function by any factor and get an equivalent MDP (proof: just look at the Bellman equations).

Two (short) attempts I make below are unsuccessful. Both of them fail in a similar ways: by making the uniform distribution over rewards "too central", and by not having the full support over all probability distributions.

### Lifting uniform distribution by Giry monad

We have one canonical distribution over distributions: take the uniform distribution $p$:
$$
1 \xrightarrow{p} X
$$
and apply the Giry monad:
$$
G(1) \simeq 1 \xrightarrow{Gp} GX
$$
But unfortunately, $Gp$ is just the uniform distribution over the Dirac deltas in $\delta_{R} \in GX$. Then, the situation is equivalent to calculating the measure of the set $POWER_\mu(s_1) \geq POWER_\mu(s_2)$ in one case of the uniform distribution over rewards (that is, a Diract delta on the uniform distribution of the rewards).

### Dirichlet process

We can take Dirichlet process $\mathcal{D} := D(\mathcal{U}[0, 1]^n, \alpha)$ centered around the uniform distribution, for a fixed $\alpha$.

**Definition**: For a given distribution $H$ on $X$ and $\alpha \in (0, \infty)$, the Dirichlet process $D(H, \alpha)$ on the space $\Delta(X)$ is characterized by the following property: for any partition $A_1, A_2, \dots A_n$ of $X$, we have that if $\mu \sim D(H, \alpha)$, then:
$$
(\mu(A_1), \mu(A_2), \dots \mu(A_n)) \sim Dir(\alpha H(A_1), \alpha H(A_2), \dots \alpha H(A_n))
$$

Applying the definition to our $D(\mathcal{U}[0, 1]^n, \alpha)$ and the partition $\{A, A^c\}$, we get
$$
(\mu(A), \mu(A^c)) = (\mu(A), 1-\mu(A)) \sim Dir(\alpha\lambda(A), \alpha(1-\lambda(A))) = Beta(\alpha\lambda(A), \alpha(1-\lambda(A)))
$$

where $\lambda(A)$ the measure of $A$ under uniform distribution $\mathcal{U}[0, 1]^n$.

Therefore, we have the mean:
$$
\mathbb{E}_{\mu \sim \mathcal{D}}[\mu(A)] = \lambda(A)
$$

Unfortunately, this is again quite uninteresting. All the choice was in centering the Dirichlet process at $H = \mathcal{U}[0, 1]^n$. (Another drawback is that Dirichlet process has support only on discrete distributions, which is at odds with our initial goal of having a broad prior over all distributions of rewards - in particular, continuous ones).

## Conclusion

We proved that the event $POWER_\mu(s_1) \geq POWER_\mu(s_2)$ is measurable in the space of distibutions $\mu \sim \mathcal{D}$. We failed to point out any interesting concrete priors over distributions of reward functions, but I am optimistic that it should be possible with better understanding of the [metric](https://arxiv.org/pdf/2006.13900.pdf) (topological, linear, differential, ...?) structure of the reward functions on an MDP.