Question:
represent the article for finding a single-sentence tldr

E.g. Reinforcement learning algorithms, though successful, tend to over-fit to training environments, thereby hampering their application to the real-world. This paper proposes $\text{W}\text{R}^{2}\text{L}$ -- a robust reinforcement learning algorithm with significant robust performance on low and high-dimensional control tasks. Our method formalises robust reinforcement learning as a novel min-max game with a Wasserstein constraint for a correct and convergent solver. Apart from the formulation, we also propose an efficient and scalable solver following a novel zero-order optimisation method that we believe can be useful to numerical optimisation in general. We empirically demonstrate significant gains compared to standard and robust state-of-the-art algorithms on high-dimensional MuJuCo environments Reinforcement learning (RL) has become a standard tool for solving decision-making problems with feedback, and though significant progress has been made, algorithms often over-fit to training environments and fail to generalise across even slight variations of transition dynamics . Robustness to changes in transition dynamics is a crucial component for adaptive and safe RL in real-world environments. Motivated by real-world applications, recent literature has focused on the above problems, proposing a plethora of algorithms for robust decisionmaking (; ;). Most of these techniques borrow from game theory to analyse, typically in a discrete state and actions spaces, worst-case deviations of agents' policies and/or environments, see;;; and references therein. These methods have also been extended to linear function approximators , and deep neural networks showing (modest) improvements in performance gain across a variety of disturbances, e.g., action uncertainties, or dynamical model variations. In this paper, we propose a generic framework for robust reinforcement learning that can cope with both discrete and continuous state and actions spaces. Our algorithm, termed Wasserstein Robust Reinforcement Learning (WR 2 L), aims to find the best policy, where any given policy is judged by the worst-case dynamics amongst all candidate dynamics in a certain set. This set is essentially the average Wasserstein ball around a reference dynamics P 0. The constraints makes the problem well-defined, as searching over arbitrary dynamics can only in non-performing system. The measure of performance is the standard RL objective, the expected return. Both the policy and the dynamics are parameterised; the policy parameters θ k may be the weights of a deep neural network, and the dynamics parameters φ j the parameters of a simulator or differential equation solver. The algorithm performs estimated descent steps in φ space and -after (almost) convergence -performs an update of policy parameters, i.e., in θ space. Since φ j may be high-dimensional, we adapt a zero'th order sampling method based extending to make estimations of gradients, and in order to define the constraint set which φ j is bounded by, we generalise the technique to estimate Hessians (Proposition 2). We emphasise that although access to a simulator with parameterisable dynamics is required, the actual reference dynamics P 0 need not be known explicitly nor learnt by our algorithm. Put another way, we are in the "RL setting", not the "MDP setting" where the transition probability matrix is known a priori. The difference is made obvious, for example, in the fact that we cannot perform dynamic programming, and the determination of a particular probability transition can only be estimated from sampling, not retrieved explicitly. Hence, our algorithm is not model-based in the traditional sense of learning a model to perform planning. We believe our contribution is useful and novel for two main reasons. Firstly, our framing of the robust learning problem is in terms of dynamics uncertainty sets defined by Wasserstein distance. Whilst we are not the first to introduce the Wasserstein distance into the context of MDPs (see, e.g., or), we believe our formulation is amongst the first suitable for application to the demanding application-space we desire, that being, high-dimensional, continuous state and action spaces. Secondly, we believe our solution approach is both novel and effective (as evidenced by experiments below, see Section 5), and does not place a great demand on model or domain knowledge, merely access to a simulator or differentiable equation solver that allows for the parameterisation of dynamics. Furthermore, it is not computationally demanding, in particular, because it does not attempt to build a model of the dynamics, and operations involving matrices are efficiently executable using the Jacobian-vector product facility of automatic differentiation engines. A Markov decision process (MDP) 1 is denoted by M = S, A, P, R, γ, where S ⊆ R d denotes the state space, A ⊆ R n the action space, P: S × A × S → is a state transition probability describing the system's dynamics, R: S × A → R is the reward function measuring the agent's performance, and γ ∈ specifies the degree to which rewards are discounted over time. At each time step t, the agent is in state s t ∈ S and must choose an action a t ∈ A, transitioning it to a new state s t+1 ∼ P (s t+1 |s t, a t), and yielding a reward R(s t, a t). A policy π: S × A → is defined as a probability distribution over state-action pairs, where π(a t |s t) represents the density of selecting action a t in state s t. Upon subsequent interactions with the environment, the agent collects a trajectory τ of state-action pairs. The goal is to determine an optimal policy π by solving: where p π (τ) denotes the trajectory density function, and R Total (τ) the return, that is, the total accumulated reward: with µ 0 (·) denoting the initial state distribution. We make use of the Wasserstein distance to quantify variations from a reference transition density P 0 (·). The latter being a probability distribution, one may consider other divergences, such as Kullback-Leibler (KL) or total variation (TV). Our main intuition for choosing Wasserstein distance is explained below, but we note that it has a number of desirable properties: Firstly, it is symmetric (W p (µ, ν) = W p (ν, µ), a property K-L lacks). Secondly, it is well-defined for measures with different supports (which K-L also lacks). Indeed, the Wasserstein distance is flexible in the forms of the measures that can be compared -discrete, continuous or a mixture. Finally, it takes into account the underlying geometry of the space the distributions are defined on, which can encode valuable information. It is defined as follows: Let X be a metric space with metric d(·, ·). Let C(X) be the space of continuous functions on X and let M(X) be the set of probability measures on X. Let µ, ν ∈ M(X). Let K(µ, ν) be the set of couplings between µ, ν: That is, the set of joint distributions κ ∈ M(X × X) whose marginals agree with µ and ν respectively. Given a metric (serving as a cost function) d(·, ·) for X, the p'th Wasserstein distance W p (µ, ν) for p ≥ 1 between µ and ν is defined as: (in this paper, and mostly for computational convenience, we use p = 2, though other values of p are applicable). The desirable properties of Wasserstein distance aside, our main intuition for choosing it is described thus: Per the definition, constraining the possible dynamics to be within an -Wasserstein ball of a reference dynamics P 0 (·) means constraining it in a certain way. Wasserstein distance has the form mass × distance. If this quantity is constrained to be less than a constant, then if the mass is large, the distance is small, and if the distance is large, the mass is small. Intuitively, when modelling the dynamics of a system, it may be reasonable to concede that there could be a systemic erroror bias -in the model, but that bias should not be too large. It is also reasonable to suppose that occasionally, the behaviour of the system may be wildly different to the model, but that this should be a low-probability event. If the model is frequently wrong by a large amount, then there is no use in it. In a sense, the Wasserstein ball formalises these assumptions. Due to the continuous nature of the state and action spaces considered in this work, we resort to deep neural networks to parameterise policies, which we write as π θ (a t |s t), where θ ∈ R d1 is a set of tunable hyper-parameters to optimise. For instance, these policies can correspond to multilayer perceptrons for MuJoCo environments, or to convolutional neural networks in case of highdimensional states depicted as images. Exact policy details are ultimately application dependent and, consequently, provided in the relevant experiment sections. In principle, one can similarly parameterise dynamics models using deep networks (e.g., LSTMtype models) to provide one or more action-conditioned future state predictions. Though appealing, going down this path led us to agents that discover worst-case transition models which minimise training data but lack any valid physical meaning. For instance, original experiments we conducted on CartPole ended up involving transitions that alter angles without any change in angular velocities. More importantly, these effects became more apparent in high-dimensional settings where the number of potential minimisers increases significantly. It is worth noting that we are not the first to realise such an artifact when attempting to model physics-based dynamics using deep networks. Authors in remedy these problems by introducing Lagrangian mechanics to deep networks, while others (; argue the need to model dynamics given by differential equation structures directly. Though incorporating physics-based priors to deep networks is an important and challenging task that holds the promise of scaling model-based reinforcement learning for efficient solvers, in this paper we rather study an alternative direction focusing on perturbing differential equation solvers and/or simulators with respect to the dynamic specification parameters φ ∈ R d2 . Not only would such a consideration reduce the dimension of parameter spaces representing transition models, but would also guarantee valid dynamics due to the nature of the simulator. Though tackling some of the above problems, such a direction arrives with a new set of challenges related to computing gradients and Hessians of black-box solvers. In Section 4, we develop an efficient and scalable zero-order method for valid and accurate model updates. Unconstrained Loss Function: Having equipped agents with the capability of representing policies and perturbing transition models, we are now ready to present an unconstrained version of WR 2 L's loss function. Borrowing from robust optimal control, we define robust reinforcement learning as an algorithm that learns best-case policies under worst-case transitions: where p φ θ (τ) is a trajectory density function parameterised by both policies and transition models, i.e., θ and φ, respectively: specs vector and diff. solver Under review as a conference paper at ICLR 2020 At this stage, it should be clear that our formulation, though inspired from robust optimal control, is, truthfully, more generic as it allows for parameterised classes of transition models without incorporating additional restrictions on the structure or the scope by which variations are executed 2. Constraints & Complete Problem Definition: Clearly, the problem in Equation 5 is ill-defined due to the arbitrary class of parameterised transitions. To ensure well-behaved optimisation objectives, we next introduce constraints to bound search spaces and ensure convergence to feasible transition models. For a valid constraint set, our method assumes access to samples from a reference dynamics model P 0 (·|s, a), and bounds learnt transitions in an -Wasserstein ball around P 0 (·|s, a), i.e., the set defined as: where ∈ R + is a hyperparameter used to specify the "degree of robustness" in a similar spirit to maximum norm bounds in robust optimal control. It is worth noting, that though we have access to samples from a reference simulator, our setting is by no means restricted to model-based reinforcement learning in an MDP setting. That is, our algorithm operates successfully given only traces from P 0 accompanied with its specification parameters, e.g., pole lengths, torso masses, etc. -a more flexible framework that does not require full model learners. Though defining a better behaved optimisation objective, the set in Equation 6 introduces infinite number of constraints when considering continuous state and/or actions spaces. To remedy this problem, we target a relaxed version that considers a constraint of average Wasserstein distance bounded by a hyperparameter: The sampling (s, a) in the expectation is done as follows: We sample trajectories using reference dynamics P 0 and a policy π that chooses actions uniformly at random (uar). Then (s, a) pairs are sampled uar from those collected trajectories. For a given pair (s, a), W 2 2 (P φ (·|s, a), P 0 (·|s, a)) is approximated through the empirical distribution: we use the state that followed (s, a) in the collected trajectories as a data point. Estimating Wasserstein distance using empirical data is standard, see, e.g., Peyré et al.. One approach which worked well in our experiments, was to assume that the dynamics are given by deterministic functions plus Gaussian noise with diagonal convariance matrices. This makes estimation easier in high dimensions since sampling in each dimension is independent of others, and the total samples needed is a constant factor of the number of dimensions. Gaussian distributions also have closed-form expressions for Wasserstein distance, given in terms of mean and covariance. We thus arrive at WR 2 L's optimisation problem allowing for best policies under worst-case yet bounded transition models: Wasserstein Robust Reinforcement Learning Objective: Our solution alternates between updates of θ and φ, keeping one fixed when updating the other. Fixing dynamics parameters φ, policy parameters θ can be updated by solving, which is the formulation of a standard RL problem. Consequently, one can easily adapt any policy search method for updating policies under fixed dynamical models. As described later in Section 4, we make use of proximal policy optimisation . When updating φ given a set of fixed policy parameters θ, the Wasserstein constraints must be respected. Unfortunately, even with the simplification introduced in Section 3.1 the constraint is still difficult to compute. To alleviate this problem, we propose to approximate the constraint in by its Taylor expansion up to second order. That is, defining W (φ):= E (s,a) W 2 2 (P φ (·|s, a), P 0 (·|s, a)), the above can be approximated around φ 0 by a second-order Taylor as: Recognising that W (φ 0) = 0 (the distance between the same probability densities), and ∇ φ W (φ 0) = 0 since φ 0 minimises W (φ), we can simplify the Hessian approximation by writing:. Substituting our approximation back in the original problem in Equation 8, we reach the following optimisation problem for determining model parameter given fixed policies: where is the Hessian of the expected squared 2-Wasserstein distance evaluated at φ 0. Optimisation problems with quadratic constraints can be efficiently solved using interior-point methods. To do so, one typically approximates the loss with a first-order expansion and determines a closed-form solution. Consider a pair of parameters θ [k] and φ [j] (which will correspond to parameters of the j'th inner loop of the k'th outer loop in the algorithm we present). To find φ [j+1], we solve: It is easy to show that a minimiser to the above equation can derived in a closed-form as: with g [k,j] denoting the gradient 3 evaluated at θ [k] and φ [j], i.e., g Generic Algorithm: Having described the two main steps needed for updating policies and models, we now summarise these findings in the pseudo-code in Algorithm 1. As the Hessian 4 of the Wasserstein distance is evaluated based on reference dynamics and any policy π, we pass it, along with and φ 0 as inputs. Then Algorithms 1 operates in a descent-ascent fashion in two main phases. In the first, lines 5 to 10 in Algorithm 1, dynamics parameters are updated using the closed-form solution in Equation 10, while ensuring that learning rates abide by a step size condition (we used the Wolfe conditions , though it can be some other method). With this, the second phase (line 11) utilises any state-of-the-art reinforcement learning method to adapt policy parameters generating θ [k+1]. Regarding the termination condition for the inner loop, we leave this as a decision for the user. It could be, for example, a large finite time-out, or the norm of the gradient g [k,j] being below a threshold, or whichever happens first., and j ← 0 Phase I: Update model parameter while fixing the policy: while termination condition not met do 7: Compute descent direction for the model parameters as given by Equation 10: Update candidate solution, while satisfying step size conditions (see discussion below) on the learning rate α: end while 10: Perform model update setting 11: Phase II: Update policy given new model parameters: Use any standard reinforcement learning algorithm for ascending in the gradient direction, e.g.,, with β [k] a learning rate. 13: end for Consider a simulator (or differential equation solver) S φ for which the dynamics are parameterised by a real vector φ, and for which we can execute steps of a trajectory (i.e., the simulator takes as input an action a and gives back a successor state and reward). For generating novel physics-grounded transitions, one can simply alter φ and execute the instruction in S φ from some a state s ∈ S, while applying an action a ∈ A. Not only does this ensure valid (under mechanics) transitions, but also promises scalability as specification parameters typically reside in lower dimensional spaces compared to the number of tuneable weights when using deep networks as transition models. Gradient Estimation: Recalling the update rule in Phase I of Algorithm 1, we realise the need for, estimating the gradient of the loss function with respect to the vector specifying the dynamics of the environment, i.e., g at each iteration of the inner-loop j. Handling simulators as black-box models, we estimate the gradients by sampling from a Gaussian distribution with mean 0 and σ 2 I co-variance matrix. Our choice for such estimates is not arbitrary but rather theoretically grounded as one can easily prove the following proposition: Proposition 1 (Zero-Order Gradient Estimate). For a fixed θ and φ, the gradient can be computed as: Hessian Estimation: Having derived a zero-order gradient estimator, we now generalise these notions to a form allowing us to estimate the Hessian. It is also worth reminding the reader that such a Hessian estimator needs to be performed one time only before executing the instructions in Algorithm 1 (i.e., H 0 is passed as an input). Precisely, we prove the following proposition: Proposition 2 (Zero-Order Hessian Estimate). The hessian of the Wasserstein distance around φ 0 can be estimated based on function evaluations. Recalling that, and defining W (s,a) (φ):= W 2 2 (P φ (·|s, a), P 0 (·|s, a)), we prove: Proofs of these propositions are given in Appendix A. They allow for a procedure where gradient and Hessian estimates can be simply based on simulator value evaluations while perturbing φ and φ 0. It is important to note that in order to apply the above, we are required to be able to evaluate An empirical estimate of the p-Wasserstein distance between two measures µ and ν can be performed by computing the p-Wasserstein distance between the empirical distributions evaluated at sampled data. That is, one can approximation µ by µ n = 1 n n i=1 δ xi where x i are identically and independently distributed according to µ. Approximating ν n similarly, we then realise that We evaluate WR 2 L on a variety of continuous control benchmarks from the MuJoCo environment. Dynamics in our benchmarks were parameterised by variables defining physical behaviour, e.g., density of the robot's torso, friction of the ground, and so on. We consider both low and high dimensional dynamics and demonstrate that our algorithm outperforms state-of-the-art from both standard and robust reinforcement learning. We are chiefly interested in policy generalisation across environments with varying dynamics, which we measure using average test returns on novel systems. The comparison against standard reinforcement learning algorithms allows us to understand whether lack of robustness is a critical challenge for sequential decision making, while comparisons against robust algorithms test if we outperform state-of-the-art that considered a similar setting to ours. From standard algorithms, we compare against proximal policy optimisation (PPO) , and trust region policy optimisation (TRPO) (b); an algorithm based on natural actor-crtic . From robust algorithms, we demonstrate how WR 2 L favours against robust adversarial reinforcement learning (RARL) , and action-perturbed Markov decision processes (PR-MDP) proposed in . Due to space constraints, the are presented in Appendix B.2. It is worth noting that we attempted to include deep deterministic policy gradients (DDPG) in our comparisons. Results including DDPG were, however, omitted as it failed to show any significant robustness performance even on relatively simple systems, such as the inverted pendulum; see reported in Appendix B.3. During initial trials, we also performed experiments parameterising models using deep neural networks. Results demonstrated that these models, though minimising training data error, fail to provide valid physics-grounded dynamics. For instance, we arrived at inverted pendula models that vary pole angles without exerting any angular speeds. This problem became even more apparent in high-dimensional systems, e.g., Hopper, Walker, etc due to the increased number of possible minima. As such, presented in this section make use of our zero-order method that can be regarded as a scalable alternative for robust solutions. We evaluate our method both in low and high-dimensional MuJuCo tasks . We consider a variety of systems including CartPole, Hopper, and Walker2D; all of which require direct joint-torque control. Keeping with the generality of our method, we utilise these frameworks as-is with no additional alterations, that is, we use the exact setting of these benchmarks as that shipped with OpenAI gym without any reward shaping, state-space augmentation, feature extraction, or any other modifications of-that-sort. Details are given in section B. Due to space constraints, for one-dimensional parameter variations are given in Appendix B.2, where it can be seen that WR 2 L outperforms both robust and non-robust algorithms when onedimensional simulator variations are considered. Figure 1 shows for dynamics variations along two dimensions. Here again, our methods demonstrates considerable robustness. The fourth column, "PPO mean", refers to experiments where PPO is trained on a dynamics sampled uniformly at random from the Wasserstein constraint set. It displayes more robustness than when trained on just the reference dynamics, however, as can be seen from Fig. 2, our method performs noticably better in high dimensions, which is the main strength of our algorithm. Results with High-Dimensional Model Variation: Though above demonstrate robustness, an argument against a min-max objective can be made especially when only considering lowdimensional changes in the simulator. Namely, one can argue the need for such an objective as opposed to simply sampling a set of systems and determining policies performing-well on average similar to the approach proposed in . A counter-argument to the above is that a gradient-based optimisation scheme is more efficient than a sampling-based one when high-dimensional changes are considered. In other words, a sampling procedure is hardly applicable when more than a few parameters are altered, while WR 2 L can remain suitable. To assess these claims, we conducted two additional experiments on the Hopper and HalfCheetah benchmarks. In the first, we trained robustly while changing friction and torso densities, and tested on 1,000 systems generated by varying all 11 dimensions of the Hopper dynamics, and 21 dimensions of the HalfCheetah system. The histogram Figures 2(b) and (f) demonstrate that the empirical densities of the average test returns are mostly centered around 3,000 for the Hopper, and around 4,500 for the Cheetah, which improves that of PPO trained on reference (Figures 2(a) and (e)) with return masses mostly accumulated at around 1,000 in the case of the Hopper and almost equally distributed when considering HalfCheetah. Such improvements, however, can be an artifact of the careful choice of the low-dimensional degrees of freedom allowed to be modified during Phase I of Algorithm 1. To get further insights, Figures 2(c) and (g) demonstrate the effectiveness of our method trained and tested while allowing to tune all 11 dimensional parameters of the Hopper sim- ulator, and the 21 dimensions of the HalfCheetah. Indeed, our are in accordance with these of the previous experiment depicting that most of the test returns' mass remains around 3,000 for the Hopper, and improves to accumulate around 4,500 for the HalfCheetah. Interestingly, however, our algorithm is now capable of acquiring higher returns on all systems 6 since it is allowed to alter all parameters defining the simulator. As such, we conclude that WR 2 L outperforms others when high-dimensional simulator variations are considered. In Figures 2(d) and (h), we see the for PPO trained with dynamics sampled uar from the Wasserstein constraint set. We see that although in the two-dimensional variation case this training method worked well (see Figures 1(d), (h), (l)), it does not scale well to high dimensions, and our method does better. Previous work on robust MDPs, e.g., , whilst valuable in its own right, is not sufficient for the RL setting due to the need in the latter case to give efficient solutions for large state and action spaces, and the fact that the dynamics are not known a priori. Closer to our own work, approaches the robustness problem by training on an ensemble of dynamics in order to be deployed on a target environment. The algorithm introduced, Ensemble Policy Optimisation (EPOpt), alternates between two phases: (i) given a distribution over dynamics for which simulators (or models) are available (the source domain), train a policy that performs well for the whole distribution; (ii) gather data from the deployed environment (target domain) to adapt the distribution. The objective is not max-min, but a softer variation defined by conditional value-at-risk (CVaR). The algorithm samples a set of dynamics {φ k} from a distribution over dynamics P ψ, and for each dynamics φ k, it samples a trajectory using the current policy parameter θ i. It then selects the worst performing -fraction of the trajectories to use to update the policy parameter. Clearly this process bears some resemblance to our algorithm, but there is a crucial difference: our algorithm takes descent steps in the φ space. The difference if important when the dynamics parameters sit in a high-dimensional space, since in that case, optimisationfrom-sampling could demand a considerable number of samples. In any case, our experiments demonstrate our algorithm performs well even in these high dimensions. We note that we were were unable to find the code for this paper, and did not attempt to implement it ourselves. The CVaR criterion is also adopted in , in which, rather than sampling trajectories and finding a quantile in terms of performance, two policies are trained simultaneously: a "protagonist" which aims to optimise performance, and an adversary which aims to disrupt the protagonist. The protagonist and adversary train alternatively, with one being fixed whilst the other adapts. We made comparisons against this algorithm in our experiments. More recently, studies robustness with respect to action perturbations. There are two forms of perturbation addressed: (i) Probabilistic Action Robust MDP (PR-MDP), and (ii) Noisy Action Robust MDP (NR-MDP). In PR-MDP, when an action is taken by an agent, with probability α, a different, possibly adversarial action is taken instead. In NR-MDP, when an action is taken, a perturbation is added to the action itself. and , the algorithm is suitable for applying deep neural networks, and the paper reports experiments on InvertedPendulum, Hopper, Walker2d and Humanoid. We tested against PR-MDP in some of our experiments, and found it to be lacking in robustness (see Section 5). a non-stationary Markov Decision Process model is considered, where the dynamics can change from one time step to another. The constraint is based on Wasserstein distance, specifically, the Wasserstein distance between dynamics at time t and t is bounded by L|t − t |, i.e., is L-Lipschitz with respect to time, for some constant L. They approach the problem by treating nature as an adversary and implement a Minimax algorithm. The basis of their algorithm is that due to the fact that the dynamics changes slowly (due to the Lipschitz constraint), a planning algorithm can project into the future the scope of possible future dynamics and plan for the worst. The ing algorithm, known as Risk Averse Tree Search, isas the name implies -a tree search algorithm. It operates on a sequence "snapshots" of the evolving MDP, which are instances of the MDP at points in time. The algorithm is tested on small grid world, and does not appear to be readily extendible to the continuous state and action scenarios our algorithm addresses. To summarise, our paper uses the Wasserstein distance for quantifying variations in possible dynamics, in common with , but is suited to applying deep neural networks for continuous state and action spaces. Our algorithm does not require a full dynamics available to it, merely a parameterisable dynamics. It competes well with the above papers, and operates well for high dimensional problems, as evidenced by the experiments. In this paper, we proposed a robust reinforcement learning algorithm capable of outperforming others in terms of test returns on unseen dynamics. The algorithm makes use of Wasserstein constraints for policies generalising across varying domains, and considers a zero-order method for scalable solutions. Empirically, we demonstrated superior performance against state-of-the-art from both standard and robust reinforcement learning on low and high-dimensional MuJuCo environments. In future work, we aim to consider robustness in terms of other components of MDPs, e.g., state representations, reward functions, and others. Furthermore, we will implement WR 2 L on real hardware, considering sim-to-real experiments. -Sub-Case III when indices are all distinct: We have Diagonal Elements Conclusion: Using the above we conclude that • Off-Diagonal Elements (i.e., when i = j): The above analysis is now repeated for computing the expectation of the off-diagonal elements of matrix B. Similarly, this can also be split into three sub-cases depending on indices: -Sub-Case III when indices are all distinct: We have Off-Diagonal Elements Conclusion: Using the above and due to the symmetric properties of H, we conclude that Finally, analysing c, one can realise that Substituting the above back in the original approximation in Equation 11, and using the linearity of the expectation we can easily achieve the statement of the proposition. For clarity, we summarise variables parameterising dynamics in Table 1, and detail specifics next. CartPole: The goal of this classic control benchmark is to balance a pole by driving a cart along a rail. The state space is composed of the position x and velocityẋ of the cart, as well as the angle θ and angular velocities of the poleθ. We consider two termination conditions in our experiments: 1) pole deviates from the upright position beyond a pre-specified threshold, or 2) cart deviates from its zeroth initial position beyond a certain threshold. To conduct robustness experiments, we parameterise the dynamics of the CartPole by the pole length l p, and test by varying l p ∈ [0.3, 3]. Hopper: In this benchmark, the agent is required to control a hopper robot to move forward without falling. The state of the hopper is represented by positions, {x, y, z}, and linear velocities, {ẋ,ẏ,ż}, of the torso in global coordinate, as well as angles, {θ i} 2 i=0, and angular speeds, {θ i} 2 i=0, of the three joints. During training, we exploit an early-stopping scheme if "unhealthy" states of the robot were visited. Parameters characterising dynamics included densities {ρ i} 3 i=0 of the four links, armature {a i} 2 i=0 and damping {ζ i} 2 i=0 of three joints, and the friction coefficient µ g. To test for robustness, we varied both frictions and torso densities leading to significant variations in dynamics. We further conducted additional experiments while varying all 11 dimensional specification parameters. Walker2D: This benchmark is similar to Hopper except that the controlled system is a biped robot with seven bodies and six joints. Dimensions for its dynamics are extended accordingly as reported in Table 1. Here, we again varied the torso density for performing robustness experiments in the range ρ 0 ∈. Halfcheetah: This benchmark is similar to the above except that the controlled system is a twodimensional slice of a three-dimensional cheetah robot. Parameters specifying the simulator consist of 21 dimensions, with 7 representing densities. In our two-dimensional experiments we varied the torso-density and floor friction, while in high-dimensional ones, we allowed the algorithm to control all 21 variables. Our experiments included training and a testing phases. During the training phase we applied Algorithm 1 for determining robust policies while updating transition model parameters according to the min-max formulation. Training was performed independently for each of the algorithms on the relevant benchmarks while ensuring best operating conditions using hyper-parameter values reported elsewhere (; ;). For all benchmarks, policies were represented using parametrised Gaussian distributions with their means given by a neural network and standard derivations by a group of free parameters. The neural network consisted of two hidden layers with 64 units and hyperbolic tangent activations in each of the layers. The final layer exploited linear activation so as to output a real number. Following the actor-critic framework, we also trained a standalone critic network having the same structure as that of the policy. For each policy update, we rolled-out in the current worst-case dynamics to collect a number of transitions. The number associated to these transitions was application-dependent and varied between benchmarks in the range of 5,000 to 10,000. The policy was then optimised (i.e., Phase II of Algorithm 1) using proximal policy optimization with a generalised advantage estimation. To solve the minimisation problem in the inner loop of Algorithm 1, we sampled a number of dynamics from a diagonal Gaussian distribution that is centered at the current worst-case dynamics model. The number of sampled dynamics and the variance of the sampled distributions depended on both the benchmark itself, and well as the dimensions of the dynamics. Gradients needed for model updates were estimated using the in Propositions 7 and 8. Finally, we terminated training when the policy entropy dropped below an application-dependent threshold. When testing, we evaluated policies on unseen dynamics that exhibited simulator variations as described earlier. We measured performance using returns averaged over 20 episodes with a maximum length of 1,000 time steps on testing environments. We note that we used non-discounted mean episode rewards to compute such averages. Figure 3 shows the robustness of policies on various taks. For a fair comparison, we trained two standard policy gradient methods (TRPO (b) and PPO ), and two robust RL algorithms (RARL , PR-MDP ) with the reference dynamics preset by our algorithm. The range of evaluation parameters was intentionally designed to include dynamics outside of the -Wasserstein ball. Clearly, WR 2 L outperforms all baselines in this benchmark. As mentioned in the experiments section of the main paper, we refrained from presenting involving deep deterministic policy gradients (DDPG) due to its lack in robustness even on simple systems, such as the CartPole. In Section 3.2 we presented a closed form solution to the following optimisation problem: which took the form of: == An RL algorithm that learns to be robust to changes in dynamics != We introduce two approaches for efficient and scalable inference in stochastic simulators for which the density cannot be evaluated directly due to, for example, rejection sampling loops.

Question: We propose a novel unsupervised generative model, Elastic-InfoGAN, that learns to disentangle object identity from other low-level aspects in class-imbalanced datasets. We first investigate the issues surrounding the assumptions about uniformity made by InfoGAN, and demonstrate its ineffectiveness to properly disentangle object identity in imbalanced data. Our key idea is to make the discovery of the discrete latent factor of variation invariant to identity-preserving transformations in real images, and use that as the signal to learn the latent distribution's parameters. Experiments on both artificial (MNIST) and real-world (YouTube-Faces) datasets demonstrate the effectiveness of our approach in imbalanced data by: (i) better disentanglement of object identity as a latent factor of variation; and (ii) better approximation of class imbalance in the data, as reflected in the learned parameters of the latent distribution. Generative models aim to model the true data distribution, so that fake samples that seemingly belong to the modeled distribution can be generated (; ;). Recent deep neural network based models such as Generative Adversarial Networks (; ;) and Variational Autoencoders have led to promising in generating realistic samples for high-dimensional and complex data such as images. More advanced models show how to discover disentangled representations;;;; ), in which different latent dimensions can be made to represent independent factors of variation (e.g., pose, identity) in the data (e.g., human faces). InfoGAN ) in particular, tries to learn an unsupervised disentangled representation by maximizing the mutual information between the discrete or continuous latent variables and the corresponding generated samples. For discrete latent factors (e.g., digit identities), it assumes that they are uniformly distributed in the data, and approximates them accordingly using a fixed uniform categorical distribution. Although this assumption holds true for many existing benchmark datasets (e.g.,), real-word data often follows a long-tailed distribution and rarely exhibits perfect balance between the categories. Indeed, applying InfoGAN on imbalanced data can in incoherent groupings, since it is forced to discover potentially non-existent factors that are uniformly distributed in the data; see Fig. 1. In this work, we augment InfoGAN to discover disentangled categorical representations from imbalanced data. Our model, Elastic-InfoGAN, makes two modifications to InfoGAN which are simple and intuitive. First, we remodel the way the latent distribution is used to fetch the latent variables; we lift the assumption of any knowledge about class imbalance, where instead of deciding and fixing them beforehand, we treat the class probabilities as learnable parameters of the optimization process. To enable the flow of gradients back to the class probabilities, we employ the Gumbel-Softmax distribution , which acts as a proxy for the categorical distribution, generating differentiable samples having properties similar to that of categorical samples. Second, we enforce our network to assign the same latent category for an image I and its transformed image I, which induces the discovered latent factors to be invariant to identity-preserving transformations like illumination, translation, rotation, and scale changes. Although there are multiple meaningful ways to partition unlabeled data-e.g., with digits, one partitioning could be based Samples generated with an InfoGAN model learned with a fixed uniform categorical distribution Cat(K = 10, p = 0.1) on balanced and imbalanced data, respectively. Each row corresponds to a different learned latent category. (Right): Samples generated with Elastic-InfoGAN using its automatically learned latent categorical distribution. Although InfoGAN discovers digit identities in the balanced data, it produces redundant/incoherent groupings in the imbalanced data. In contrast, our model is able to discover digit identities in the imbalanced data. on identity, whereas another could be based on stroke width-we aim to discover the partitioning that groups objects according to a high-level factor like identity while being invariant to low-level "nuisance" factors like lighting, pose, and scale changes. Such partitionings focusing on object identity are more likely to be useful for downstream visual recognition applications (e.g., semi-supervised object recognition). In sum, our modifications to InfoGAN lead to better disentanglement and categorical grouping of the data (Fig. 1), while at the same time enabling the discovery of the original imbalance through the learned probability parameters of the Gumbel softmax distribution. Importantly, these modifications do not impede InfoGAN's ability to jointly model both continuous and discrete factors in either balanced or imbalanced data scenarios. Our contributions can be summarized as follows: To our knowledge, our work is the first to tackle the problem of unsupervised generative modeling of categorical disentangled representations in imbalanced data. We show qualitatively and quantitatively our superiority in comparison to Info-GAN and other relevant baselines. Our work takes a step forward in the direction of modeling real data distributions, by not only explaining what modes of a factor of variation are present in the data, but also discovering their respective proportions. Disentangled representation learning Learning disentangled representations of the data has a vast literature . InfoGAN ) is one of the most popular unsupervised GAN based disentanglement methods, which learns disentanglement by maximizing the mutual information between the latent codes and generated images. It has shown promising for discovering meaningful latent factors in balanced datasets like MNIST , CelebA , and SVHN . The recent method of JointVAE extends beta-VAE by jointly modeling both continuous and discrete factors, using Gumbel-Softmax sampling. However, both InfoGAN and JointVAE assume uniformly distributed data, and hence fail to be equally effective in imbalanced data, evident by Fig. 1 and our experiments. Our work proposes modifications to InfoGAN to enable it to discover meaningful latent factors in imbalanced data. Learning from imbalanced data Real world data have a long-tailed distribution , which can impede learning, since the model can get biased towards the dominant categories. To alleviate this issue, researchers have proposed re-sampling (; ; ; ;) and class reweighting techniques (; ; ;) to oversample rare classes and down-weight dominant classes. These methods have shown to be effective for the supervised setting, in which the class distributions are known a priori. There are Figure 2: Elastic-InfoGAN takes a sampled categorical code from a Gumbel-Softmax distribution and a noise vector to generate fake samples. Apart from the original InfoGAN ) loss functions, we have two additional constraints: We take real images x and create a transformed version x using identity-preserving operations (e.g., small rotation), and force their inferred latent code distributions to be close; We also constrain their entropy to be low. The use of differentiable latent variables from the Gumbel-Softmax enables gradients to flow back to the class probabilities to update them. also unsupervised clustering methods that deal with imbalanced data in unknown class distributions (e.g., ;). Our model works in the same unsupervised setting; however, unlike these methods, we propose an unsupervised generative model method that learns to disentangle latent categorical factors in imbalanced data. Leveraging data augmentation for unsupervised image grouping Some works (; ; ;) use data augmentation for image transformation invariant unsupervised clustering or representation learning. The main idea is to maximize the mutual information or similarity between the features of an image and its corresponding transformed image. However, unlike our approach, these methods do not target imbalanced data and do not perform generative modeling. Let X = {x 1, x 2, . . ., x N} be a dataset of N unlabeled images from k different classes. No knowledge about the nature of class imbalance is known beforehand. Our goal is twofold: (i) learn a generative model G which can learn to disentangle object category from other aspects (e.g., digits in MNIST , face identity in YouTube-Faces ); (ii) recover the unknown true class imbalance distribution via the generative modeling process. In the following, we first briefly discuss InfoGAN ), which addressed this problem for the balanced setting. We then explain how InfoGAN can be extended to the scenario of imbalanced data. Learning disentangled representations using the GAN framework was introduced in InfoGAN ). The intuition is for generated samples to retain the information about latent variables, and consequently for latent variables to gain control over certain aspects of the generated image. In this way, different types of latent variables (e.g., discrete categorical vs. continuous) can control properties like discrete (e.g., digit identity) or continuous (e.g., digit rotation) variations in the generated images. Formally, InfoGAN does this by maximizing the mutual information between the latent code c and the generated samples G(z, c), where z ∼ P noise (z) and G is the generator network. The mutual information I(c, G(c, z)) can then be used as a regularizer in the standard GAN training objective. Computing I(c, G(c, z)) however, requires P (c|x), which is intractable and hard to compute. The authors circumvent this by using a lower bound of I(c, G(c, z)), which can approximate P (c|x) via a neural network based auxiliary distribution Q(c|x). The training objective hence becomes: Figure 3: Different ways for unsupervised learning based methods to group unlabeled data; based on rotation (left) vs. digit identity (right). Here, we show two different groups for each grouping. where D is the discriminator network, and H(c) is the entropy of the latent code distribution. Training with this objective in latent codes c having control over the different factors of variation in the generated images G(z, c). To model discrete variations in the data, InfoGAN employs nondifferentiable samples from a uniform categorical distribution with fixed class probabilities; i.e., c ∼ Cat(K = k, p = 1/k) where k is the number of discrete categories to be discovered. As shown in Fig. 1, applying InfoGAN to an imbalanced dataset in suboptimal disentanglement, since the uniform prior assumption does not match the actual ground-truth data distribution of the discrete factor (e.g., digit identity). To address this, we propose two augmentations to InfoGAN. The first is to enable learning of the latent distribution's parameters (class probabilities), which requires gradients to be backpropagated through latent code samples c, and the second is to enforce identity-preserving transformation invariance in the learned latent variables so that the ing disentanglement favors groups that coincide with object identities. Learning the prior distribution To learn the prior distribution, we replace the fixed categorical distribution in InfoGAN with the Gumbel-Softmax distribution , which enables sampling of differentiable samples. The continuous Gumbel-Softmax distribution can be smoothly annealed into a categorical distribution. Specifically, if p 1, p 2..., p k are the class probabilities, then sampling of a k-dimensional vector c can be done in a differentiable way: Here g i, g j are samples drawn from Gumbel, and τ (softmax temperature) controls the degree to which samples from Gumbel-Softmax resemble the categorical distribution. Low values of τ make the samples possess properties close to that of a one-hot sample. In theory, InfoGAN's behavior in the class balanced setting (Fig. 1 left) can be replicated in the imbalanced case (where grouping becomes incoherent, Fig. 1 center), by simply replacing the fixed uniform categorical distribution with Gumbel-Softmax with learnable class probabilities p i's; i.e. gradients can flow back to update the class probabilities (which are uniformly initialized) to match the true class imbalance. And once the true imbalance gets reflected in the class probabilities, the possibility of proper categorical disentanglement (Fig. 1 right) becomes feasible. Empirically, however, this ideal behavior is not observed in a consistent manner. As shown in Fig. 3 (left), unsupervised grouping can focus on non-categorical attributes such as rotation of the digit. Although this is one valid way to group unlabeled data, our goal in this work is to prefer groupings that correspond to class identity as in Fig. 3 (right). Learning object identities To capture object identity as the factor of variation, we make another modification to InfoGAN. Specifically, to make the model focus on high level object identity and be invariant to low level factors like rotation, thickness, illumination, etc., we explicitly create these identity-preserving transformations on real images, and enforce the latent prediction Q(c|x) to be invariant to these transformations. Note that such transformations (aka data augmentations) are standard for learning invariant representations for visual recognition tasks. Formally, for any real image x ∼ P data (x), we apply a set of transformations δ to obtain a transformed image x = δ(x). It is important to point out that these transformations are not learned over the optimization process. Instead we use fixed simple transformations which guarantee that the human defined object identity label for the original image x and the transformed image x image remain the same. For example, the digit identity of a'one' from MNIST will remain the same if a transformation of rotation (±10 degree) is applied. Similarly, a face identity will remain the same upon horizontal flipping. We hence formulate our transformation constraint loss function: where d(·) is a distance metric (e.g., cosine distance), and Q(c x |x), Q(c x |x), are the latent code predictions for real image x and transformed image x, respectively. Note that ideally Q(c|x), for either x ∼ P data (x) or x ∼ P g (G), should have low entropy (peaky class distribution) for proper inference about the latent object category. Eq. 2 automatically enforces a peaky class distribution for Q(c|x) for x ∼ P g (G), because the sampled input latent code c from Gumbel-Softmax is peaky. For x ∼ P data (x) though, Eq. 4 alone isn't sufficient as it can be optimized in a sub-optimal manner (e.g., if c x ≈ c x, but both have high entropy). We hence add an additional entropy loss which forces c x and c x to have low entropy (s) class distributions: The losses L trans and L ent, along with Gumble-Softmax, constitute our overall training objective: V Inf oGAN plays the role of generating realistic images and associating the latent variables to correspond to some factor of variation in the data, while the addition of L trans will push the discovered factor of variation to be close to object identity. Finally, L ent's objective is to ensure Q behaves similarly for real and fake image distributions. The latent codes sampled from Gumbel-softmax, generated fake images, and losses operating on fake images are all functions of class probabilities p i's too. Thus, during the minimization phase of Eqn. 6, the gradients are used to optimize the class probabilities along with G and Q in the backward pass. In this section, we perform quantitative and qualitative analyses to demonstrate the advantage of Elastic-InfoGAN in discovering categorical disentanglement for imbalanced datasets. We use: MNIST and YouTube-Faces . MNIST is by default a balanced dataset with 70k images, with a similar number of training samples for each of 10 classes. We artificially introduce imbalance over 50 random splits (max imbalance ratio 10:1 between the largest and smallest class). YouTube-Faces is a real world imbalanced video dataset with varying number of training samples (frames) for the 40 face identity classes (as used in). The smallest/largest class has 53/695 images, with a total of 10,066 tightly-cropped face images. All are reported over the average of: (i) 50 runs (over 50 random imbalances) for MNIST, (ii) 5 runs over the same imbalanced dataset for YouTube-Faces. We use MNIST to provide a proof-of-concept of our approach. For example, one of the ways in which different'ones' in MNIST vary is rotation, which can be used as a factor (as opposed to object identity) to group data in imbalanced cases (recall Fig. 3 left). Thus, using rotation as a transformation in L trans should alleviate this problem. We ultimately care most about the YouTube-Faces since it is more representative of real world data, both in terms of challenging visual variations (e.g., facial pose, scale, expression, and lighting changes) a well as inherent class imbalance. For this reason, the effect of augmentations in L trans will be more reflective of how well our model can work in real world data. We design different baselines to show the importance of having learnable priors for different latent variables and applying our transformation constraints. • Uniform InfoGAN ): This is the original InfoGAN with fixed and uniform categorical distribution. • Ground-truth InfoGAN: This is InfoGAN with a fixed, but imbalanced categorical distribution where the class probabilities reflect the ground-truth class imbalance. • Ground-truth InfoGAN + Transformation constraint: Similar to the previous baseline but with our data transformation constraint (L trans). • Gumbel-softmax: In this case, InfoGAN does not have a fixed prior for the latent variables. Instead, the priors are learned using the Gumbel-softmax technique . • Gumbel-softmax + Transformation constraint: Apart from having a learnable prior we also apply our transformation constraint (L trans). This is a variant of our final approach that does not have the entropy loss (L ent). • Gumbel-softmax + Transformation constraint + Entropy Loss (Elastic-InfoGAN): This is our final model with all the losses, L trans and L ent, in addition to V Inf oGAN (D, G, Q). • JointVAE : We also include this VAE based baseline, which performs joint modeling of disentangled discrete and continuous factors. Our evaluation should capture: how well we learn class-specific disentanglement for the imbalanced dataset, and recover the ground-truth class distribution of the imbalanced dataset. To capture these aspects, we apply three evaluation metrics: • Average Entropy (ENT): Evaluates two properties: (i) whether the images generated for a given categorical code belong to the same ground-truth class i.e., whether the ground-truth class histogram for images generated for each categorical code has a low entropy; (ii) whether each ground-truth class is associated with a single unique categorical code. We generate 1000 images for each of the k latent categorical codes, compute class histograms using a pre-trained classifier 2 to get a k × k matrix (where rows index latent categories and columns index ground-truth categories). We report the average entropy across the rows (tests (i)) and columns (tests (ii)). • Normalized Mutual Information (NMI) : We treat our latent category assignments of the fake images (we generate 1000 fake images for each categorical code) as one clustering, and the category assignments of the fake images by the pre-trained classifier as another clustering. NMI measures the correlation between the two clusterings. The value of NMI will vary between 0 to 1; higher the NMI, stronger the correlation. • Root Mean Square Error (RMSE) between predicted and actual class distributions: measures the accuracy of approximating the true class distribution of the imbalanced dataset. Since the learned latent distribution may not be aligned to the ground-truth distribution (e.g., the first dimension for the learned distribution might capture 9's in MNIST whereas the first dimension for the groundtruth distribution may be for 0's), we need a way to align the two. For this, we use the pre-trained classifier to classify the generated images for a latent variable and assign the variable to the most frequent class. If more than one latent variable is assigned to the same class, then their priors are added before computing its distance with the known prior of the ground-truth class. We first evaluate disentanglement quality as measured by NMI and average entropy (ENT); see Figure 4: Representative image generations on a random imbalanced MNIST split. Each row corresponds to a learned latent variable. Our approach generates inconsistent images in only row 2 whereas Uniform InfoGAN does so in rows 1,2,6,8 and JointVAE does so in rows 3,5,6,7,9,10. particular, our full model obtains significant boosts of 0.101 and 0.104 in NMI, and -0.222 and -0.305 in ENT compared to the Uniform InfoGAN baseline for MNIST and YouTube-Faces, respectively. The boost is even more significant when compared to JointVAE: 0.1977, 0.3380 in NMI, and -0.4658, -0.9963 in ENT for MNIST and YouTube-Faces, respectively. This again is a of the assumption of a uniform categorical prior by JointVAE, along with poorer quality generations. We see that our transformation constraint generally improves the performance for both when the ground-truth prior is known (Ground-truth InfoGAN vs. Ground-truth InfoGAN + Transformation constraint) as well as when the prior is learned (Gumbel-softmax vs. Gumbel-softmax + Transformation constraint). This shows that enforcing the network to learn groupings that are invariant to identity-preserving transformations helps it to learn a disentangled representation in which the latent dimensions correspond more closely to identity-based classes. Also, learning the prior using the Gumbel-softmax leads to better categorical disentanglement than fixed uniform priors, which demonstrates the importance of learning the prior distribution in imbalanced data. Overall, our approach using Gumbel-softmax to learn the latent prior distribution together with our transformation constraint works better than applying them individually, which demonstrates their complementarity. Interestingly, using a fixed ground-truth prior (Ground-truth InfoGAN) does not in better disentanglement than learning the prior (Gumbel-softmax). This requires further investigation, but we hypothesis that having a rigid prior makes optimization more difficult compared to allowing the network to converge to a distribution on its own, as there are multiple losses that need to be simultaneously optimized. Finally, in Table 2, we evaluate how well the Gumbel-softmax can recover the ground-truth prior distribution. For this, we compute the RMSE between the learned prior distribution and ground- truth prior distribution. Our full model (transformation constraint + entropy loss) produces the best estimate of the true class imbalance for both datasets, as evident through lowest RMSE. Our improvement over the Gumbel-Softmax baseline indicates the importance of our tranformation L trans and entropy L ent losses in approximating the class imbalance. We next qualitatively evaluate the disentanglement achieved by our approach. Figs. 4, 5, and 7 show for MNIST and YouTube-Faces. Overall, Elastic-InfoGAN generates more consistent images for each latent code compared to Uniform InfoGAN and JointVAE. For example, in Fig. 4, ElasticInfoGAN only generates inconsistent images in the second row whereas the baseline approaches generate inconsistent images in several rows. Similarly, in Fig. 7, Elastic-InfoGAN generates faces of the same person corresponding to a latent variable more consistently than the baselines. Both Uniform InfoGAN and JointVAE on the other hand tend to mix up identities within the same categorical code because they incorrectly assume a prior uniform distribution. Finally, we demonstrate that Elastic-InfoGAN does not impede modeling of continuous factors in the imbalanced setting. Specifically, one can augment the input with continuous latent codes (e.g. r1, r2 ∼ Unif(-1, 1)) along with the existing categorical and noise vectors. In Fig. 6, we show the of continuous code interpolation; we can see that each of the two continuous codes largely captures a particular continuous factor (stroke width on left, and digit rotation on the right). In this work, we proposed a new unsupervised generative model that learns categorical disentanglement in imbalanced data. Our model learns the class distribution of the imbalanced data and enforces invariance to be learned in the discrete latent variables. Our demonstrate superior performance over alternative baselines. We hope this work will motivate other researchers to pursue this interesting research direction in generative modeling of imbalanced data. For MNIST, we operate on the original 28x28 image size, with 10-dimensional categorical code to represent 10 digit categories. For YouTube-Faces, we crop the faces using bounding box annotations provided, and then resize them to 64x64 resolution, and use a 40-dimensional categorical code to represent 40 face identities (first 40 categories sorted in alphabetical manner), as done in. Pre-trained classification architecture used for evaluation for MNIST: 2 Conv + 2 FC layers, with max pool and ReLU after every convolutional layer. For YouTube-Faces classification, we fine-tune a ResNet-50 network pretrained on VGGFace2, for face recognition. We set λ 1 = 1 (for L 1), λ 2 = 10 (for L trans), and λ 3 = 1 (for L ent). These hyperparameters were chosen to balance the magnitude of the different loss terms. Finally, one behavior we observe is that if the random initialization of class probabilities is too skewed (only few classes have high probability values), then it becomes very difficult for them to get optimized to the ideal state. We hence initialize them with the uniform distribution, which makes training much more stable. We follow the exact architecture as described in InfoGAN ): The generator network G takes as input a 64 dimensional noise vector z ∼ N and 10 dimensional samples from Gumbel-Softmax distribution. The discriminator D and the latent code prediction network Q share most of the layers except the final fully connected layers. Elastic-InfoGAN architecture for YouTube Faces We operate on cropped face images resized to 64x64 resolution. Our architecture is based on the one proposed in StackGANv2 , where we use its 2-stage version for generating 64x64 resolution images. The input is a 100 dimensional noise vector z ∼ N and 40 dimensional samples (c) from the Gumbel-Softmax distribution. There is an initial fully connected layer which maps the input (concatenation of z and c) to an intermediate feature representation. A series of a combination of upsampling + convolutional (interleaved with batch normalization and Gated Linear Units) increase the spatial resolution of the feature representation, starting from 1024 (feature size: 4 x 4 x 1024) channels to 64 (feature size: 64 x 64 x 64) channels. For the first stage, a convolutional network transforms the feature representation into a 3 channel output, while maintaining the spatial resolution; this serves as the fake image from the first stage. The next stage uses the 64 x 64 x 64 resolution features, forwards it through a network containing residual blocks and convolutional layers, while again maintaining the spatial resolution of 64 x 64. For the second stage, again a convolutional layer maps the ing feature into a 64 x 64 resolution fake image, which is the one used by the model for evaluation purposes. The discriminator networks are identical at both stages. It consists of 4 convolutional layers interleaved with batch normalization and leaky ReLU layers, which serve as the common layers for both the D and Q networks. After that, D has one non-shared convolutional layer which maps the feature representation into a scalar value reflecting the real/fake score. For Q, we have a pair of non-shared convolutional layers which map the feature representation into a 40 dimensional latent code prediction. We employ a similar way of training the generative and discriminative modules as described in. We first update the discriminator based on the real/fake adversarial loss. In the next step, after computing the remaining losses (mutual information + L trans + L ent), we update the generator (G) + latent code predictor (Q) + latent distribution parameters at once. Our optimization process alternates between these two phases. For MNIST, we train all baselines for 200 epochs, with a batch size of 64. For YouTube-Faces, we train until convergence, as measured via qualitative realism of the generated images. We use a batch size of 50. τ = 0.1 when used for sampling from Gumbel-Softmax, which in samples having very low entropy (very close to one hot vectors from a categorical distribution). Here we describe the exact class imbalance used in our experiments. For MNIST, we include below the 50 random imbalances created. For YouTube-Faces, we include the true ground truth class imbalance in the first 40 categories. The imbalances reflect the class frequency. A.2.1 MNIST 147, 0.037, 0.033, 0.143, 0.136, 0.114, 0.057, 0.112, 0.143, 0.078 • 0.061, 0.152, 0.025, 0.19, 0.12, 0.036, 0.092, 0.185, 0.075, 0.064 • 0.173, 0.09, 0.109, 0.145, 0.056, 0.114, 0.075, 0.03, 0.093, 0.116 • 0.079, 0.061, 0.033, 0.139, 0.145, 0.135, 0.057, 0.062, 0.169, 0.121 • 0.053, 0.028, 0.111, 0.142, 0.13, 0.121, 0.107, 0.066, 0.125, 0.118 • 0.072, 0.148, 0.092, 0.081, 0.119, 0.172, 0.05, 0.109, 0.085, 0.073 • 0.084, 0.143, 0.07, 0.082, 0.059, 0.163, 0.156, 0.063, 0.074, 0.105 • 0.062, 0.073, 0.065, 0.183, 0.099, 0.08, 0.05, 0.16, 0.052, 0.177 • 0.139, 0.113, 0.074, 0.06, 0.068, 0.133, 0.142, 0.13, 0.112, 0.03 • 0.046, 0.128, 0.059, 0.112, 0.135, 0.164, 0.142, 0.125, 0.051, 0.037 • 0.107, 0.057, 0.154, 0.122, 0.05, 0.111, 0.032, 0.044, 0.136, 0.187 • 0.129, 0.1, 0.039, 0.112, 0.119, 0.095, 0.047, 0.14, 0.156, 0.064 • 0.146, 0.08, 0.06, 0.072, 0.051, 0.119, 0.176, 0.11, 0.158, 0.028 A.3 DISCUSSION ABOUT EVALUATING PREDICTED CLASS IMBALANCE IN SEC. 4.2 To measure the ability of a generative model to approximate the class imbalance present in the data, we derive a metric in Section 4.2 of the main paper, the of which are presented in Table 2. Even though we do get better as measured by RMSE between the approximated and the original imbalance distribution, we would like to discuss certain flaws associated with this metric. In its current form, we compute the class histogram (using the pre-trained classifier, which classifies each fake image into one of the ground-truth categories) for a latent code and associate the latent code to the most frequent class. If multiple latent codes get associated to the same ground-truth class, there will be ground-truth classes for which the predicted class probability will be zero. This is rarely an issue for MNIST, as it only has 10 ground-truth classes, and thus in most cases both our method and the baselines assign each latent code to a unique ground-truth class. However, for YouTube-Faces, after associating latent codes to the ground truth categories in this manner, roughly 10-13 ground-truth classes (out of 40) get associated with 0 probability for both our approach and the baselines (due to multiple latent codes being associated to the same majority ground-truth class). Our metric therefore may be too strict, especially for difficult settings with many confusing groundtruth categories. The tricky part about evaluating how well the model is approximating the class imbalance is that there are two key aspects that need to be simultaneously measured. Specifically, not only should (i) the raw probability values discovered match the ground-truth class imbalance distribution, but (ii) the class probabilities approximated by the latent codes must correspond to the correct ground-truth classes. For example, if the original data had 80% samples from class A and 20% from class B, the generative model should not only estimate the imbalance as 80%-20%, but the model must associate 80% to class A and 20% to class B (instead of 80% to class B and 20% to class A). Another way to evaluate whether a model is capturing the ground-truth class imbalance could be the FID score, but it's worth noting that a method can still have a good FID score without disentangling the different factors of variations. Given the limitation with our metric on YouTube-Faces, we have also measured the min/max of predicted prior values. For YouTube-Faces, the min/max of predicted and ground-truth priors are: Gumbel-Softmax: Min 2.76748415e-05, Max: 0.0819286481; Ours without L ent: Min 0.00211485, Max: 0.06152404; Ours complete: Min 0.00336615, Max: 0.06798439; and GroundTruth: Min 0.005265, Max: 0.069044. Our full method's min/max more closely matches that of the ground-truth, and the overall ordering of the methods follows that of Table 2 using our RMSE based metric. In sum, we have made an effort to evaluate accurate class imbalance prediction in multiple ways, but it is important to note that this is an area which calls for better metrics to evaluate the model's ability to approximate the class imbalance distribution.

Elastic-InfoGAN is a modification of InfoGAN that learns, without any supervision, disentangled representations in class imbalanced data

Conditional VAE on top of latent spaces of pre-trained generative models that enables  transfer between drastically different domains while preserving locality and semantic alignment.