Disentangled (Un)Controllable Features

To isolate controllable features in the latent representation, $z^{c}_{t}$ is used to make an action-conditioned forward prediction in latent space. In the context of a continuous latent space and deterministic transitions, $z^{c}$ is updated using a mean squared error (MSE) forward prediction loss $\mathcal{L}_{c}=\big{|}\hat{z}^{c}_{t+1}-z^{c}_{t+1}\big{|}^{2}$, where $\hat{z}^{c}_{t+1}$ is the action-conditioned residual forward prediction of the parameterized function $T_{c}(z,a;\theta_{c}):\mathcal{Z}\times\mathcal{A}\rightarrow\mathcal{Z}$:

and the prediction target $z^{c}_{t+1}$ is part of the encoder output $f(s_{t+1};\theta_{enc})$. Note that the full latent state $z_{t}$ is necessary in order to predict $\hat{z}^{c}_{t+1}$ (e.g. the uncontrollable features could represent a wall or other static structure that is necessary for the prediction of the controllable features). Furthermore, the uncontrollable latent partition input $z^{u}_{t}$ is accompanied by a stop gradient to discourage the presence of controllable features in $z^{u}$. When minimizing $\mathcal{L}_{c}$, both the encoder ($\theta_{enc}$) as well as the predictor ($\theta_{c}$) are updated, which allows shaping the representation $z^{c}$ as well as learning the internal dynamics.

To express uncontrollable features in the latent space, $z^{u}_{t}$ is used to make a state-only (not conditioned on the action $a_{t}$) forward prediction in latent space. This enforces uncontrollable features within the uncontrollable latent partition $z^{u}$, since features that are action-dependent cannot be accurately predicted with the preceding state only. Following a residual prediction, $z^{u}$ is then updated using a MSE forward prediction loss $\mathcal{L}_{u}=\big{|}\hat{z}^{u}_{t+1}-z^{u}_{t+1}\big{|}^{2}$, with $\hat{z}^{u}_{t+1}$ defined as:

and $T_{u}(z^{u};\theta_{u}):\mathcal{Z}\rightarrow\mathcal{Z}$ representing the parameterized prediction function. The target $z^{u}_{t+1}$ is part of the output of the encoder $f(s_{t+1};\theta_{enc})$. When minimizing $\mathcal{L}_{u}$, both $\theta_{enc}$ and $\theta_{u}$ are updated. In this way the loss $\mathcal{L}_{u}$ drives the latent representation $z^{u}$, which is conditioned on $\theta_{enc}$ according to $(z^{c}_{t},z^{u}_{t})=f(s_{t};\theta_{enc})$, to only represent the features of $s_{t}$ that are not conditioned on the action $a_{t}$.

Minimizing a forward prediction loss in latent space $\mathcal{Z}$ is prone to collapse [9, 16], due to the convergence of $\mathcal{L}_{c}$ and $\mathcal{L}_{u}$ when $f(s_{t};\theta_{enc})$ is a constant    $\forall\hskip 5.69054pts_{t}\in\mathcal{S}$. To avoid representation collapse when using forward predictors, a contrastive loss is used to enforce sufficient diversity in the latent representation:

where $C_{d}$ represents a constant hyperparameter and $\bar{z}_{t}$ is a ‘negative’ batch of latent states $z_{t}$, which is obtained by shifting each position of latent states in the batch by a random number between 0 and the batch size. In the random maze environment, an additional contrastive loss is added to further diversify the controllable representation:

where $z^{c}_{t}$ is obtained from randomly sampled trajectories. This additional regularizer proved neccessary to avoid collapse of $z^{c}$ when moving to a near infinite number of possible mazes. More information on this subject can be found in Appendix A-D. The resulting contrastive loss $\mathcal{L}_{H}$ for the random maze environment then consists of $0.5\mathcal{L}_{H_{1}}+0.5\mathcal{L}_{H_{2}}$. The total loss used to update the encoder’s parameters now consists of $\mathcal{L}_{enc}=\mathcal{L}_{c}+\mathcal{L}_{u}+\mathcal{L}_{H}$.

When using a controllable latent space $z^{c}\in\mathbb{R}^{x},x\in\mathbb{N}$, where $x>g$, with $g$ representing the number of dimensions needed to portray the controllable features, some information about the uncontrollable features in the controllable latent representation might be present (see Appendix C-B). This is due to the non-enforcing nature of $\mathcal{L}_{c}$, as the uncontrollable features are equally predictable with or without the action. To ensure that no information about the uncontrollable features is kept in the controllable latent representation, an adversarial component is added to the architecture in Fig. 2. This is done by updating the encoder with an adversarial loss $\mathcal{L}_{adv}$ and reversing the gradient [28]. The adversarial loss is defined as

with $\hat{z}^{u}_{t}=T_{adv}(z^{c}_{t};\theta_{adv})$, where $\hat{z}^{u}_{t}$ is the uncontrollable prediction of the parameterized function $T_{adv}(z^{c};\theta_{adv}):\mathcal{Z}\rightarrow\mathcal{Z}$ and $z^{u}_{t}$ is the target. Intuitively, since the parameters of $T_{adv}(z^{c};\theta_{adv})$ are being updated with $\mathcal{L}_{adv}$ and the parameters of $f(s;\theta_{enc})$ are being updated with $-\mathcal{L}_{adv}$, the prediction function can be seen as the discriminator and the encoder can be seen as the generator [29]. The discriminator tries to give an accurate prediction of the uncontrollable latent $z^{u}$ given the controllable latent $z^{c}$, while the generator tries to counteract the discriminator by removing any uncontrollable features from the controllable representation. In our case, the predictor is a multi-layer perceptron (MLP), which means that minimizing $\mathcal{L}_{adv}$ enforces that no nonlinear relation between $z^{c}$ and $z^{u}$ can be learned. We hypothesize that this is a deterministic approximation of minimizing the Mutual Information (MI) between $z^{u}$ and $z^{c}$. When using the adversarial loss, the combined loss propagating through the encoder consists of $\mathcal{L}_{enc}=\mathcal{L}_{c}+\mathcal{L}_{u}+\mathcal{L}_{H}-\mathcal{L}_{adv}$. Here the minus term in $-\mathcal{L}_{adv}$ represents a gradient reversal to the encoder. Note that the losses are not scaled, as this did not prove to be necessary for the experiments conducted.