On the Transfer of Disentangled Representations in Realistic Settings

Learning meaningful representations that disentangle the underlying structure of the data generating process is considered to be of key importance in machine learning . While disentangled representations were found to be useful for diverse tasks such as abstract reasoning and fair classification , their scalability and real-world impact remain questionable . We introduce a new high-resolution dataset with 1M simulated images and over 1,800 annotated real-world images of the same setup . In contrast to previous work , this new dataset exhibits correlations , a complex underlying structure , and allows to evaluate transfer to unseen simulated and realworld settings where the encoder i ) remains in distribution or ii ) is out of distribution . We propose new architectures in order to scale disentangled representation learning to realistic high-resolution settings and conduct a large-scale empirical study of disentangled representations on this dataset . We observe that disentanglement is a good predictor for out-of-distribution ( OOD ) task performance . 1 INTRODUCTION Disentangled representations hold the promise of generalization to unseen scenarios ( Higgins et al. , 2017b ) , increased interpretability ( Adel et al. , 2018 ; Higgins et al. , 2018 ) and faster learning on downstream tasks ( van Steenkiste et al. , 2019 ; Locatello et al. , 2019a ) . However , most of the focus in learning disentangled representations has been on small synthetic datasets whose ground truth factors exhibit perfect independence by design . More realistic settings remain largely unexplored . We hypothesize that this is because real-world scenarios present several challenges that have not been extensively studied to date . Important challenges are scaling ( much higher resolution in observations and factors ) , occlusions , and correlation between factors . Consider , for instance , a robotic arm moving a cube : Here , the robot arm can occlude parts of the cube , and its end-effector position exhibits correlations with the cube ’ s position and orientation , which might be problematic for common disentanglement learners ( Träuble et al. , 2020 ) . Another difficulty is that we typically have only limited access to ground truth labels in the real world , which requires robust frameworks for model selection when no or only weak labels are available . ∗Equal contribution . Correspondence to : < adit @ dtu.dk > , < frederik.traeuble @ tuebingen.mpg.de > . †Work done during an internship at the Max Planck Institute for Intelligent Systems . The goal of this work is to provide a path towards disentangled representation learning in realistic settings . First , we argue that this requires a new dataset that captures the challenges mentioned above . We propose a dataset consisting of simulated observations from a scene where a robotic arm interacts with a cube in a stage ( see Fig . 1 ) . This setting exhibits correlations and occlusions that are typical in real-world robotics . Second , we show how to scale the architecture of disentanglement methods to perform well on this dataset . Third , we extensively analyze the usefulness of disentangled representations in terms of out-of-distribution downstream generalization , both in terms of held-out factors of variation and sim2real transfer . In fact , our dataset is based on the TriFinger robot from Wüthrich et al . ( 2020 ) , which can be built to test the deployment of models in the real world . While the analysis in this paper focuses on the transfer and generalization of predictive models , we hope that our dataset may serve as a benchmark to explore the usefulness of disentangled representations in real-world control tasks . The contributions of this paper can be summarized as follows : • We propose a new dataset for disentangled representation learning , containing 1M simulated high-resolution images from a robotic setup , with seven partly correlated factors of variation . Additionally , we provide a dataset of over 1,800 annotated images from the corresponding real-world setup that can be used for challenging sim2real transfer tasks . These datasets are made publicly available.1 • We propose a new neural architecture to successfully scale VAE-based disentanglement learning approaches to complex datasets . • We conduct a large-scale empirical study on generalization to various transfer scenarios on this challenging dataset . We train 1,080 models using state-of-the-art disentanglement methods and discover that disentanglement is a good predictor for out-of-distribution ( OOD ) performance of downstream tasks . 2 RELATED WORK . Disentanglement methods . Most state-of-the-art disentangled representation learning approaches are based on the framework of variational autoencoders ( VAEs ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) . A ( high-dimensional ) observation x is assumed to be generated according to the latent variable model pθ ( x|z ) p ( z ) where the latent variables z have a fixed prior p ( z ) . The generative model pθ ( x|z ) and the approximate posterior distribution qφ ( z|x ) are typically parameterized by neural networks , which are optimized by maximizing the evidence lower bound ( ELBO ) : LV AE = Eqφ ( z|x ) [ log pθ ( x|z ) ] −DKL ( qφ ( z|x ) ‖p ( z ) ) ≤ log p ( x ) ( 1 ) As the above objective does not enforce any structure on the latent space except for some similarity to p ( z ) , different regularization strategies have been proposed , along with evaluation metrics to gauge the disentanglement of the learned representations ( Higgins et al. , 2017a ; Kim & Mnih , 2018 ; Burgess et al. , 2018 ; Kumar et al. , 2018 ; Chen et al. , 2018 ; Eastwood & Williams , 2018 ) . Recently , Locatello et al . ( 2019b , Theorem 1 ) showed that the purely unsupervised learning of disentangled representations is impossible . This limitation can be overcome without the need for explicitly labeled data by introducing weak labels ( Locatello et al. , 2020 ; Shu et al. , 2019 ) . Ideas related to disentangling the factors of variation date back to the non-linear ICA literature ( Comon , 1994 ; Hyvärinen & Pajunen , 1999 ; Bach & Jordan , 2002 ; Jutten & Karhunen , 2003 ; Hyvarinen & Morioka , 2016 ; Hyvarinen et al. , 2019 ; Gresele et al. , 2019 ) . Recent work combines non-linear ICA with disentanglement ( Khemakhem et al. , 2020 ; Sorrenson et al. , 2020 ; Klindt et al. , 2020 ) . Evaluating disentangled representations . The BetaVAE ( Higgins et al. , 2017a ) and FactorVAE ( Kim & Mnih , 2018 ) scores measure disentanglement by performing an intervention on the factors of variation and predicting which factor was intervened on . The Mutual Information Gap ( MIG ) ( Chen et al. , 2018 ) , Modularity ( Ridgeway & Mozer , 2018 ) , DCI Disentanglement ( Eastwood & Williams , 2018 ) and SAP scores ( Kumar et al. , 2018 ) are based on matrices relating factors of variation and codes ( e.g . pairwise mutual information , feature importance and predictability ) . 1http : //people.tuebingen.mpg.de/ei-datasets/iclr_transfer_paper/robot_ finger_datasets.tar ( 6.18 GB ) Datasets for disentanglement learning . dSprites ( Higgins et al. , 2017a ) , which consists of binary low-resolution 2D images of basic shapes , is one of the most commonly used synthetic datasets for disentanglement learning . Color-dSprites , Noisy-dSprites , and Scream-dSprites are slightly more challenging variants of dSprites . The SmallNORB dataset contains toy images rendered under different lighting conditions , elevations and azimuths ( LeCun et al. , 2004 ) . Cars3D ( Reed et al. , 2015 ) exhibits different car models from Fidler et al . ( 2012 ) under different camera viewpoints . 3dshapes is a popular dataset of simple shapes in a 3D scene ( Kim & Mnih , 2018 ) . Finally , Gondal et al . ( 2019 ) proposed MPI3D , containing images of physical 3D objects with seven factors of variation , such as object color , shape , size and position available in a simulated , simulated and highly realistic rendered simulated variant . Except MPI3D which has over 1M images , the size of the other datasets is limited with only 17 , 568 to 737 , 280 images . All of the above datasets exhibit perfect independence of all factors , the number of possible states is on the order of 1M or less , and due to their static setting they do not allow for dynamic downstream tasks such as reinforcement learning . In addition , except for SmallNORB , the image resolution is limited to 64x64 and there are no occlusions . Other related work . Locatello et al . ( 2020 ) probed the out-of-distribution generalization of downstream tasks trained on disentangled representations . However , these representations are trained on the entire dataset . Generalization and transfer performance especially for representation learning has likewise been studied in Dayan ( 1993 ) ; Muandet et al . ( 2013 ) ; Heinze-Deml & Meinshausen ( 2017 ) ; Rojas-Carulla et al . ( 2018 ) ; Suter et al . ( 2019 ) ; Li et al . ( 2018 ) ; Arjovsky et al . ( 2019 ) ; Krueger et al . ( 2020 ) ; Gowal et al . ( 2020 ) . For the role of disentanglement in causal representation learning we refer to the recent overview by Schölkopf et al . ( 2021 ) . Träuble et al . ( 2020 ) systematically investigated the effects of correlations between factors of variation on disentangled representation learners . Transfer of learned disentangled representations from simulation to the real world has been recently investigated by Gondal et al . ( 2019 ) on the MPI3D dataset , and previously by Higgins et al . ( 2017b ) in the context of reinforcement learning . Sim2real transfer is of major interest in the robotic learning community , because of limited data and supervision in the real world ( Tobin et al. , 2017 ; Rusu et al. , 2017 ; Peng et al. , 2018 ; James et al. , 2019 ; Yan et al. , 2020 ; Andrychowicz et al. , 2020 ) . 3 SCALING DISENTANGLED REPRESENTATIONS TO COMPLEX SCENARIOS . A new challenging dataset . Simulated images in our dataset are derived from the trifinger robot platform introduced by Wüthrich et al . ( 2020 ) . The motivation for choosing this setting is that ( 1 ) it is challenging due to occlusions , correlations , and other difficulties encountered in robotic settings , ( 2 ) it requires modeling of fine details such as tip links at high resolutions , and ( 3 ) it corresponds to a robotic setup , so that learned representations can be used for control and reinforcement learning in simulation and in the real world . The scene comprises a robot finger with three joints that can be controlled to manipulate a cube in a bowl-shaped stage . Fig . 1 shows examples of scenes from our dataset . The data is generated from 7 different factors of variation ( FoV ) listed in Table 1 . Unlike in previous datasets , not all FoVs are independent : The end-effector ( the tip of the finger ) can collide with the floor or the cube , resulting in infeasible combinations of the factors ( see Appendix B.1 ) . We argue that such correlations are a key feature in real-world data that is not present in existing datasets . The high FoV resolution results in approximately 1.52 billion feasible states , but the dataset itself only contains one million of them ( approximately 0.065 % of all possible FoV combinations ) , realistically rendered into 128 × 128 images . Additionally , we recorded an annotated dataset under the same conditions in the real-world setup : we acquired 1,809 camera images from the same viewpoint and recorded the labels of the 7 underlying factors of variation . This dataset can be used for out-of-distribution evaluations , few-shot learning , and testing other sim2real aspects . Model architecture . When scaling disentangled representation learning to more complex datasets , such as the one proposed here , one of the main bottlenecks in current VAE-based approaches is the flexibility of the encoder and decoder networks . In particular , using the architecture from Locatello et al . ( 2019b ) , none of the models we trained correctly captured all factors of variation or yielded high-quality reconstructions . While the increased image resolution already presents a challenge , the main practical issue in our new dataset is the level of detail that needs to be modeled . In particular , we identified the cube rotation and the lower joint position to be the factors of variation that were the hardest to capture . This is likely because these factors only produce relatively small changes in the image and hence the reconstruction error . To overcome these issues , we propose a deeper and wider neural architecture than those commonly used in the disentangled representation learning literature , where the encoder and decoder typically have 4 convolutional and 2 fully-connected layers . Our encoder consists of a convolutional layer , 10 residual blocks , and 2 fully-connected layers . Some residual blocks are followed by 1x1 convolutions that change the number of channels , or by average pooling that downsamples the tensors by a factor of 2 along the spatial dimensions . Each residual block consists of two 3x3 convolutions with a leaky ReLU nonlinearity , and a learnable scalar gating mechanism ( Bachlechner et al. , 2020 ) . Overall , the encoder has 23 convolutional layers and 2 fully connected layers . The decoder mirrors this architecture , with average pooling replaced by bilinear interpolation for upsampling . The total number of parameters is approximately 16.3M . See Appendix A for further implementation details . Experimental setup . We perform a large-scale empirical study on the simulated dataset introduced above by training 1,080 β-VAE models.2 For further experimental details we refer the reader to Appendix A . The hyperparameter sweep is defined as follows : • We train the models using either unsupervised learning or weakly supervised learning ( Locatello et al. , 2020 ) . In the weakly supervised case , a model is trained with pairs of images that differ in k factors of variation . Here we fix k = 1 as it was shown to lead to higher disentanglement by Locatello et al . ( 2020 ) . The dataset therefore consists of 500k pairs of images that differ in only one FoV . • We vary the parameter β in { 1 , 2 , 4 } , and use linear deterministic warm-up ( Bowman et al. , 2015 ; Sønderby et al. , 2016 ) over the first { 0 , 10000 , 50000 } training steps . • The latent space dimensionality is in { 10 , 25 , 50 } . • Half of the models are trained with additive noise in the input image . This choice is mo- tivated by the fact that adding noise to the input of neural networks has been shown to be beneficial for out-of-distribution generalization ( Sietsma & Dow , 1991 ; Bishop , 1995 ) . • Each of the 108 resulting configurations is trained with 10 random seeds . Can we scale up disentanglement learning ? Most of the trained VAEs in our empirical study fully capture all the elements of a scene , correctly model heavy occlusions , and generate detailed , 2Training these models requires approximately 2.8 GPU years on NVIDIA Tesla V100 PCIe . high-quality samples and reconstructions ( see Appendix B.2 ) . From visual inspections such as the latent traversals in Fig . 2 , we observe that many trained models fully disentangle the ground-truth factors of variation . This , however , appears to only be possible in the weakly supervised scenario . The fact that models trained without supervision learn entangled representations is in line with the impossibility result for the unsupervised learning of disentangled representations from Locatello et al . ( 2019b ) . Latent traversals from a selection of models with different degrees of disentanglement are presented in Appendix B.3 . Interestingly , the high-disentanglement models seem to correct for correlations and interpolate infeasible states , i.e . the fingertip traverses through the cube or the floor . Summary : The proposed architecture can scale disentanglement learning to more realistic settings , but a form of weak supervision is necessary to achieve high disentanglement . How useful are common disentanglement metrics in realistic scenarios ? The violin plot in Fig . 3 ( left ) shows that DCI and MIG measure high disentanglement under weak supervision and lower disentanglement in the unsupervised setting . This is consistent with our qualitative conclusion from visual inspection of the models ( Appendix B.3 ) and with the aforementioned impossibility result . Many of the models trained with weak supervision exhibit a very high DCI score ( 29 % of them have > 99 % DCI , some of them up to 99.89 % ) . SAP and Modularity appear to be ineffective at capturing disentanglement in this setting , as also observed by Locatello et al . ( 2019b ) . Finally , note that the BetaVAE and FactorVAE metrics are not straightforward to be evaluated on datasets that do not contain all possible combinations of factor values . According to Fig . 3 ( right ) , DCI and MIG strongly correlate with test accuracy of GBT classifiers predicting the FoVs . In the weakly supervised setting , these metrics are strongly correlated with the ELBO ( positively ) and with the reconstruction loss ( negatively ) . We illustrate these relationships in more detail in Appendix B.4 . Such correlations were also observed by Locatello et al . ( 2020 ) on significantly less complex datasets , and can be exploited for unsupervised model selection : these unsupervised metrics can be used as proxies for disentanglement metrics , which would require fully labeled data . Summary : DCI and MIG appear to be useful disentanglement metrics in realistic scenarios , whereas other metrics seem to fall short of capturing disentanglement or can be difficult to compute . When using weak supervision , we can select disentangled models with unsupervised metrics .

Summary: The paper presents a new, more complex, dataset for the use of disentangled representation learning. The dataset is based on real and simulated images of the trifinger robot platform. There are 7 factors of variation with high-resolution measurements of these factors. The dataset contains over 1 million simulated images and another ~1000 annotated images of a real trifinger robot arm.