WDIBS: Wasserstein deterministic information bottleneck for state abstraction to balance state-compression and performance

As an important branch of reinforcement learning, Apprenticeship learning studies how an agent learns good behavioral decisions by observing an expert policy from the environment. It has made many encouraging breakthroughs in real-world applications. State abstraction is typically used to compress the state space of the environment to eliminate redundant information, thereby improving learning efficiency. However, excessive compression results in poor decision performance. Therefore, it is important to balance the compression degree and decision performance. Deterministic Information Bottleneck for State abstraction (DIBS) attempts to solve this problem. Specifically, DIBS uses the information rate to represent the compression degree at first. Then, decision performance after compression is measured using the Kullback-Leibler (KL) divergence of distributions between the policy after state compression and the expert policy. However, if the two distributions do not have exactly overlapping support sets, then the KL divergence is usually infinity, which leads to poor decision performance under the low information rate. In this paper, we propose the Wasserstein DIBS (WDIBS) algorithm to optimize the trade-off between the compression degree and decision performance. Specifically, we use the Wasserstein distance to calculate the difference of the distributions between the policy after state compression and the expert policy. Even if the two distributions do not have precisely overlapping support sets, the Wasserstein distance can still reflect their actual difference, thereby ensuring that WDIBS has good decision performance under the low information rate. Theoretical analyses and experiments demonstrate that our method provides a better trade-off between the compression degree and decision performance than DIBS.

This work is supported in part by The National Nature Science Foundation of China under Grant Nos. 61772120. We thank Maxine Garcia, PhD, from Liwen Bianji (Edanz) (www.liwenbianji.cn/) for editing the English text of a draft of this manuscript.

Authors and Affiliations

1. Institute of Fundamental and Frontier Sciences, University of Electronic Science and Technology of China, Chengdu, China

   Xianchao Zhu, Tianyi Huang, Ruiyuan Zhang & William Zhu

Corresponding author

Correspondence to William Zhu.

Publisher’s note

Springer Nature remains neutral with regard to jurisdictional claims in published maps and institutional affiliations.

Appendix: A

In this section, we prove theorems and lemmas involved in algorithms.

1.1 A.1 Proofs of Theorem 1

Firstly, we provide the proof of Theorem 1. Specifically, to prove Theorem 1, we need two lemmas.

Usually, it is difficult to calculate the size of the abstract state space \(\left | {{S}_{\phi }} \right |\) in \(\mathcal {J}\). We use the size of the abstract state space \(\left | {{S}_{\phi }} \right |_{{{\rho }_{\phi }}\left (s\right )}^{{{\delta }_{{\min \limits } }}}\) in \({{\hat {\mathcal {J}}}_{WDIB}}\) to approximate \(\left | {{S}_{\phi }} \right |\). \(\left | {{S}_{\phi }} \right |_{{{\rho }_{\phi }}\left (s\right )}^{{{\delta }_{{\min \limits } }}}\) can be calculated by the size of the non-ignorable part of the alphabet under PMF defined by the following formula.

Definition 3

PMF-Used Alphabet Size: The PMF-used alphabet size of \(\mathcal {X}\) is the number of the elements whose probability under p(x) is higher than some negligibility threshold \({{\delta }_{{\min \limits } }}\in \left (0,1 \right )\):

Therefore, we use the first lemma, which associates the entropy of PMF with the maximum size of the alphabet used by that PMF [43]:

Lemma 1

Consider a discrete stochastic variable X, with alphabet \(\mathcal {X}\) and some PMF p(x). For a given threshold \({{\delta }_{{\min \limits } }}\in \left (0,1 \right )\), the PMF-used alphabet size of the alphabet is bounded, relative to some PMF p(x):

This bound is relatively large: it is not difficult to know that \(H\left (X \right )\le {{\log }_{2}}\left | \mathcal {X} \right |\) [3]. Therefore, in the worst case, the boundary can be \(\tilde {O}({1}/{{{\delta }_{{\min \limits } }}})\) times larger than the real alphabet. Nevertheless, the result allows us to correlate the entropy of a stochastic variable with the size of the alphabet it uses. In addition, by definition, the entropy \(H\left ({{\rho }_{\phi }} \right )\) of the abstract state distribution provides a lower bound for the number of bits to represent the used part of S_ϕ. In this way, the entropy is a measure of the degree of compression, which uses the fact that the entropy of the most probable state can be expressed as 0. Therefore, lower entropy means a higher compression degree.

Proof

The entropy of a random variable X is as follows:

Then, for a given maximum entropy \(H\left (X \right )\le N\), we seek to understand the largest possible \(\left | \mathcal {X} \right |_{p\left (x \right )}^{{{\delta }_{\min \limits }}}\). That is:

If the alphabet \(\mathcal {X}\) takes the maximum size, the variable X obeys a uniform distribution, where the probability of each element is \({{\delta }_{{\min \limits } }}\) such that \(H\left (X \right )=N\):

Therefore, for a given PMF p(x) with entropy \(H\left (X \right )\), and a minimum threshold of probability, the minimum size of the alphabet \(\mathcal {X}\) is upper bounded:

□

Next, we introduce a second lemma that relates the expected Wasserstein distance between two policies to the difference in value realized by the policies, in expectation under some state distribution:

Lemma 2

We denote by π the set of 1-Lipschitz-policies, s.t. for two stochastic policies, π₁ and π₂ on state space S, and a fixed probability distribution over S, p(s). If, for \(\varepsilon \in {{\mathbb {R}}_{\ge 0}}\):

then:

where VMAX is an upper bound on the value function.

The above boundary allows us to correlate the distortion metric controlled by WIB with the distortion metric of the CVA objective.

Proof

Recall the Wasserstein distance between policies π₁ and π₂ for a given state s is defined as:

where \({{\rho }_{{{\pi }},{{s}_{0}}}}\) denotes the stationary distribution over states under policy π, starting in state s₀ ∈ S.

We first bound the difference between the two state distributions after t steps:

From the above bound, we get the following inequality:

then:

The expectation of the value function loss is as follows:

Then, we apply the upper bound on the possible value \(VMAX=\frac {RMAX}{1-\gamma }\) to the above equation:

Then, by Pinsker’s inequality, we conclude:

□

Now we use Lemma 1 and Lemma 2 to prove Theorem 1.

Theorem 2

The function f of WDIB objective \({{\hat {\mathcal {J}}}_{WDIB}}\) is an upper bound of the CVA Objective \(\mathcal {J}\) that is tighter than \({{\hat {\mathcal {J}}}_{DIB}}\). That is:

Proof

Consider the ϕ that minimizes \({{\hat {\mathcal {J}}}_{WDIB}}\), yielding the value of at most N + βε, where:

Then, by Lemma 1, we know:

By Lemma 2, we know:

Therefore, we conclude:

Thus, we can obtain the upper bound of \(\mathcal {J}\), which can be seen as a function of the \({{\hat {\mathcal {J}}}_{WDIB}}\). □

The DIBS algorithm uses the following function as the upper bound of:

According to Pinsker’s inequality, the following inequality can be derived:

Then,

1.2 A.2 Extensions

Lastly, we briefly discuss the situation where the agent is located in high-dimensional observation space. In this case, the state abstractions are usually stochastic. We utilize a known result relating the KL divergence to mutual information by variational autoencoder (VAE) [36]. During the process of training VAE, the KL divergence between \({{q}_{\psi }}\left (z\left | x \right . \right )\) and some prior distribution over latent codes, \(p\left (z \right )\), is minimized in expectation over the data distribution, p(x); by the way, \({\mathbb {E}_{p\left (x \right )}}\left [ {{D}_{KL}}\left ({{q}_{\psi }}\left (z\left | x \right . \right )\left \| p\left (z \right ) \right . \right ) \right ]\) is minimized. We treat x as the ground state representation S and z as the abstract state S_ϕ.

Proof

To prove that our objective function \({{{\hat {\mathcal {J}}}}_{WSIB}}\) is a tighter upper bound of \({{{\hat {\mathcal {J}}}}_{SIB}}\) than \({{{\hat {\mathcal {J}}}}_{SIB\_N}}\), we utilize a known result [44]; [37] whose proof we replicate here for completeness with \(q\left (z,x \right )=p\left (x \right ){{q}_{\psi }}\left (z\left | x \right . \right )\) and \(q\left (z \right )={\mathbb {E}_{p\left (x \right )}}\left [ {{q}_{\psi }}\left (z\left | x \right . \right ) \right ]\):

So the following inequality can be derived:

This formula is equivalent to the following form:

It usually supposes that both \({{q}_{\psi }}\left (z\left | x \right . \right )\) and \(p\left (z \right )\) obey the Gaussian distribution to use the Gaussian reparameterization technique [36]. Let us suppose that these distributions have dimension k, means μ_i, and covariance matrices \(\sum \nolimits _{i}\), for i = 1, 2. The following two formulas can be obtained:

and

For simplicity, we consider \(\sum \nolimits _{1}{\text {=}\sum \nolimits _{2}{\text {=}w{{I}_{k}}}}\), where \(w \in \left [ 0,1 \right ]\) is the weight and μ₁≠μ₂. When combined with the − k term, the trace term in Wasserstein distance, the trace term in the KL divergence and the log-determinant ratio will be 0. So these two equations become:

and

Note that Wasserstein distance does not change if the variance changes whereas the KL divergence does. Therefore, the following inequality can be derived:

At same time, according to Equation (46), the following inequality can be derived:

So the following inequality can be derived:

Therefore, we conclude:

□

Reprints and permissions