Disentangling the independently controllable factors of variation by interacting with the world
Valentin Thomas Emmanuel Bengio William Fedus Jules Pondard Philippe Beaudoin Hugo Larochelle Joelle Pineau Doina Precup Yoshua Bengio


^*Equal contribution, random order, 
^1MILA, Université de Montréal,
^2Element AI,
^3McGill University,
^4ENS Paris,
^5Google Brain,
^6Facebook AI Research,
^7Google Deepmind,
^8CIFAR Senior Fellow


Adaptive Geospatial Joins for Modern Hardware
    Andreas Kipf1,
Harald Lang1,
Varun Pandey1,
Raul Alexandru Persa1,

Peter Boncz2,
Thomas Neumann1,
Alfons Kemper1

0.51Technical University of Munich

Munich, Germany
{first.last}@in.tum.de
0.52Centrum Wiskunde & Informatica

Amsterdam, The Netherlands

boncz@cwi.nl



    December 30, 2023
=================================================================================================================================================================================================================================================================================





    

  It has been postulated that a good representation is one that disentangles the underlying explanatory factors of variation. However, it remains an open question what kind of training framework could potentially achieve that. Whereas most previous work focuses on the static setting (e.g., with images), we postulate that some of the causal factors could be discovered if the learner is allowed to interact with its environment. The agent can experiment with different actions and observe their effects. More specifically, we hypothesize that some of these factors correspond to aspects of the environment which are independently controllable, i.e., that there exists a policy and a learnable feature for each such aspect of the environment, such that this policy can yield changes in that feature with minimal changes to other features that explain the statistical variations in the observed data. We propose a specific objective function to find such factors, and verify experimentally that it can indeed disentangle independently controllable aspects of the environment without any extrinsic reward signal.




§ INTRODUCTION











When solving Reinforcement Learning problems, what separates great results from random policies is often having the right feature representation. Even with function approximation, learning the right features can lead to faster convergence than blindly attempting to solve given problems <cit.>.

The idea that learning good representations is vital for solving most kinds of real-world problems is not new, both in the supervised learning literature <cit.>, and in the RL literature <cit.>. An alternate idea is that these representations do not need to be learned explicitly, and that learning can be guided through internal mechanisms of reward, usually called intrinsic motivation <cit.>. 

We build on a previously studied <cit.> mechanism for representation learning that has close ties to intrinsic motivation mechanisms and causality. This mechanism explicitly links the agent's control over its environment to the representation of the environment that is learned by the agent. More specifically, this mechanism's hypothesis is that most of the underlying factors of variation in the environment can be controlled by the agent independently of one another. 

We propose a general and easily computable objective for this mechanism, that can be used in any RL algorithm that uses function approximation to learn a latent space. We show that our mechanism can push a model to learn to disentangle its input in a meaningful way, and learn to represent factors which take multiple actions to change and show that these representations make it possible to perform model-based predictions in the learned latent space, rather than in a low-level input space (e.g. pixels).



§ LEARNING DISENTANGLED REPRESENTATIONS










The canonical deep learning framework to learn representations is the autoencoder framework <cit.>. There, an encoder f:S→ H and a decoder g:H → S are trained to minimize the reconstruction error, s-g(f(s))_2^2. H is called the latent (or representation) space, and is usually constrained in order to push the autoencoder towards more desirable solutions. For example, imposing that H∈ℝ^K, S∈ℝ^N,  K≪ N pushes f to learn to compress the input; there the bottleneck often forces f to extract the principal factors of variation from S. However, this does not necessarily imply that the learned latent space disentangles the different factors of variations. Such a problem motivates the approach presented in this work.






Other authors have proposed mechanisms to disentangle underlying factors of variation. Many deep generative models, including variational autoencoders <cit.> and other descendants of the Helmholtz machine <cit.>, generative adversarial
networks <cit.> or non-linear versions of ICA <cit.> attempt to disentangle the underlying factors of variation by assuming that their joint distribution (marginalizing out the observed s) factorizes, i.e., that they are marginally independent.

Here we explore another direction, trying to exploit the ability of a learning agent to act in the world in order to impose a further constraint on the representation. We hypothesize that interactions can be the key to learning how to disentangle the various causal factors of the stream of observations that an agent is faced with, and that such learning can be done in an unsupervised way. 



















§ THE SELECTIVITY OBJECTIVE

We consider the classical reinforcement learning setting but in the case where extrinsic rewards are not available. We introduce the notion of controllable factors of variation ϕ∈ℝ^K which are generated from a neural network Φ(h, z), z ∼𝒩(0, 1)^m where h = f(s) is the current latent state. The factor ϕ represents an embedding of a policy π_ϕ whose goal is to realize the variation ϕ in the environment.


To discover meaningful factors of variation ϕ and their associated policies π_ϕ, we consider the following general quantity 𝒮 which we refer to as selectivity and that is used as a reward signal for π_ϕ:

    𝒮(h, ϕ) = 𝔼[logA(h', h, ϕ)/𝔼_p(φ|h)[A(h', h, φ)] |  s' ∼𝒫^π_ϕ_ss']



Here h=f(s) is the encoded initial state before executing π_ϕ and   h' = f(s') is the encoded terminal state. ϕ and φ represent factors of variation a factor.
A(h', h, ϕ) should be understood as a score describing how close ϕ is to the variation it caused in (h', h). For example in the experiments of section 4.1, we choose A to be a gaussian kernel between h'-h and ϕ, while in the experiments of section 4.2, we choose A(h', h, ϕ) = max{0,⟨ h'-h, ϕ⟩}.
The intuition behind these objectives is that in expectation, a factor ϕ should be close to the variation it caused (h', h) when following π_ϕ compared to other factors φ that could have been sampled and followed thus encouraging independence within the factors.


Conditioned on a scene representation h, a distribution of policies are feasible. Samples from this distribution represent ways to modify the scene and thus may trigger an internal selectivity reward signal. For instance, h might represent a room with objects such as a light switch. ϕ = ϕ(h,z) can be thought of as the distributed representation for the “name” of an underlying factor, to which is associated a policy and a value.  In this setting, the light in a room could be a factor that could be either on or off. It could be associated with a policy to turn it on, and a binary value referring to its state, called an attribute or a feature value.
We wish to jointly learn the policy π_ϕ(·|s) that modifies the scene, so as to control the corresponding value of the attribute in the scene, whose variation is computed by a scoring function A(h', h, ϕ) ∈ℝ. In order to get a distribution of such embeddings, we compute ϕ(h,z) as a function of h and some random noise z.






The goal of a selectivity-maximizing model is to find the density of factors p(ϕ |h), the latent representation h, as well as the policies π_ϕ that maximize 𝔼_p(ϕ|h) [𝒮(h, ϕ)].












 §.§ Link with mutual information and causality







The selectivity objective, while intuitive, can also be related to information theoretical quantities defined in the latent space.
From <cit.> we have
pq = sup_A∈ℒ^∞(q)𝔼_p [ log A] - log𝔼_q [A].
Applying this equality to the mutual information ℐ_p(ϕ, h' | h) = 𝔼_p(h'|h)[ p(ϕ|h', h)p(ϕ|h)]
gives

    ℐ_p(ϕ, h' | h) ≥sup_θ𝔼_p(ϕ|h)[ 𝒮(h, ϕ)]

where θ is the set of weights shared by the factor generator, the policy network and the encoder.

Thus, our total objective along entire trajectories is a lower bound on the causal <cit.> or directed <cit.> information ℐ_p(ϕ↦ h) = ∑_t ℐ_p(ϕ_1:t, h_t|h_t-1) which is a measure of the causality the process ϕ exercises on the process h. See Appendix <ref> for details.







§ EXPERIMENTS

 So do we


  * Keep only this experiment (we wouldn't have to worry much about the space but that is light on the experiments)

  * Add the one with planning/policy inference? <ref> (small example from the paper)

  * Add the one with factors that seem to like to toggle switch in multistep <ref> (though not the behaviour we would really expect)

  * Add the discrete one? <ref>

Seems the experiments is working, doing a figure....
We use MazeBase <cit.> to assess the performance of our approach. We do not aim to solve the game.
In this setting, the agent (a red circle) can move in a small environment (64 × 64 pixels) and perform the actions . The agent can go anywhere except on the orange blocks.










 §.§ Learned representations


pink and white for , light blue for , green for , purple black  and night blue for .

After jointly training the reconstruction and selectivity losses, our algorithm disentangles four directed factors of variations as seen in Figure <ref>: ± x-position and ± y-position of the agent. For visualization purposes we chose the bottleneck of the autoencoder to be of size K = 2.
To complicate the disentanglement task, we added the redundant action  as well as the action  in this experiment.

The disentanglement appears clearly as the latent features corresponding to the x and y position are orthogonal in the latent space. Moreover, we notice that our algorithm assigns both actions  (white and pink dots in Figure <ref>.a) to the same feature. It also does not create a significant mode for the feature corresponding to the action  (light blue dots in Figure <ref>.a) as this feature is already explained by features  and .











 §.§ Multistep embedding of policies

In this experiment, ϕ are embeddings of 3-steps policies π_ϕ. We add a model-based loss ℒ_MB=||h_t+3 - T_θ(h_t, ϕ)||^2 defined only in the latent space, and jointly train a decoder alongside with the encoder. Notice that we never train our model-based cost at pixel level.
While we currently suffer from mode collapsing of some factors of variations, we show that we are successfully able to do predictions in latent space, reconstruct the latent prediction with the decoder, and that our factor space disentangles several types of variations.









§ CONCLUSION, SUCCESS AND LIMITATIONS

Pushing representations to model independently controllable features currently yields some encouraging success. Visualizing our features clearly shows the different controllable aspects of simple environments, yet, our learning algorithm is unstable. What seems to be the strength of our approach could also be its weakness, as the independence prior forces a very strict separation of concerns in the learned representation, and should maybe be relaxed.

Some sources of instability also seem to slow our progress: learning a conditional distribution on controllable aspects that often collapses to fewer modes than desired, learning stochastic policies that often optimistically converge to a single action, tuning many hyperparameters due to the multiple parts of our model. Nonetheless, we are hopeful in the steps that we are now taking. Disentangling happens, but understanding our optimization process as well as our current objective function will be key to further progress. 





plainnat



§ ADDITIONAL DETAILS



 §.§ Architecture

Our architecture is as follows: the encoder, mapping the raw pixel state to a latent representation, is a 4-layer convolutional neural network with batch normalization <cit.> and leaky ReLU activations. The decoder uses the transposed architecture with ReLU activations. The noise z is sampled from a 2-dimensional gaussian distribution and both the generator Φ(h,z) and the policy π(h,ϕ) are neural networks consisting of 2 fully-connected layers. In practice, a minibatch of n = 256 or 1024 vectors ϕ_1, …, ϕ_n is sampled at each step. The agent randomly choses one ϕ = ϕ_behavior and samples actions from its policy a ∼π(h, ϕ_behavior). Our model parameters are then updated using policy gradient with the REINFORCE estimator and a state-dependent baseline and importance sampling. For each selectivity reward, the term 𝔼_ϕ'[  A(h', h,ϕ')] is estimated as 1n∑_i = 1^n A(h', h,ϕ_i).

In practice, we don't use concatenation of vectors when feeding two vectors as input for a network (like (h, z) for the factor generator or (h, ϕ) for the policy).
For vectors a,b ∈ℝ^n_a × n_b. We use a bilinear operation bil(a, b) = (a_i * b_j)_i ∈ [[n_a]], j ∈ [[n_b]] as in <cit.>.  We observe the bilinear integrated input to more strongly enforce dependence on both vectors; in contrast, our models often ignored one input when using a simple concatenation.

Through our research, we experiment with different outputs for our generator Φ(h,z).  We explored embedding the ϕ-vectors into a hypercube, a hypersphere, a simplex and also a simplex multiplied by the output of a tanh(·) operation on a scalar.



 §.§ First experiment

In the first experiment, figure <ref>, we used a gaussian similarity kernel i.e A(h', h, ϕ) = exp(- ||h' - (h+ϕ)||^2/2 σ^2) with σ = √(dim(h)).
In this experiment only, for clarity of the figure, we only allowed permissible actions in the environment (no no-op action).






§ ADDITIONAL FIGURES



 §.§ Discrete simple case


Here we consider the case where we learn a latent space H of size K, with K factors corresponding to the coordinates of h (h_i,  i∈[k]), and learn K separately parameterized policies π_i(a|h), i∈[k]. We train our model with the selectivity objective, but no autoencoder loss, and find that we correctly recover independently controllable features on a simple environment.

Albeit slower than when jointly training an autoencoder, this shows that the objective we propose is strong enough to provide a learning signal for discovering a disentangled latent representation.


We train such a model on a gridworld MNIST environment, where there are two MNIST digits lecun1998mnist. The two digits can be moved on the grid via 4 directional actions (so there are 8 actions total), the first digit is always odd and the second digit always even, so they are distiguishable. In Figure <ref> we plot each latent feature h_k as a curve, as a function of each ground truth. For example we see that the black feature recovers +x_1, the horizontal position of the first digit, or that the purple feature recovers -y_2, the vertical position of the second digit.





 §.§ Planning and policy inference example in 1-step



This disentangled structure could be used to address many challenging issues in reinforcement learning. We give two examples in figure <ref>: 


  * Model-based predictions: Given an initial state, s_0, and an action sequence a_{0:T-1}, 
we want to predict the resulting state s_T.


  * A simplified deterministic policy inference problem: Given an initial state s_start and a terminal state s_goal, we aim to find a suitable action sequence a_{0:T-1} such that s_goal can be reached from s_start by following it.

Because of the tanh activation on the last layer of Φ(h, z), the different factors of variation dh = h' - h are placed on the vertices of a hypercube of dimension K, and we can think of the the policy inference problem as finding a path in that simpler space, where the starting point is h_start and the goal is h_goal. We believe this could prove to be a much easier problem to solve.





 §.§ Multistep Example


We demonstrate an instance of ICF operating in a 4×4 Mazebase enviroment over five time steps in Figure <ref>. We consistently witness a failure of mode collapse in our generator Φ and therefore the generator only produces a subset of all possible ϕ-variations.  In Figure <ref>, we observe the ϕ governing the agent's policy π_ϕ appears to correspond to moving two positions down and then to repeatedly toggle the switch.  A random action due to ϵ-greedy led to the agent moving up and off the switch at time step-4. This perturbation is corrected by the policy π_ϕ by moving down in order to return to toggling the relevant switch.







§ VARIATIONAL BOUND AND THE SELECTIVITY



Let us call p(h_t+1 | ϕ_t+1, h_t) =  𝒫^ϕ_h' ,h the probability distribution over final hidden states starting from h and using the policy parametrized by the embedding ϕ.

p(h_t+1 | ϕ_t+1, h_t) = Π_k=1^Kπ_ϕ_t+1( a_t+k-1/K | h_t+k-1/K) p_env(s_t+k/K | a_t+k-1/K, s_t+k-1/K).
where p_env is the transition probability of the environment.


For simplicity, let's refer to h_t as h, h_t+1 as h' and ϕ_t+1 as ϕ.


 §.§ Lower bound on the mutual information

The bound 
    ℐ_p(ϕ, h' | h) ≥sup_θ𝔼_p(ϕ|h)[ 𝒮(h, ϕ)]
 can be proven by using Donsker-Varadhan variational representation of the KL divergence <cit.>:


    pq = sup_T ∈ℒ^∞(q)𝔼_p [ T  ] - log𝔼_q [ e^T ]


For A = e^T and using the identity ℐ_p(X, Y) = 𝔼_p(y)[ p(x|y)p(x)] with X=ϕ|h and Y=h'|h, we have:


    ℐ_p(ϕ, h' | h)    =   𝔼_h'|hsup_A 𝔼_ϕ|h, h'[ log A(h', h, ϕ) ] - log𝔼_φ|h[ A(h', h, φ) ]
       =   𝔼_h'|hsup_A 𝔼_ϕ|h, h'[ log A(h', h, ϕ)/𝔼_φ|h[ A(h', h, φ)]]
       ≥   sup_A 𝔼_ϕ|h𝔼_h'|ϕ, h[ log A(h', h, ϕ)/𝔼_φ|h[ A(h', h, φ)]]
       ≥   sup_θ𝔼_ϕ|h𝔼_h'|ϕ, h[ log A(h', h, ϕ; θ)/𝔼_φ|h[ A(h', h, φ; θ)]]

for parametric A functions.




















As we sample the factors ϕ uniformly, our total objective is then a lower bound on ∑_t ℐ(ϕ_t, h_t | h_t-1) which corresponds here to the directed information <cit.> <cit.> as ϕ_t is sampled independently from ϕ_1:t-1.







§ ADDITIONAL INFORMATION ON THE TRAINING


In our experiments, we use the selectivity objective, an autoencoding loss and an entropy regularization loss ℋ(π_ϕ) for each of the policies π_ϕ. Furthermore, in experiment 4.2 we added the model-based cost ||h' - T(h, ϕ)||^2 with T a learned two layer fully connected neural network.

The selectivity is used to update the parameters of the encoder, factor generator and policy networks.
We use the following equation for computing the gradients


    ∇_θ𝔼_π_θ[ f_θ] =  𝔼_π_θ[ ∇_θ f_θ + f_θ∇_θlogπ_θ]
 
We also use a state dependent baseline V as a control variate to reduce the variance of the REINFORCE estimator.

Furthermore, to be able to train the factor generator efficiently, we train all ϕ sampled in a mini-batch (of size 1024) by importance sampling on the probability ratio of the trajectory under each ϕ

