summaries: The paper proposes a framework of jointly learning a state and action embedding using the model of the environment, eventually using those embeddings to learn a parameterized control policy using standard policy gradient (PG) methods. Joint learning of state and action embeddings allows us to capture the interactions between actions in different states. The framework proposes to learn an internal (embedding) policy, a state embedding, an inverse function on action embeddings, combining all the parts to form an overall policy. The paper theoretically shows that optimizing the internal policy leads to an optimal overall policy.

summaries: The paper proposes a method to jointly learn: (a) a latent state embedding; (b) a latent action embedding; (c) a state transition model; and (d) an RL policy.  The latent models should allow for better generalization over states and actions, and therefore result in improved learning, particularly for discrete action domains. The method shows improved performance over vanilla policy gradient on a grid-world task, a slot machine task, a recommender system, and half-cheetah locomotion.

summaries: Learning on environments with large state-action spaces can be difficult. This paper addresses this issue by learning a joint state-action embedding and learn an internal policy(\pi_i) on this embedded state-action space instead of the original state-action space. There are three parts of learning, 1. learning the embedding model that learns mapping from state to state embedding, 2. learning the internal policy, and 3. learning the mapping from action embedding to action space. The authors justify this approach by showing that the overall policy (\pi_o) can be expressed in terms of the internal policy (\pi_i). Furthermore, there is equivalence between the internal state-action-value function and overall state-action-value function and the authors show that updating \pi_i is equivalent to updating \pi_o.

abstractText: While reinforcement learning has achieved considerable successes in recent years, state-of-the-art models are often still limited by the size of state and action spaces. Model-free reinforcement learning approaches use some form of state representations and the latest work has explored embedding techniques for actions, both with the aim of achieving better generalization and applicability. However, these approaches consider only states or actions, ignoring the interaction between them when generating embedded representations. In this work, we propose a new approach for jointly learning embeddings for states and actions that combines aspects of model-free and model-based reinforcement learning, which can be applied in both discrete and continuous domains. Specifically, we use a model of the environment to obtain embeddings for states and actions and present a generic architecture that uses these to learn a policy. In this way, the embedded representations obtained via our approach enable better generalization over both states and actions by capturing similarities in the embedding spaces. Evaluations of our approach on several gaming, robotic control, and recommender systems show it significantly outperforms state-of-the-art models in both discrete/continuous domains with large state/action spaces, thus confirming its efficacy and the overall superior performance.

authors:
Prithviraj Ammanabrolu
Mark O Riedl

title: Playing text-adventure games with graph-based deep reinforcement learning

venue: arXiv preprint arXiv:1812.01628,

title: Scalable methods for computing state similarity in deterministic markov decision processes

venue: In Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Yash Chandak
Georgios Theocharous
James Kostas
Scott Jordan
Philip S. Thomas

title: Learning Action Representations for Reinforcement Learning

venue: 36th International Conference on Machine Learning (ICML),

authors:
Tim de Bruin
Jens Kober
Karl Tuyls
Robert Babuška

title: Integrating state representation learning into deep reinforcement learning

venue: IEEE Robotics and Automation Letters,

authors:
Gabriel Dulac-Arnold
Richard Evans
Hado van Hasselt
Peter Sunehag
Timothy Lillicrap
Jonathan Hunt
Timothy Mann
Theophane Weber
Thomas Degris
Ben Coppin

title: Deep reinforcement learning in large discrete action spaces

venue: arXiv preprint arXiv:1512.07679,

authors:
Richard Evans
Jim Gao

title: Deepmind AI reduces google data centre cooling bill by 40

venue: DeepMind blog,

authors:
Norm Ferns
Prakash Panangaden
Doina Precup

title: Metrics for finite markov decision processes

venue: In UAI,

authors:
Vincent Francois-Lavet
Yoshua Bengio
Doina Precup
Joelle Pineau

title: Combined Reinforcement Learning via Abstract Representations. pp. 3582–3589

venue: Association for the Advancement of Artificial Intelligence (AAAI),

authors:
Robert Givan
Thomas Dean
Matthew Greig

title: Equivalence notions and model minimization in markov decision processes

venue: Artificial Intelligence,

authors:
David Ha
Jürgen Schmidhuber

title: Recurrent world models facilitate policy evolution

venue: In Advances in Neural Information Processing Systems,

authors:
Tuomas Haarnoja
Aurick Zhou
Pieter Abbeel
Sergey Levine

title: Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor

venue: arXiv preprint arXiv:1801.01290,

authors:
Danijar Hafner
Timothy Lillicrap
Mohammad Norouzi
Jimmy Ba

title: Mastering atari with discrete world models

venue: arXiv preprint arXiv:2010.02193,

authors:
Zhengyao Jiang
Dixing Xu
Jinjun Liang

title: A deep reinforcement learning framework for the financial portfolio management problem

venue: arXiv preprint arXiv:1706.10059,

authors:
Lukasz Kaiser
Mohammad Babaeizadeh
Piotr Milos
Blazej Osinski
Roy H Campbell
Konrad Czechowski
Dumitru Erhan
Chelsea Finn
Piotr Kozakowski
Sergey Levine

title: Model-based reinforcement learning for atari

authors:
Yu Lei
Wenjie Li

title: Interactive recommendation with user-specific deep reinforcement learning

venue: ACM Transactions on Knowledge Discovery from Data (TKDD),

authors:
Lihong Li
Thomas J Walsh
Michael L Littman

title: Towards a unified theory of state abstraction for mdps

venue: In ISAIM,

authors:
Feng Liu
Ruiming Tang
Xutao Li
Weinan Zhang
Yunming Ye
Haokun Chen
Huifeng Guo
Yuzhou Zhang

title: Deep reinforcement learning based recommendation with explicit user-item interactions modeling

authors:
Hongzi Mao
Mohammad Alizadeh
Ishai Menache
Srikanth Kandula

title: Resource management with deep reinforcement learning

venue: In Proceedings of the 15th ACM Workshop on Hot Topics in Networks,

authors:
Tomas Mikolov
Kai Chen
Greg Corrado
Jeffrey Dean

title: Efficient estimation of word representations in vector space

venue: arXiv preprint arXiv:1301.3781,

authors:
Volodymyr Mnih
Koray Kavukcuoglu
David Silver
Alex Graves
Ioannis Antonoglou
Daan Wierstra
Martin Riedmiller

title: Playing atari with deep reinforcement learning

venue: arXiv preprint arXiv:1312.5602,

authors:
Volodymyr Mnih
Koray Kavukcuoglu
David Silver
Andrei A Rusu
Joel Veness
Marc G Bellemare
Alex Graves
Martin Riedmiller
Andreas K Fidjeland
Georg Ostrovski

title: Human-level control through deep reinforcement learning

authors:
Jelle Munk
Jens Kober
Robert Babuška

title: Learning state representation for deep actor-critic control

venue: In 2016 IEEE 55th Conference on Decision and Control (CDC),

authors:
Karthik Narasimhan
Tejas Kulkarni
Regina Barzilay

title: Language understanding for text-based games using deep reinforcement learning

venue: arXiv preprint arXiv:1506.08941,

authors:
Jason Pazis
Ronald Parr

title: Generalized value functions for large action sets

venue: In Proceedings of the 28th International Conference on International Conference on Machine Learning (ICML),

authors:
Brian Sallans
Geoffrey E Hinton

title: Reinforcement learning with factored states and actions

venue: Journal of Machine Learning Research,

title: On learning to think: Algorithmic information theory for novel combinations of reinforcement learning controllers and recurrent neural world models

venue: arXiv preprint arXiv:1511.09249,

authors:
John Schulman
Filip Wolski
Prafulla Dhariwal
Alec Radford
Oleg Klimov

title: Proximal policy optimization algorithms

venue: arXiv preprint arXiv:1707.06347,

authors:
Guy Shani
David Heckerman
Ronen I Brafman

title: An mdp-based recommender system

authors:
Sahil Sharma
Aravind Suresh
Rahul Ramesh
Balaraman Ravindran

title: Learning to factor policies and action-value functions: Factored action space representations for deep reinforcement learning

venue: arXiv preprint arXiv:1705.07269,

authors:
Guy Tennenholtz
Shie Mannor

title: The natural language of actions. volume

venue: Proceedings of Machine Learning Research,

authors:
Emanuel Todorov
Tom Erez
Yuval Tassa

title: Mujoco: A physics engine for model-based control

venue: In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems,

authors:
Hado Van Hasselt
Marco A Wiering

title: Using continuous action spaces to solve discrete problems

venue: In 2009 International Joint Conference on Neural Networks (IJCNN),

authors:
Amy Zhang
Rowan McAllister
Roberto Calandra
Yarin Gal
Sergey Levine

title: Learning invariant representations for reinforcement learning without reconstruction

venue: arXiv preprint arXiv:2006.10742,

heading: 1 INTRODUCTION

text:
Reinforcement learning (RL) has been successfully applied to a range of tasks, including challenging gaming scenarios (Mnih et al., 2015). However, the application of RL in many real-world domains is often hindered by the large number of possible states and actions these settings present. For instance, resource management in computing clusters (Mao et al., 2016; Evans & Gao, 2016), portfolio management (Jiang et al., 2017), and recommender systems (Lei & Li, 2019; Liu et al., 2018) all suffer from extremely large state/action spaces, thus challenging to be tackled by RL.
In this work, we investigate efficient training of reinforcement learning agents in the presence of large state-action spaces, aiming to improve the applicability of RL to real-world domains. Previous work attempting to address this challenge has explored the idea of learning representations (embeddings) for states or actions. Specifically, for state embeddings, using machine learning to obtain meaningful features from raw state representations is a common practice in RL, e.g. through the use of convolutional neural networks for image input (Mnih et al., 2013). Previous works such as by Ha & Schmidhuber (2018b) have explored the use of environment models, termed world models, to learn abstract state representations, and several pieces of literature explore state aggregation using bisimulation metrics (Castro, 2020). While for action embeddings, the most recent works by Tennenholtz & Mannor (2019) and Chandak et al. (2019) propose methods for learning embeddings for discrete actions that can be directly used by an RL agent and improve generalization over actions. However, these works consider the state representation and action representation as isolated tasks, which ignore the underlying relationships between them. In this regard, we take a different approach and propose to jointly learn embeddings for states and actions, aiming for better generalization over both states and actions in their respective embedding spaces.
To this end, we propose an architecture consisting of two models: a model of the environment that is used to generate state and action representations and a model-free RL agent that learns a policy using the embedded states and actions. By using these two models, our approach combines aspects
of model-based and model-free reinforcement learning and effectively bridges the gap between both approaches. In contrast to model-based RL, however, we do not use the environment model for planning, but to learn state and action embeddings. One key benefit of this approach is that state and action representations can be learned in a supervised manner, which greatly improves sampling efficiency and potentially enables their use for transfer learning. In sum, our key contributions are:
• We formulate an embedding model for states and actions, along with an internal policy πi that leverages the learned state/action embeddings, as well as the corresponding overall policy πo. We show the existence of an overall policy πo that achieves optimality in the original problem domain.
• We further prove the equivalence between updates to the internal policy πi acting in embedding space and updates to the overall policy πo.
• We present a supervised learning algorithm for the proposed embedding model that can be combined with any policy gradient based RL algorithms.
• We evaluate our methodology on some game-based as well as real-world tasks and find that it outperforms state-of-the-art models in both discrete/continuous domains.
The remainder of this paper is structured as follows: In Section 2, we provide some background on RL. We then give an overview of related work in Section 3, before presenting our proposed methodology in Section 4, which we evaluate in Section 5.

heading: 2 BACKGROUND

text:
We consider an agent interacting with its environment over discrete time steps, where the environment is modelled as a discrete-time Markov decision process (MDP), defined by a tuple (S,A, T ,R, γ). Specifically, S and A are the sets of all possible states and actions, referred to as the state space and action space, respectively. In this work, we consider both discrete and continuous state and action spaces. The transition function from one state to another, given an action, is T : S × A 7→ S, which may be deterministic or stochastic. The agent receives a reward at each time step defined by R : S × A 7→ R. γ ∈ [0, 1] denotes the reward discounting parameter. The state, action, and reward at time t ∈ {0, 1, 2...} are denoted by the random variables St, At, and Rt. The initial state of the environment comes from an initial state distribution d0. Thereafter, the agent follows a policy π, defined as a conditional distribution over actions given states, i.e., π(a|s) = P (At = a|St = s). The goal of the reinforcement learning agent is to find an optimal policy π∗ that maximizes the expected sum of discounted accumulated future rewards for a given environment, i.e., π∗ ∈ arg maxπ E[ ∑∞ t=0 γ
tRt|π]. For any policy, we also define the state value function vπ(s) = E[ ∑∞ k=0 γ
kRt+k|π, St = s] and the state-action value function Qπ(s, a) = E[ ∑∞ k=0 γ kRt+k|π, St = s,At = a].

heading: 3 RELATED WORK

text:
For the application of state embeddings in reinforcement learning, there are two dominant strands of research, namely world models and state aggregation using bisimulation metrics. World model approaches train an environment model in a supervised fashion from experience collected in the environment, which is then used to generate compressed state representations (Ha & Schmidhuber, 2018a) or to train an agent using the learned world model (Ha & Schmidhuber, 2018b; Schmidhuber, 2015). Further applications of world models, e.g. for Atari 2000 domains, show that abstract state representations learned via world models can substantially improve sample efficiency (Kaiser et al., 2019; Hafner et al., 2020). Similar to this idea, Munk et al. (2016) pre-train an environment model and use it to provide state representations for an RL agent. Furthermore, de Bruin et al. (2018) investigate additional learning objectives to learn state representations, and Francois-Lavet et al. (2019) propose the use of an environment model to generate abstract state representations; their learned state representations are then used by a Q-learning agent. By using a learned model of the environment to generate abstract states, these approaches capture structure in the state space and reduce the dimensionality of its representation – an idea similar to our proposed embedding model.
In contrast to world models, we considers both states and actions, rather than just states. Additionally, our approach trains an RL agent on the original environment using its embedded representation, rather than learning a policy using a surrogate world model. Bisimulation, on the other hand, is a method for aggregating states that are “behaviorally equivalent” (Li et al., 2006) and can improve convergence speed by grouping similar states into abstract states. Several works on bisimulationbased state aggregation, e.g. Ferns et al. (2004); Givan et al. (2003), present different metrics to measure state similarity. Furthermore, Zhang et al. (2020) and Castro (2020) propose deep learning methods for generating bisimulation-based state aggregations that scale beyond the tabular methods proposed in several earlier works (Castro, 2020). While there are parallels between bisimulation and our approach, we do not propose the aggregation of states. Instead, our embedding technique projects states into a continuous state embedding space, similar to Zhang et al. (2020), where their behavioral similarity is captured in their proximity in embedding space. Furthermore, our method embeds both states and actions and does not employ an explicit similarity metric such as a bisimulation metric, but instead learns the relationships among different states and actions via an environment model. State representations are also used in RL-based NLP tasks, such as Narasimhan et al. (2015), who jointly train an LSTM-based state representation module and a DQN agent, and Ammanabrolu & Riedl (2018), who propose the use of a knowledge graph based state representation.
In addition to state representations, previous work has explored the use of additional models to learn meaningful action representations. In particular, Van Hasselt & Wiering (2009) investigate the use of a continuous actor in a policy gradient based approach to solve discrete action MDPs, where the policy is learned in continuous space and actions are discretized before execution in the environment. Dulac-Arnold et al. (2015) propose a similar methodology, where a policy is learned in continuous action embedding space and then mapped to discrete actions in the original problem domain. Both Van Hasselt & Wiering (2009) and Dulac-Arnold et al. (2015) only consider actions and assume that embeddings are known a priori. Tennenholtz & Mannor (2019) propose a methodology called Act2Vec, where they introduce an embedding model similar to the Skip-Gram model (Mikolov et al., 2013) that is trained using data from expert demonstrations and then combined with a DQN agent. One significant drawback of this approach is that information on the semantics of actions has to be injected via expert demonstration and is not learned automatically. In contrast, Chandak et al. (2019) propose a method that enables the self-supervised learning of action representations from experience collected by the RL agent, using an embedding model that resembles CBOW. The approach presented in (Chandak et al., 2019) is most closely related to our method in the way of integrating embeddings with policy gradient based RL algorithms. However, they only consider action embeddings, while we jointly embed states and actions and improve the performance significantly. As an alternative to embedding methods, Pazis & Parr (2011) and Sallans & Hinton (2004) represent actions in binary format, where they learn a value function associated with each bit, and Sharma et al. (2017) propose the use of factored actions, where actions are expressed using underlying primitive actions. Again, these three methods rely on handcrafted action decomposition.
Previous embedding approaches consider either states or actions in isolation, embedding only one of the two or using separate models for each. By contrast, in addition to capturing structure in the state or action space respectively, we leverage interdependencies between the two by jointly learning embeddings for them, and thus improve the quality of embeddings.

heading: 4 METHODOLOGY

text: We propose an embedding model to jointly learn state and action embeddings and present a generic framework that uses the learned embeddings in conjunction with any policy gradient algorithms.

heading: 4.1 ARCHITECTURE

text:
Our proposed approach has three components: (i) state-action embedding model, (ii) policy gradient based RL agent that takes in the embeded states and outputs actions in action embedding space, and (iii) action mapping function that maps embedded actions to actions in the original problem domain. We define two embedding spaces, i.e., a continuous state embedding space X ∈ Rm and a continuous action embedding space E ∈ Rd; states are projected to the state embedding space via a function φ : S 7→ X . On the other hand, we define a function f : E 7→ A that maps points in the embedding space to actions in the original problem domain. With these embedding spaces, we
further define an internal policy πi that takes in the embedded state and outputs the embedded action, i.e., πi : X 7→ E . Furthermore, the mapping functions φ and f , together with the internal policy πi form the overall policy πo : S 7→ A, i.e., πo is characterized by the following three components:
Xt = φ(St), Et ∼ πi(·|Xt), and At = f(Et),
where random variables St ∈ S, Xt ∈ X , Et ∈ E , and At ∈ A denote the state in the original problem domain, state in the embedding domain, action in the embedding domain, and action in the original domain at time t, respectively. The overall model architecture is illustrated in Figure 1a. In Figure 1a, function φ is directly learned using an embedding model illustrated in Figure 1b (detailed in Section 4.3.1). However, function f is not directly part of this embedding model, but uses the obtained action embeddings to learn a mapping from e ∈ E to a ∈ A (see Section 4.3.2 for details). For an RL problem, the ultimate goal is to optimize the policy π in the original problem domain, i.e., π in the S×A space. For the proposed architecture, the idea is to first establish efficient embeddings for both states and actions in S and A. Then we reduce the problem of finding the optimal π∗ in the original problem domain to the problem of internal policy πi optimization that acts purely in the embedding space. However, it is still unclear how the optimization of πi is related to π∗. In addition, the overall policy πo relying on φ and f further complicates the policy learning. In this regard, we establish the relationships among π, πo, and πi in the next section, which is the foundation to justify the validity of leveraging both state-action embeddings in RL.
4.2 THEORETICAL FOUNDATIONS: RELATIONSHIPS AMONG π, πo , AND πi
To understand the relationships among π, πo, and πi, we first derive how πo is related to π, after which we prove the existence of πo that achieves optimality in the original problem domain. For this goal, two further assumptions are required on the nature of the state embedding function φ and the action mapping function f .
Assumption 1. Given an action embedding Et, At is deterministic and defined by a function f : E 7→ A, i.e., there exists a such that P (At = a|Et = e) = 1. Assumption 2. Given a state St, Xt is deterministic and defined by a function φ : S 7→ X , i.e., there exists x such that P (Xt = x|St = s) = 1.
We validate Assumption 1, which defines f as a many-to-one mapping, empirically via all experiments conducted in Section 5, and find that no two actions share exactly the same embedded representation, i.e., Assumption 1 holds in practice. Note that we also assume the Markov property for our environment, which is a standard assumption for reinforcement learning problems. With slight abuse of notation, we denote the inverse mapping from an action a ∈ A to its corresponding points in the embedding space (one-to-many mapping) by f−1(a) := {e ∈ E : f(e) = a}. Lemma 1. Under Assumptions 1 and 2, for policy π in the original problem domain, there exists πi such that
vπ(s) = ∑ a∈A ∫ {e}=f−1(a) πi(e|x = φ(s))Qπ(s, a) de .
Proof Sketch. Based on the Bellman equation for the value function vπ in the original domain, we introduce the embedded state x = φ(s) using Assumption 2. By the law of total probability and Assumption 1, we then introduce the embedded action e. From the Markovian property and the definition of our model, we can derive Claims 1 - 5 on (conditional) independence in Appendix A.2, which then allow us to remove the original state and action from the expression of the policy, leaving us with πi. See Appendix A.2 for the complete proof.
By Assumptions 1 and 2, Lemma 1 allows us to express the overall policy πo in terms of the internal policy πi as
πo(a|s) = ∫ {e}=f−1(a) πi(e|x = φ(s)) de , (1)
bridging the gap between the original problem domain and the policy in embedding spaces X and E . Under π∗ in the original domain, we define v∗ := vπ ∗ and Q∗ := Qπ ∗ for the ease of discussions. Then using Lemma 1, we now prove the existence of an overall policy πo that is optimal. Theorem 1. Under Assumptions 1 and 2, there exists an overall policy πo that is optimal, such that vπo = v∗.
Proof. Under finite state and action sets, bounded rewards, and γ ∈ [0, 1), at least one optimal policy π∗ exists. From Lemma 1, we then have
v∗(s) = ∑ a∈A ∫ {e}=f−1(a) πi(e|φ(s))Q∗(s, a) de . (2)
Thus, ∃ φ, f , and πi, representing an overall policy πo, which is optimal, i.e., vπo = v∗.
Theorem 1 suggests that in order to get the optimal policy π∗ in the original domain, we can focus on the optimization of the overall policy πo, which will be discussed in the next section.

heading: 4.3 ARCHITECTURE EMBODIMENT AND TRAINING

text: In this section, we first present the implementation and the joint-training of the state-action embedding model (illustrated in Figure 1b). Based on this method and according to Theorem 1, we then propose a strategy to train the overall policy πo, where functions φ and f are iteratively updated.

heading: 4.3.1 JOINT TRAINING OF THE STATE-ACTION EMBEDDING

text:
In this section, we elaborate on our proposed embedding model for jointly learning state and action embeddings. Specifically, we require two components: (i) function φ that projects St into embedding space X and (ii) function f that maps each point in embedding space E to an action At in the original problem domain. Since these functions are not known in advance, we train estimators using a model of the environment (see Figure 1b). This environment model requires two further components: (iii) function g : A 7→ E and (iv) transition model T : E ×X 7→ S that predicts the next state from the concatenated embeddings of the current state and action. Note that g is a one-to-one mapping, since according to Assumption 1, no two actions have exactly the same embedding. Denote the estimators of components (i)-(iv) by φ̂ , f̂ , ĝ , T̂ ; the target of the environment model is
P̂ (St+1|St, At) = T̂ (St+1|Xt, Et)ĝ(Et|At)φ̂(Xt|St) . (3) The difference between the true transition probabilities P (St+1|St, At) and the estimated probabilities P̂ (St+1|St, At) can be measured using the Kullback-Leibler (KL) divergence, where the expectation is over the true distribution P (St+1|St, At), i.e.,
DKL(P ||P̂ ) =− ESt+1∼P (St+1|St,At) [ ln ( P̂ (St+1|St, At) P (St+1|St, At) )] . (4)
Since we observe the full tuple (St+1, St, At), we can compute a sample estimate of this using Equation (4). In Equation (4), the denominator does not depend on φ̂, ĝ, or T̂ . Therefore, we define the loss function for the embedding model as
L(φ̂, ĝ, T̂ ) = −ESt+1∼P (St+1|St,At) [ ln(P̂ (St+1|St, At)) ] . (5)
Note that the estimator f̂ is not directly included in the embedding model. Instead, we find f̂ by minimizing the reconstruction error between the original action and the action reconstructed from the embedding. We can again define this using the KL divergence as
L(f̂) = −EAt∼f(At|Et) [ ln(f̂(At|Et)) ] . (6)
Then, all components of our embedding model can be learned by minimizing the loss functions L(φ̂, ĝ, T̂ ) and L(f̂). Note that the embedding model is trained in two steps, firstly by minimizing the loss in Equation (5) to update φ̂, ĝ, and T̂ and secondly by minimizing the loss in Equation (6) to update f̂ . Here, we give the target of the embedding model for the case of discrete state and action spaces. However, the model is equally applicable to continuous domains. In this case, the loss functions L(φ̂, ĝ, T̂ ) and L(f̂) must be replaced with loss functions suitable for continuous domains. For continuous domains, we adopt a mean squared error loss instead of the losses given in Equations (5) and (6), but other loss functions may also be applicable. For our experiments, we parameterize all components of the embedding model illustrated in Figure 1b as neural networks. More details on this can be found in Appendix A.4.

heading: 4.3.2 STATE-ACTION EMBEDDING AND POLICY LEARNING

text:
Theorem 1 only shows that optimizing πo can help us achieve the optimality in the original domain. Then a natural question is whether we can optimize πo by directly optimizing the internal policy πi using the state/action embeddings in Section 4.3.1. To answer this question, we first mathematically derive that updating πi is equivalent to updating πo.We then proceed to present an iterative algorithm for learning the state-action embeddings and the internal policy.
Suppose πi is parameterized by θ. Then with the objective of optimizing πo by only updating πi, we define the performance function of πo as1
Jo(φ, θ, f) = ∑ s∈S d0(s) ∑ a∈A πo(a|s)Qπo(s, a) . (7)
Let the state-action-value function for the internal policy beQπi(x, e) = E[ ∑∞ k=0 γ
kRt+k|πi, Xt = x,Et = e]. We can then define the performance function of the internal policy as
Ji(θ) = ∫ x∈X d0(x) ∫ e∈E πi(e|x)Qπi(x, e) de dx . (8)
The parameters of the overall policy πo can then be learned by updating its parameters θ in the direction of ∂Jo(φ, θ, f)/∂θ, while the parameters of the internal policy can be learned by updating θ in the direction of ∂Ji(θ)/∂θ. With the following assumption, we then have Lemma 2. Assumption 3. The state embedding function φ(s) maps each state s to a unique state embedding x = φ(s), i.e., ∀si 6= sj , P (φ(si) = φ(sj)) = 0.
Note that Assumption 3, defining φ as a one-to-one mapping, is not theoretically guaranteed by the definition of our embedding model. Nevertheless, we test this empirically in Section 5 and find that no two states share the same embedded representation in any of our experiments, thus justifying Assumption 3 in practical scenarios considered in this work. Lemma 2. Under Assumptions 1–3, for all deterministic functions f and φ, which map each point e ∈ E and s ∈ S to an action a ∈ A and to an embedded state x ∈ X , and the internal policy πi parameterized by θ, the gradient of the internal policy’s performance function ∂Ji(θ) ∂θ equals the gradient of the overall policy’s performance function ∂Jo(φ,θ,f)∂θ , i.e.,
∂Ji(θ) ∂θ = ∂Jo(φ, θ, f) ∂θ .
Proof Sketch. Assumptions 1 and 3 allow us to show the equivalence between the internal stateaction-value function and the overall state-action-value function. Using this equivalence, we can then remove functions φ and f from ∂Jo(φ, θ, f)/∂θ and are left with the gradient of the internal policy πi only. The complete proof is deferred to Appendix A.3.
1Equation (7) is for the overall policy with discrete states and actions. Note that our approach is also applicable in continuous domains. ∑ s∈S and ∑ a∈A would then be replaced by ∫ s and ∫ a , respectively.
Algorithm 1: Joint Training of State-Action Embeddings and the Internal Policy Initialize the state and action embeddings (optional pre-training); for Epoch = 0, 1, 2 ... do
Lemma 2 shows that the updates to the internal policy πi and the overall policy πo are equivalent. This allows us to optimize the overall policy by making updates directly to the internal policy πi, thereby avoiding the potentially intractable computation of the inverse functions f−1 and φ−1. Since there are no special restrictions on the internal policy πi, we can use any policy gradient algorithm designed for continuous control to optimize the policy. In addition, we can also iteratively update parameters for jointly learning φ, f , and πi, as shown in Algorithm 1.

heading: 5 EMPIRICAL EVALUATION

text: We evaluate our proposed architecture on game-based applications, robotic control, and a real-world recommender system, covering different combinations of discrete and continuous state and action spaces. Our methodology is evaluated using Vanilla Policy Gradient (VPG) (Achiam, 2018), Proximal Policy Optimization (PPO) (Schulman et al., 2017) and Soft Actor Critic (SAC) (Haarnoja et al., 2018) algorithms. We benchmark the performance against these algorithms without embeddings and with action embeddings generated by the method called policy gradients with Representations for Actions (RA) proposed by Chandak et al. (2019). We pre-train the embedding model on randomly collected samples for all experiments and enable continuous updates for the Ant-v2 and recommender system environments. Complete parameterizations are presented in Appendix A.4.

heading: 5.1 PROOF-OF-CONCEPT: GRIDWORLD AND SLOTMACHINE

text:
Gridworld: It is similar to that used by Chandak et al. (2019). States are given as a continuous coordinate, while actions are defined via n actuators, equally spaced around the agent, which move the agent in the direction they are pointing towards. Then each combination of actuators forms an action, resulting in 2n unique actions. We run two sets of experiments in this environment: (i) we use the continuous state directly and (ii) we discretize the coordinate in a 40 × 40 grid. The agent receives small step and collision penalties and a large positive reward for reaching the goal state.
Figure 2 illustrates the learned embeddings for this environment. The obtained embeddings for states and actions illustrate that our approach is indeed able to obtain meaningful representations. Actions are embedded according to the displacement in the Cartesian coordinates they represent (Figure 2d). Figure 2b shows that states are embedded according to the coordinate they represent and that the embeddings capture the L-shaped obstacle in the original problem domain.
Slotmachine: It consists of several reels, each of which can be individually spun by the agent for a fraction of a full turn. Then, each unique combination of these fractions constitutes an action. The discrete state is given by the numbers on the front of the reels.
Both scenarios yield large discrete action spaces, rendering them well-suited for evaluating our approach. The results for the gridworld and slotmachine environments are shown in Figure 3a, 3b, and 3c. Our approach outperforms both benchmarks on the discretized gridworld and the slotmachine environments and is comparable for the continuous-state gridworld. Such observation suggests that our approach is particularly well suited for discrete domains. Intuitively, this is because the relationship between different states or action in discrete domains is usually not apparent, e.g., in one-hot encoding. In continuous domains, however, structure in the state and action space is already captured to some extent, rendering it harder to uncover additional structure by embeddings; nevertheless, when a continuous problem becomes more complicated, our embedding method again shows effectiveness in capturing useful information (see Section 5.2).

heading: 5.2 ROBOTIC CONTROL

text: We use the half-cheetah-v2 and the Ant-v2 environments from the robotic control simulator Mujoco (Todorov et al., 2012). Here, the agent observes continuous states (describing the current position of the cheetah or ant respectively) and then decides on the force to apply to each of the joints of the robot (i.e., continuous actions). Rewards are calculated based on the distance covered in an episode. From the results presented in Figure 3d, we observe that our method does not outperform the benchmark on the half-cheetah-v2 environment. On the Ant-v2 environment, however, our approach outperforms the benchmarks. Similar to the experiments in a continuous gridworld, it is less likely that embedding uncovers additional structure on top of the structure inherent to the continuous state and action, explaining the comparable performance with and without embeddings on the halfcheetah-v2 domain. Nevertheless, in the more challenging Ant-v2 environment, the state and action representation is more complex and the agent might thus benefit from the reduced dimensionality in embedding space. Additionally, the learned embeddings appear to be able to capture sufficient structure in the state and action space, despite using an embedding space with lower dimensionality.

heading: 5.3 RECOMMENDER SYSTEM

text: In addition, we also test our methodology on a real-world application of a recommender system. We use data on user behavior from an e-commerce store collected in 2019 (Kechinov, 2019). The environment is constructed as an n-gram based model, following Shani et al. (2005). Based on the n previously purchased items, a user’s purchasing likelihood is computed for each item in the store. Recommending an item then scales the likelihood for that item by a pre-defined factor. States are represented as a concatenation of the last n purchased items and each item forms an action. In this environment we have 835 items (actions) and approx. 700, 000 states. The results obtained under various RL methods are reported in Figure 3g. We find that our approach leads to significant improvements in both the convergence speed and final performance compared to benchmarks. This result confirms that our method is particularly useful in the presence of large discrete state and action spaces. Interestingly, the PPO and VPG benchmarks outperform the benchmarks using the RA methodology. We conjecture that the action embeddings generated in RA on this environment in fact obfuscate the effect an action has, and thus limit the agent’s performance.

heading: 5.4 APPLICATION IN TRANSFER LEARNING

text: In addition to improving the convergence speed and the accumulated reward, we test whether the policies learned in one environment can be transferred to a similar one. We consider two environments (old and new) that differ by only the state space or the action space. To leverage the previous experience, if the two environments have the same state (or action) space, the state (or action) embedding component in the new environment is initialized with the weights learned from the old environment; we then train our model in this new environment. For evaluation, we use two gridworlds, where only the numbers of actions differ, i.e., 512 (old) and 2048 (new). Results in Figure 3e show that the model with the transferred knowledge outperforms the case where the entire model is trained from scratch, in terms of the convergence speed. Therefore, it confirms that the previously learned policy serves as a critical initialization for the new gridworld environment and sheds light on the potential of our approach in transfer learning; further evaluations are left for future work.

heading: 6 CONCLUSION

text: In this paper, we presented a new architecture for jointly training states/action embeddings and combined this with common reinforcement learning algorithms. Our theoretical results confirm the validity of the proposed approach, i.e., the existence of an optimal policy in the embedding space. We empirically evaluated our method on several environments, where it outperforms state-of-theart RL approaches in complex large-scale problems. Our approach is easily extensible as it can be combined with most existing RL algorithms and there are no special restrictions on the parameterizations of the embedding model. Consequently, we conclude that the proposed methodology works towards the applicability of reinforcement learning to real-world problems.

heading: A APPENDIX

text:
A.1 CLAIMS
In addition to Assumptions 1 and 2, we derive three claims on conditional independence from the definition of our embedding model and the Markovian property of the environment, listed below.
Claim 1. Action At is conditionally independent of St given a state embedding Xt:
P (At = a, St = s|Xt = x) = P (At = a|Xt = x)P (St = s|Xt = x) .
Claim 2. Action At is conditionally independent of St given state and action embeddings Xt and Et:
P (At = a, St = s|Xt = x,Et = e) = P (At = a|Xt = x,Et = e)P (St = s|Xt = x,Et = e) .
Claim 3. Next state St+1 is conditionally independent of Et and Xt given At and St:
P (St+1 = s ′, Et = e,Xt = x|At = a, St = s) =
P (St+1 = s ′|At = a, St = s)P (Et = e,Xt = x|At = a, St = s) .
Using Claims 1 - 3, we can derive two further auxiliary claims that will be used in the proof of Lemma 1.
Claim 4. The probability of next state St+1 is independent of Et and Xt, given state St and action At, i.e.,
P (St+1 = s ′|Et = e,Xt = x,At = a, St = s) = P (St+1 = s′|At = a, St = s) .
Proof. From Claim 3, we have
P (s′, e, x|a, s) =P (s′|a, s)P (e, x|a, s) P (s′|e, x, a, s)P (e, x|a, s) =P (s′|a, s)P (e, x|a, s)
P (s′|e, x, a, s) =P (s′|a, s) .
Claim 5. The probability of actionAt is independent of St andXt, given action embedding Et, i.e.,
P (At = a|St = s,Xt = x,Et = e) =P (At = s|Et = e) .
Proof. From Claim 2, we have
P (a, s|x, e) =P (a|x, e)P (s|x, e) P (a|s, x, e) =P (a|x, e) .
Since action a only depends on the action embedding e in our model, this becomes
P (a|s, x, e) =P (a|e) .
A.2 PROOF OF LEMMA 1
Lemma 1. Under Assumptions 1 and 2, for policy π in the original problem domain, there exists πi such that
vπ(s) = ∑ a∈A ∫ {e}=f−1(a) πi(e|x = φ(s))Qπ(s, a) de
Proof. The Bellman equation for a MDP is given by
vπ(s) = ∑ a∈A π(a|s) ∑ s′∈S P (s′|s, a)G ,
where G denotes the return, i.e. [R(s, a) + γvπ(s′)], which is a function of s, a, and s′. By rearranging terms we get
vπ(s) = ∑ a∈A ∑ s′∈S π(a|s)P (s′|s, a)G
= ∑ a∈A ∑ s′∈S π(a|s)P (s ′, s, a) P (s, a) G
= ∑ a∈A ∑ s′∈S π(a|s) P (s ′, s, a) π(a|s)P (s) G
= ∑ a∈A ∑ s′∈S P (a|s′, s)P (s′, s) P (s) G .
Since s can be deterministically mapped to x via x = φ(s), by Assumption 2, we have
vπ(s) = ∑ a∈A ∑ s′∈S P (x, a|s′, s)P (s′, s) P (s) G
= ∑ a∈A ∑ s′∈S P (x, a|s′, s)P (s′, s)P (x|s) P (x|s)P (s) G
= ∑ a∈A ∑ s′∈S P (x, a, s′, s)P (x|s) P (x, s) G
= ∑ a∈A P (x|s) ∑ s′∈S P (x, a, s′, s) P (x, s) G
= ∑ a∈A P (x|s) ∑ s′∈S P (a, s′|x, s)P (x, s) P (x, s) G
= ∑ a∈A P (x|s) ∑ s′∈S P (a, s′|x, s)G
= ∑ a∈A P (x|s) ∑ s′∈S P (s′|a, x, s)P (a|s, x)G .
From Claim 1, we know that P (a|s, x) = P (a|x). Therefore,
vπ(s) = ∑ a∈A P (x|s) ∑ s′∈S P (s′|a, x, s)P (a|x)G .
Since P (x|s) is deterministic by Assumption 2 and evaluates to 1 for the representation of φ(s) = x, we can rewrite the equation above to
vπ(s) = ∑ a∈A ∑ s′∈S P (s′|a, x, s)P (a|x)G .
We now proceed to establish the relationship with the action embedding. From above we have vπ(s) = ∑ a∈A ∑ s′∈S P (a|x)P (s′|a, x, s)G
= ∑ a∈A ∑ s′∈S P (a|x)P (s ′, a, x, s) P (a, x, s) G .
From Claim 1, a and s are conditionally independent given x. Therefore, P (a, x, s) = P (a|x)P (s|x)P (x), which allows us to rewrite the above equation as
vπ(s) = ∑ a∈A ∑ s′∈S P (a|x) P (s ′, a, x, s) P (a|x)P (x, s) G
= ∑ a∈A ∑ s′∈S P (s′, a, x, s) P (x, s) G
= ∑ a∈A ∑ s′∈S P (a|s′, x, s)P (s′, x, s) P (x, s) G .
By the law of total probability, we can now introduce the new variable e, which is the embedded action. Then
vπ(s) = ∑ a∈A ∑ s′∈S ∫ e P (a, e|s′, x, s)P (s′, x, s) P (x, s) Gde
= ∑ a∈A ∑ s′∈S ∫ e P (e|x, s)P (a, e|s ′, x, s)P (s′, x, s) P (e|x, s)P (x, s) Gde .
Since e is uniquely determined by x, we can drop s in P (e|x, s), and thus vπ(s) = ∑ a∈A ∑ s′∈S ∫ e P (e|x)P (a, e, s ′, x, s) P (e, x, s) Gde
= ∑ a∈A ∑ s′∈S ∫ e P (e|x)P (a, s′|e, x, s)Gde
= ∑ a∈A ∑ s′∈S ∫ e P (e|x)P (s′|a, e, x, s)P (a|x, e, s)Gde .
Using the previously derived Claims 4 and 5, the above equation can be simplified to vπ(s) = ∑ a∈A ∑ s′∈S ∫ e P (e|φ(s))P (s′|a, s)P (a|e)Gde .
Since the function f , mapping e to a, is deterministic by Assumption 1 and only evaluates to 1 for a particular a and 0 elsewhere, we can rewrite this further as
vπ(s) = ∑ a∈A ∑ s′∈S ∫ f−1(a) P (e|φ(s))P (s′|a, s)Gde .
Summarizing the terms, this becomes vπ(s) = ∑ a∈A ∫ f−1(a) πi(e|φ(s))Qπ(a, s) de .
A.3 PROOF OF LEMMA 2
Lemma 2. Under Assumptions 1–3, for all deterministic functions f and φ which map each point e ∈ E and s ∈ S to an action a ∈ A and to an embedded state x ∈ X , and the internal policy πi parameterized by θ, the gradient of the internal policy’s performance function ∂Ji(θ) ∂θ equals the gradient of the overall policy’s performance function ∂Jo(φ,θ,f)∂θ , i.e.,
∂Ji(θ) ∂θ = ∂Jo(φ, θ, f) ∂θ .
Proof. Recall from Lemma 1 that the overall policy is defined using the internal policy πo(a|s) = ∫ f−1(a) πi(e|φ(s)) de .
We can then define the performance function of the overall policy using the internal policy as Jo(φ, θ, f) = ∑ s∈S d0(s) ∑ a∈A ∫ f−1(a) πi(e|φ(s))Qπo(s, a) de .
The gradient of this performance function w.r.t. θ is
∂Jo(φ, θ, f)
∂θ =
∂
∂θ [∑ s∈S d0(s) ∑ a∈A ∫ f−1(a) πi(e|φ(s))Qπo(s, a) de ] . (9)
Now consider another policy π′i as illustrated in Figure 4, i.e., a policy that takes the raw state input St and outputs Et in embedding space. By Assumption 1, Qπo(s, a) = Qπ ′ i(s, e) as e is deterministically mapped to a. However, Qπi(x, e) = Es∼φ−1(x)[Qπ ′ i(s, e)] = Es∼φ−1(x)[Qπo(s, a)], where φ−1(·) is the inverse mapping from x to s. Nevertheless, by Assumption 3, Es∼φ−1(x)[Qπo(s, a)] = Qπo(s, a). Therefore, we can replace the overall Qπo(a, s) in Equation (9) by the internal Qπi(φ(s), e) as follows
∂Jo(φ, θ, f)
∂θ =
∂
∂θ [∑ s∈S d0(s) ∑ a∈A ∫ f−1(a) (πi(e|φ(s))Qπi(φ(s), e) ] de .
Then the deterministic mapping of e to a by function f further allows us to replace the integral over f−1(a) and the summation over A with an integral over e. Moreover, by Assumption 3, we have d0(s) = d0(x = φ(s)). Therefore, the above can be rewritten as
∂Jo(φ, θ, f)
∂θ =
∂
∂θ [ ∫ x∈X d0(x) ∫ e∈E πi(e|x)Qπi(x, e) de dx ]
= ∂Ji(θ)
∂θ .
A.4 PARAMETERIZATION
A.4.1 IMPLEMENTATION AND ENVIRONMENTS
We implement all evaluated models in Python and use the PyTorch library for neural network implementation (Paszke et al., 2019). The implementation for the benchmark models is adapted from OpenAI’s SpinningUp package (Achiam, 2018). In the following, we elaborate on the different environments used for experimentation.
Gridworld The gridworld environments are constructed using a 1 × 1 grid to represent the underlying continuous state space. For the continuous state version of the environment, a continuous coordinate with the position of the agent in this grid is used as the state. In the discrete state version, we discretize this continuous coordinate in a 40× 40 grid. Actions in these environments are generated from n actuators, equally spaced around the agent, which move the agent in the direction they are pointing towards. Then each combination of actuators forms a unique action, resulting in 2n unique actions. An illustration of this environment can be seen in Figure 5. We limit the maximum number of steps the agent can take in an episode to 150 and set the maximum step size to 0.2. The agent receives a positive reward of 1 for reaching the goal step and receives a small penalty of−0.01 for each step taken and an additional penalty of −0.05 for colliding with an obstacle or attempting to leave the grid.
Slotmachine The slotmachine environment is defined using 4 reels, with 6 values per reel, randomly sampled from the range between 0 and 8. Note that we sample these reels once and then keep them fixed across all experiments to ensure that these run in the same environment. The action space is defined by the fraction of a full turn for which each reel is spun. We set the maximum fraction to 0.8 and use increments of 0.1 to set the fraction of a full turn. Since each combination of fractions forms an action, we get 4, 096 actions in total. The maximum number of steps per episode is set to 20. Rewards are calculated according to
Rt = ∑
i=1,2,3,4
vi × 1{ni ≥ m} ,
where vi denotes the value of the ith reel, ni is the number of times vi occurred across all 4 reels, m is the minimum number of times it has to occur to result in non-zero reward, and 1{·} is an indicator function. In our experiments, we set the threshold m = 3.
Robotic Control: Half-Cheetah and Ant-v2 For the experiments in a robotic control environment, we use the physics engine Mujoco and use the Half-Cheetah and Ant-v2 environments implemented therein (Todorov et al., 2012). We limit the maximum number of steps per episode to 250 for the Half-Cheetah environment and to 1000 for the Ant-v2 environment.
Real-world Recommender System The environment model for the recommender system is constructed from a publicly available data set containing data on approximately 43 million user interactions collected from a multi-category e-commerce store (Kechinov, 2019). We only use purchase
events to create our environment model, where we filter out items that were bought fewer than 100 times and users that made less than 3 purchases. We are then left with approximately 250, 000 data points and 835 items. The environment model is then created using an n-gram based model, with n set to 2. The state is represented as the concatenation of the n last items purchased and each item available for recommendation forms an action. We thus have 835 actions and 697, 225 possible states. In the environment model, we compute the purchasing probability for a user for all items, given the current state, i.e., the n last purchases that user made. We parameterize a categorical distribution with these transition probabilities to sample the next item. The RL agent choses which item to recommend. Recommending an item increases its transition probability by a factor α = 3, after which the sum of all transition probabilities for the current state are normalized to sum to 1. The transition function in this environment is stochastic as we have transition probabilities for each item and each state. For more details on the environment model used, please refer to Shani et al. (2005), whose methodology is adopted. As the reward function, we simply use the revenue obtained in each time step, i.e., the price of the item that the customer bought after making a certain recommendation.
A.4.2 HYPER PARAMETERS
Training for all experiments is conducted in epochs, where each epoch corresponds to a set of samples. At the end of an epoch, the actor and critic are trained for a number of iterations using their respective learning rates. For the VPG and PPO algorithms, the actor network is trained for 3 iterations and the critic network is trained for 6 iterations at the end of an epoch in all of the conducted experiments, using a single batch in each iteration, consisting of all samples collected during the epoch. The SAC updates are also run at the end of an epoch, where we train the actor and critic on i batches sampled from the replay buffer, where i is set equal to the number of samples collected during the epoch. For all experiments, training is conducted in epochs, where a fixed number of samples are collected per epoch. Rather than updating the models after every timestep or episode, we run several iterations of gradient descent at the end of an epoch (see the reference implementations in Achiam (2018)). The batch size is searched over {256, 512} in all SAC experiments. The environment configuration is kept fixed for each environment across different models as well as the hyper parameter searches to ensure comparability.
We conduct a gridsearch over different parameter combinations for each experiment and report the average per episode return and standard deviation achieved over 10 runs with different random seeds, using the hyper parameters that perform the best in the gridsearch. The best hyper parameters – apart from those we kept fixed – are reported in Tables 1 - 7. The actor and critic learning rates are searched over {0.001, 0.0003, 0.0001} for all environments. We also search the initialization value for the standard deviation of the actor over {0.3, 0.6, 0.8, 1.0} in all experiments and keep γ = 0.99 and λ = 0.97 fixed. Note that λ is only used for general advantage estimation in the VPG algorithm. For all experiments apart from those using the recommender system and the Ant-v2 environment, we pre-train the embedding model using 30, 000 randomly collected samples and keep the embeddings fixed thereafter. For the recommender system, we pre-train the embedding model using 50, 000 randomly collected samples andalso allow continuous updating of the learned embeddings. Similarly, for the Ant-v2 environment, we pre-train the embedding model using 100, 000 randomly collected samples and then enable continuous updating. In all cases, we apply a Tanh non-linearity to the learned embeddings to bound their effective range. Throughout our experiments, we keep the learning rate of the embedding model fixed at 0.01. For continuous domains, we parameterize the mapping function f as a neural network with a single layer with 64 hidden units. In discrete domains, we use a simpler nearest neighbor mapping function to parameterize f , i.e., we map the output of the internal policy πi to the action that is closest in embedding space E . For experiments on the gridworld environments, we fix the state and action embedding dimensions at 8 and 2, respectively. The network size for the actor and critic is set to (64, 64), i.e., a two-layer network with 64 neurons in each layer. We then train the agent for 700 epochs, with 1, 500 steps, i.e., a minimum of 10 episodes per epoch.
In the experiments on the slotmachine environment, the dimension for state and action embeddings is searched over {8, 16} and {2, 4}, respectively. The network size for the actor and critic is kept fixed at (64, 64). The agent is trained for 2, 500 epochs, with 200 samples (a minimum of 10 episodes per epoch).
For experiments on the Half-Cheetah environment, the state and action embedding dimensions are searched over {10, 20} and {6, 10}, respectively. We search the network size over {(64, 64), (128, 128)} and train the agent for 1, 000 epochs with 1, 500 steps per epoch, i.e., at least 6 episodes per epoch.
In the experiments on the Ant-v2 environment, the dimensions of the state and action embeddings are searched over {16, 32, 64} and {6, 8, 10}, respectively. We search the network size over {(64, 64), (128, 128)} and train the agent for 2, 000 epochs with 3, 000 steps per epoch, i.e., a minimum of 3 episodes per epoch. The embeddings are pre-trained using 100, 000 randomly generated samples. Continuous updating for the leared embedding model is enabled for the Ant-v2 environment.
On the recommender system environment, we evaluate Proximal Policy Optimisation (PPO) and Soft Actor Critic (SAC) algorithms in addition to the previously used Vanilla Policy Gradient (VPG). For the experiments using the SAC algorithm, we search the actor parameter α, which corresponds to an exploration-exploitation trade-off over {0.1, 0.3}. For experiments involving a PPO algorithm, we set the gradient clipping ratio to 0.2 and keep this fixed throughout. Embedding dimensions for states and actions are searched over {8, 16, 32} and {4, 8, 12}, respectively. In all experiments, the embedding functions are pre-trained using 200, 000 randomly generated samples. In contrast to experiments on other environments, we enable continuous updating of these learned embeddings for the recommender system. We search the size of the actor and critic network over {(64, 64), (128, 128)} and train the agent for 1, 000 epochs with 1, 000 steps per epoch. The episode length is set to 20 steps.


summaries: Learning representations using self-supervision requires domain expertise to identify diverse transformations of the data samples that label preserving. This can be expensive and hard to obtain in many data modalities. The paper proposes to automate this by learning to generate transformations tailed to each modality and sample. Specifically, an adversarial strategy is applied to learning transformations that are close to the original view in the input space but hard to classify for the self-supervision encoder.

summaries: The paper presents a generative model to automatically generate data that is needed for contrastive learning, with a focus on the SimCLR framework, while the method itself is general. Experiments were conducted across multiple modalities, including image, speech and wearable sensor data. The results demonstrate the effectiveness of the proposed model as compared to data augmentation methods relying on human domain knowledge.

summaries: The paper proposes a method for automatic generation of data views for contrastive self-supervised learning of representations. The method consists of learning an adversarial perturbation model that aims to maximize the distance between the original image and its perturbed views in the space of learned representations. To avoid collapse to a completely information-destroying perturbation model, authors propose to limit the perturbation strength in terms of the $l_p$ norm of the added noise. Authors apply their method on various image, speech and wearable sensor datasets where the proposed approach provides an improvement over other methods.

abstractText: Many recent methods for unsupervised representation learning train models to be invariant to different “views,” or distorted versions of an input. However, designing these views requires considerable trial and error by human experts, hindering widespread adoption of unsupervised representation learning methods across domains and modalities. To address this, we propose viewmaker networks: generative models that learn to produce useful views from a given input. Viewmakers are stochastic bounded adversaries: they produce views by generating and then adding an `p-bounded perturbation to the input, and are trained adversarially with respect to the main encoder network. Remarkably, when pretraining on CIFAR-10, our learned views enable comparable transfer accuracy to the welltuned SimCLR augmentations—despite not including transformations like cropping or color jitter. Furthermore, our learned views significantly outperform baseline augmentations on speech recordings (+9 points on average) and wearable sensor data (+17 points on average). Viewmaker views can also be combined with handcrafted views: they improve robustness to common image corruptions and can increase transfer performance in cases where handcrafted views are less explored. These results suggest that viewmakers may provide a path towards more general representation learning algorithms—reducing the domain expertise and effort needed to pretrain on a much wider set of domains. Code is available at https://github.com/alextamkin/viewmaker. Figure 1: Viewmaker networks generate complex and diverse input-dependent views for unsupervised learning. Examples shown are for CIFAR-10. Original image in center with pink border.

authors:
Nasir Ahmed
T Natarajan
Kamisetty R Rao

title: Discrete cosine transform

venue: IEEE transactions on Computers,

authors:
Antreas Antoniou
Amos Storkey
Harrison Edwards

title: Data augmentation generative adversarial networks, 2017

authors:
Philip Bachman
R Devon Hjelm
William Buchwalter

title: Learning representations by maximizing mutual information across views

venue: In Advances in Neural Information Processing Systems,

authors:
Suzanna Becker
Geoffrey E Hinton

title: Self-organizing neural network that discovers surfaces in random-dot stereograms

authors:
Sören Becker
Marcel Ackermann
Sebastian Lapuschkin
Klaus-Robert Müller
Wojciech Samek

title: Interpreting and explaining deep neural networks for classification of audio signals, 2018

authors:
Yoshua Bengio
Aaron Courville
Pascal Vincent

title: Representation learning: A review and new perspectives

venue: IEEE transactions on pattern analysis and machine intelligence,

title: Experiment tracking with weights and biases, 2020

venue: URL https://www. wandb.com/. Software available from wandb.com

authors:
Avishek Joey Bose
Huan Ling
Yanshuai Cao

title: Adversarial contrastive estimation

venue: arXiv preprint arXiv:1805.03642,

authors:
Mathilde Caron
Ishan Misra
Julien Mairal
Priya Goyal
Piotr Bojanowski
Armand Joulin

title: Unsupervised learning of visual features by contrasting cluster assignments

venue: arXiv preprint arXiv:2006.09882,

authors:
Olivier Chapelle
Jason Weston
Léon Bottou
Vladimir Vapnik

title: Vicinal risk minimization

venue: In Advances in neural information processing systems,

authors:
Ting Chen
Simon Kornblith
Mohammad Norouzi
Geoffrey Hinton

title: A simple framework for contrastive learning of visual representations

venue: arXiv preprint arXiv:2002.05709,

authors:
Ting Chen
Simon Kornblith
Kevin Swersky
Mohammad Norouzi
Geoffrey Hinton

title: Big selfsupervised models are strong semi-supervised learners

venue: arXiv preprint arXiv:2006.10029,

authors:
Xinlei Chen
Haoqi Fan
Ross Girshick
Kaiming He

title: Improved baselines with momentum contrastive learning

venue: arXiv preprint arXiv:2003.04297,

authors:
Ekin D Cubuk
Barret Zoph
Dandelion Mane
Vijay Vasudevan
Quoc V Le

title: Autoaugment: Learning augmentation policies from data

venue: arXiv preprint arXiv:1805.09501,

authors:
Ekin D Cubuk
Barret Zoph
Jonathon Shlens
Quoc V Le

title: Randaugment: Practical automated data augmentation with a reduced search space

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops,

authors:
Terrance DeVries
Graham W. Taylor

title: Dataset augmentation in feature space, 2017

authors:
Jeff Donahue
Karen Simonyan

title: Large scale adversarial representation learning

authors:
Jeff Donahue
Philipp Krähenbühl
Trevor Darrell

title: Adversarial feature learning

venue: arXiv preprint arXiv:1605.09782,

authors:
Alexey Dosovitskiy
Jost Tobias Springenberg
Martin Riedmiller
Thomas Brox

title: Discriminative unsupervised feature learning with convolutional neural networks. In Advances in neural information processing

authors:
Raphael Gontijo-Lopes
Sylvia J Smullin
Ekin D Cubuk
Ethan Dyer

title: Affinity and diversity: Quantifying mechanisms of data augmentation

venue: arXiv preprint arXiv:2002.08973,

title: Nips 2016 tutorial: Generative adversarial networks

venue: arXiv preprint arXiv:1701.00160,

authors:
Ian Goodfellow
Jean Pouget-Abadie
Mehdi Mirza
Bing Xu
David Warde-Farley
Sherjil Ozair
Aaron Courville
Yoshua Bengio

title: Generative adversarial nets

authors:
Jean-Bastien Grill
Florian Strub
Florent Altché
Corentin Tallec
Pierre H Richemond
Elena Buchatskaya
Carl Doersch
Bernardo Avila Pires
Zhaohan Daniel Guo
Mohammad Gheshlaghi Azar

title: Bootstrap your own latent: A new approach to self-supervised learning

venue: arXiv preprint arXiv:2006.07733,

authors:
Raia Hadsell
Sumit Chopra
Yann LeCun

title: Dimensionality reduction by learning an invariant mapping

venue: IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR’06),

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

authors:
Kaiming He
Haoqi Fan
Yuxin Wu
Saining Xie
Ross Girshick

title: Momentum contrast for unsupervised visual representation learning

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

authors:
Dan Hendrycks
Thomas Dietterich

title: Benchmarking neural network robustness to common corruptions and perturbations

venue: arXiv preprint arXiv:1903.12261,

authors:
Dan Hendrycks
Steven Basart
Norman Mu
Saurav Kadavath
Frank Wang
Evan Dorundo
Rahul Desai
Tyler Zhu
Samyak Parajuli
Mike Guo

title: The many faces of robustness: A critical analysis of out-of-distribution generalization

venue: arXiv preprint arXiv:2006.16241,

authors:
Daniel Ho
Eric Liang
Xi Chen
Ion Stoica
Pieter Abbeel

title: Population based augmentation: Efficient learning of augmentation policy schedules

venue: In International Conference on Machine Learning,

authors:
Justin Johnson
Alexandre Alahi
Li Fei-Fei

title: Perceptual losses for real-time style transfer and super-resolution

venue: Computer Vision – ECCV

authors:
Eugene Kharitonov
Morgane Rivière
Gabriel Synnaeve
Lior Wolf
Pierre-Emmanuel Mazaré
Matthijs Douze
Emmanuel Dupoux

title: Data augmenting contrastive learning of speech representations in the time domain

authors:
Minseon Kim
Jihoon Tack
Sung Ju Hwang

title: Adversarial self-supervised contrastive learning, 2020

authors:
György Kovács
László Tóth
Dirk Van Compernolle
Sriram Ganapathy

title: Increasing the robustness of cnn acoustic models using autoregressive moving average spectrogram features and channel dropout

venue: Pattern Recognition Letters,

title: Learning multiple layers of features from tiny images

venue: Technical report,

authors:
Oscar D Lara
Miguel A Labrador

title: A survey on human activity recognition using wearable sensors

venue: IEEE communications surveys & tutorials,

authors:
Siddique Latif
Rajib Rana
Sara Khalifa
Raja Jurdak
Junaid Qadir
Björn W Schuller

title: Deep representation learning in speech processing: Challenges, recent advances, and future trends

authors:
Yann LeCun
Léon Bottou
Yoshua Bengio
Patrick Haffner

title: Gradient-based learning applied to document recognition

venue: Proceedings of the IEEE,

authors:
Sungbin Lim
Ildoo Kim
Taesup Kim
Chiheon Kim
Sungwoong Kim

title: Fast autoaugment

authors:
Tsung-Yi Lin
Michael Maire
Serge Belongie
James Hays
Pietro Perona
Deva Ramanan
Piotr Dollár
C Lawrence Zitnick

title: Microsoft coco: Common objects in context

venue: In European conference on computer vision,

authors:
Raphael Gontijo Lopes
Dong Yin
Ben Poole
Justin Gilmer
Ekin D Cubuk

title: Improving robustness without sacrificing accuracy with patch gaussian augmentation

authors:
Loren Lugosch
Mirco Ravanelli
Patrick Ignoto
Vikrant Singh Tomar
Yoshua Bengio

title: Speech model pre-training for end-to-end spoken language understanding, 2019

authors:
Aleksander Madry
Aleksandar Makelov
Ludwig Schmidt
Dimitris Tsipras
Adrian Vladu

title: Towards deep learning models resistant to adversarial attacks

venue: arXiv preprint arXiv:1706.06083,

authors:
Matthias Minderer
Olivier Bachem
Neil Houlsby
Michael Tschannen

title: Automatic shortcut removal for self-supervised representation learning

venue: arXiv preprint arXiv:2002.08822,

authors:
Ishan Misra
Laurens van der Maaten

title: Self-supervised learning of pretext-invariant representations

authors:
Takeru Miyato
Shin-ichi Maeda
Masanori Koyama
Shin Ishii

title: Virtual adversarial training: a regularization method for supervised and semi-supervised learning

authors:
Fernando Moya Rueda
René Grzeszick
Gernot A Fink
Sascha Feldhorst
Michael Ten Hompel

title: Convolutional neural networks for human activity recognition using body-worn sensors

venue: In Informatics,

authors:
Arsha Nagrani
Joon Son Chung
Andrew Zisserman

title: Voxceleb: a large-scale speaker identification

authors:
Aaron van den Oord
Yazhe Li
Oriol Vinyals

title: Representation learning with contrastive predictive coding

venue: arXiv preprint arXiv:1807.03748,

authors:
Sinno Jialin Pan
Qiang Yang

title: A survey on transfer learning

venue: IEEE Transactions on knowledge and data engineering,

authors:
Vassil Panayotov
Guoguo Chen
Daniel Povey
Sanjeev Khudanpur

title: Librispeech: An asr corpus based on public domain audio

authors:
Daniel S Park
William Chan
Yu Zhang
Chung-Cheng Chiu
Barret Zoph
Ekin D Cubuk
Quoc V Le

title: Specaugment: A simple data augmentation method for automatic speech recognition

authors:
Adam Paszke
Sam Gross
Francisco Massa
Adam Lerer
James Bradbury
Gregory Chanan
Trevor Killeen
Zeming Lin
Natalia Gimelshein
Luca Antiga

title: Pytorch: An imperative style, highperformance deep learning library

authors:
Luis Perez
Jason Wang

title: The effectiveness of data augmentation in image classification using deep learning

venue: arXiv preprint arXiv:1712.04621,

authors:
Senthil Purushwalkam
Abhinav Gupta

title: Demystifying contrastive self-supervised learning: Invariances, augmentations and dataset biases

venue: arXiv preprint arXiv:2007.13916,

authors:
Aditi Raghunathan
Jacob Steinhardt
Percy Liang

title: Certified defenses against adversarial examples

venue: arXiv preprint arXiv:1801.09344,

authors:
Aditi Raghunathan
Sang Michael Xie
Fanny Yang
John C Duchi
Percy Liang

title: Adversarial training can hurt generalization

authors:
Alexander J Ratner
Henry Ehrenberg
Zeshan Hussain
Jared Dunnmon
Christopher Ré

title: Learning to compose domain-specific transformations for data augmentation

authors:
Attila Reiss
Didier Stricker

title: Introducing a new benchmarked dataset for activity monitoring

venue: In 2012 16th International Symposium on Wearable Computers,

authors:
Olaf Ronneberger
Philipp Fischer
Thomas Brox

title: U-net: Convolutional networks for biomedical image segmentation

venue: In International Conference on Medical image computing and computerassisted intervention,

authors:
Mehdi SM Sajjadi
Giambattista Parascandolo
Arash Mehrjou
Bernhard Schölkopf

title: Tempered adversarial networks

venue: arXiv preprint arXiv:1802.04374,

authors:
Christian Szegedy
Wojciech Zaremba
Ilya Sutskever
Joan Bruna
Dumitru Erhan
Ian Goodfellow
Rob Fergus

title: Intriguing properties of neural networks

venue: arXiv preprint arXiv:1312.6199,

authors:
Yonglong Tian
Dilip Krishnan
Phillip Isola

title: Contrastive multiview coding

venue: arXiv preprint arXiv:1906.05849,

authors:
Yonglong Tian
Chen Sun
Ben Poole
Dilip Krishnan
Cordelia Schmid
Phillip Isola

title: What makes for good views for contrastive learning

venue: arXiv preprint arXiv:2005.10243,

authors:
László Tóth
György Kovács
Dirk Van Compernolle

title: A perceptually inspired data augmentation method for noise robust cnn acoustic models

venue: In International Conference on Speech and Computer,

authors:
Toan Tran
Trung Pham
Gustavo Carneiro
Lyle Palmer
Ian Reid

title: A bayesian data augmentation approach for learning deep models, 2017

authors:
Eleni Triantafillou
Tyler Zhu
Vincent Dumoulin
Pascal Lamblin
Utku Evci
Kelvin Xu
Ross Goroshin
Carles Gelada
Kevin Swersky
Pierre-Antoine Manzagol
Hugo Larochelle

title: Metadataset: A dataset of datasets for learning to learn from few examples, 2019

authors:
Riccardo Volpi
Hongseok Namkoong
Ozan Sener
John Duchi
Vittorio Murino
Silvio Savarese

title: Generalizing to unseen domains via adversarial data augmentation, 2018

title: Speech commands: A dataset for limited-vocabulary speech recognition, 2018

authors:
Eric Wong
J. Zico Kolter

title: Learning perturbation sets for robust machine learning, 2020

authors:
Mike Wu
Chengxu Zhuang
Milan Mosse
Daniel Yamins
Noah Goodman

title: On mutual information in contrastive learning for visual representations

venue: arXiv preprint arXiv:2005.13149,

authors:
Zhirong Wu
Yuanjun Xiong
Stella Yu
Dahua Lin

title: Unsupervised feature learning via nonparametric instance-level discrimination

venue: arXiv preprint arXiv:1805.01978,

authors:
Chaowei Xiao
Bo Li
Jun-Yan Zhu
Warren He
Mingyan Liu
Dawn Song

title: Generating adversarial examples with adversarial networks, 2018

authors:
Han Xiao
Kashif Rasul
Roland Vollgraf

title: Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms

venue: arXiv preprint arXiv:1708.07747,

authors:
Sangdoo Yun
Dongyoon Han
Seong Joon Oh
Sanghyuk Chun
Junsuk Choe
Youngjoon Yoo

title: Cutmix: Regularization strategy to train strong classifiers with localizable features

venue: In Proceedings of the IEEE International Conference on Computer Vision,

authors:
Hongyi Zhang
Moustapha Cisse
Yann N Dauphin
David Lopez-Paz

title: mixup: Beyond empirical risk minimization

venue: arXiv preprint arXiv:1710.09412,

authors:
Xinyu Zhang
Qiang Wang
Jian Zhang
Zhao Zhong

title: Adversarial autoaugment

venue: arXiv preprint arXiv:1912.11188,

authors:
Chengxu Zhuang
Alex Lin Zhai
Daniel Yamins

title: Local aggregation for unsupervised learning of visual embeddings

heading: 1 INTRODUCTION

text:
Unsupervised representation learning has made significant recent strides, including in computer vision, where view-based methods have enabled strong performance on benchmark tasks (Wu et al., 2018; Oord et al., 2018; Bachman et al., 2019; Zhuang et al., 2019; Misra & Maaten, 2020; He et al., 2020; Chen et al., 2020a). Views here refer to human-defined data transformations, which target capabilities or invariances thought to be useful for transfer tasks. In particular, in contrastive learning of visual representations, models are trained to maximize the mutual information between different views of an image, including crops, blurs, noise, and changes to color and contrast (Bachman et al.,
2019; Chen et al., 2020a). Much work has investigated the space of possible image views (and their compositions) and understanding their effects on transfer learning (Chen et al., 2020a; Wu et al., 2020; Tian et al., 2019; Purushwalkam & Gupta, 2020).
The fact that views must be hand designed is a significant limitation. While views for image classification have been refined over many years, new views must be developed from scratch for new modalities. Making matters worse, even within a modality, different domains may have different optimal views (Purushwalkam & Gupta, 2020). Previous studies have investigated the properties of good views through the lens of mutual information (Tian et al., 2020; Wu et al., 2020), but a broadly-applicable approach for learning views remains unstudied.
In this work, we present a general method for learning diverse and useful views for contrastive learning. Rather than searching through possible compositions of existing view functions (Cubuk et al., 2018; Lim et al., 2019), which may not be available for many modalities, our approach produces views with a generative model, called the viewmaker network, trained jointly with the encoder network. This flexibility enables learning a broad set of possible view functions, including input-dependent views, without resorting to hand-crafting or expert domain knowledge. The viewmaker network is trained adversarially to create views which increase the contrastive loss of the encoder network. Rather than directly outputting views for an image, the viewmaker instead outputs a stochastic perturbation that is added to the input. This perturbation is projected onto an `p sphere, controlling the effective strength of the view, similar to methods in adversarial robustness. This constrained adversarial training method enables the model to reduce the mutual information between different views while preserving useful input features for the encoder to learn from.
In summary, we contribute:
1. Viewmaker networks: to our knowledge the first modality-agnostic method to learn views for unsupervised representation learning
2. On image data, where expert-designed views have been extensively optimized, our viewmaker-models achieve comparable transfer performance to state of the art contrastive methods while being more robust to common corruptions.
3. On speech data, our method significantly outperforms existing human-defined views on a range of speech recognition transfer tasks.
4. On time-series data from wearable sensors, our model significantly outperforms baseline views on the task of human activity recognition (e.g., cycling, running, jumping rope).

heading: 2 RELATED WORK

text:
Unsupervised representation learning Learning useful representations from unlabeled data is a fundamental problem in machine learning (Pan & Yang, 2009; Bengio et al., 2013). A recently successful framework for unsupervised representation learning for images involves training a model to be invariant to various data transformations (Bachman et al., 2019; Misra & Maaten, 2020), although the idea has much earlier roots (Becker & Hinton, 1992; Hadsell et al., 2006; Dosovitskiy et al., 2014). This idea has been expanded by a number of contrastive learning approaches which push embeddings of different views, or transformed inputs, closer together, while pushing other pairs apart (Tian et al., 2019; He et al., 2020; Chen et al., 2020a;b;c), as well as non-contrastive approaches which do not explicitly push apart unmatched views (Grill et al., 2020; Caron et al., 2020). Related but more limited setups have been explored for speech, where data augmentation strategies are less explored (Oord et al., 2018; Kharitonov et al., 2020).
Understanding and designing views Several works have studied the role of views in contrastive learning, including from a mutual-information perspective (Wu et al., 2020), in relation to specific transfer tasks (Tian et al., 2019), with respect to different kinds of invariances (Purushwalkam & Gupta, 2020), or via careful empirical studies (Chen et al., 2020a). Outside of a contrastive learning framework, Gontijo-Lopes et al. (2020) study how data augmentation aids generalization in vision models. Much work has explored different handcrafted data augmentation methods for supervised learning of images (Hendrycks et al., 2020; Lopes et al., 2019; Perez & Wang, 2017; Yun et al., 2019; Zhang et al., 2017), speech (Park et al., 2019; Kovács et al., 2017; Tóth et al., 2018; Kharitonov et al., 2020), or in feature space (DeVries & Taylor, 2017).
Adversarial methods Our work is related to and inspired by work on adversarial methods, including the `p balls studied in adversarial robustness (Szegedy et al., 2013; Madry et al., 2017; Raghunathan et al., 2018) and training networks with adversarial objectives (Goodfellow et al., 2014; Xiao et al., 2018). Our work is also connected to the vicinal risk minimization principle (Chapelle et al., 2001) and can be interpreted as producing amortized virtual adversarial examples (Miyato et al., 2018). Previous adversarial view-based pretraining methods add adversarial noise on top of existing handcrafted views (Kim et al., 2020) or require access to specific transfer tasks during pretraining (Tian et al., 2020). In contrast, our method is more general: it is neither specialized to a particular downstream task, nor requires neither human-defined view families. Outside of multi-view learning paradigms, adversarial methods have also seen use for representation learning in GANs (Donahue et al., 2016; Donahue & Simonyan, 2019) or in choosing harder negative samples (Bose et al., 2018), as well as for data augmentation (Antoniou et al., 2017; Volpi et al., 2018; Bowles et al., 2018). Adversarial networks that perturb inputs have also been investigated to improve GAN training (Sajjadi et al., 2018) and to remove “shortcut” features (e.g., watermarks) for self-supervised pretext tasks (Minderer et al., 2020).
Learning views Outside of adversarial approaches, our work is related to other studies that seek to learn data augmentation strategies by composing existing human-designed augmentations (Ratner et al., 2017; Cubuk et al., 2018; Zhang et al., 2019; Ho et al., 2019; Lim et al., 2019; Cubuk et al., 2020) or by modeling variations specific to the data distribution (Tran et al., 2017; Wong & Kolter, 2020). By contrast, our method requires no human-defined view functions, does not require first pretraining a generative model, and can generate perturbations beyond naturally-occurring variation observed in the training data (e.g. brightness or contrast), potentially conferring robustness benefits, as we explore in Section 4.3.

heading: 3 METHOD

text:
In contrastive learning, the objective is to push embeddings of positive views (derived from the same input) close together, while pushing away embeddings of negative views (derived from different inputs). We focus mainly on the simple, yet performant, SimCLR contrastive learning algorithm (Chen et al., 2020a), but we also consider a memory bank-based algorithm (Wu et al., 2018) in Section 4. As our method is agnostic to the specific pretraining loss used, it is naturally compatible with other view-based algorithms such as MoCo (He et al., 2020), BYOL (Grill et al., 2020), and SwAV (Caron et al., 2020) by similarly substituting the data transformation pipeline with a viewmaker network.
Formally, given a batch of N pairs of positive views (i, j) the SimCLR loss is
L = 1 2N
NX
k=1 [`(2k 1, 2k) + `(2k, 2k 1)] where `(i, j) = log exp(si,j/⌧)P2N k=1 [k 6=i] exp(si,k/⌧)
and sa,b is the cosine similarity of the embeddings of views a and b.
We generate views by perturbing examples with a viewmaker network V , trained jointly with the main encoder network M . There are three attributes desirable for useful perturbations, each of which motivates an aspect of our method:
1. Challenging: The perturbations should be complex and strong enough that an encoder must develop useful representations to perform the self-supervised task. We accomplish this by generating perturbations with a neural network that is trained adversarially to increase the loss of the encoder network. Specifically, we use a neural network that ingests the input X and outputs a view X + V (X).
2. Faithful: The perturbations must not make the encoder task impossible, being so strong that they destroy all features of the input. For example, perturbations should not be able to zero out the input, making learning impossible. We accomplish this by constraining the perturbations to an `p sphere around the original input. `p constraints are common in the adversarial robustness literature where perturbations are expected to be indistinguishable. In our experiments, we find the best results are achieved with an `1 sphere, which grants the viewmaker a distortion budget that it can spend on a small perturbation for a large part of the input or a more extreme perturbation for a smaller portion.
3. Stochastic: The method should be able to generate a variety of perturbations for a single input, as the encoder objective requires contrasting two different views of an input against each other. To do this, we inject random noise into the viewmaker, such that the model can learn a stochastic function that produces a different perturbed input each forward pass.
Figure 2 summarizes our method. The encoder and viewmaker are optimized in alternating steps to minimize and maximize L, respectively. We use an image-to-image neural network as our viewmaker network, with an architecture adapted from work on style transfer (Johnson et al., 2016). See the Appendix for more details. This network ingests the input image and outputs a perturbation that is constrained to an `1 sphere. The sphere’s radius is determined by the volume of the input tensor times a hyperparameter ✏, the distortion budget, which determines the strength of the applied perturbation. This perturbation is added to the input image and optionally clamped in the case of images to ensure all pixels are in [0, 1]. Algorithm 1 describes this process precisely.
Algorithm 1: Generating viewmaker views Input: Viewmaker network V , C ⇥W ⇥H image X, `1 distortion budget ✏, noise Output: Perturbed C ⇥W ⇥H image X P V (X, ) // generate perturbation

heading: P ✏CWH|P |1 P // project to `1 sphere

text: X X + P // apply perturbation X clamp(X, 0, 1) // clamp (images only)

heading: 4 IMAGES

text:
We begin by applying the viewmaker to contrastive learning for images. In addition to SimCLR (Chen et al., 2020a), we also consider a memory bank-based instance discrimination framework (Wu et al., 2018, henceforth InstDisc).
We pretrain ResNet-18 (He et al., 2015) models on CIFAR-10 (Krizhevsky, 2009) for 200 epochs with a batch size of 256. We train a viewmaker-encoder system with a distortion budget of ✏ = 0.05. We tried distortion budgets ✏ 2 {0.1, 0.05, 0.02} and found 0.05 to work best; however, we anticipate that further tuning would yield additional gains. As we can see in Figure 1, the learned views are diverse, consisting of qualitatively different kinds of perturbations and affecting different parts of the input. We compare the resulting encoder representations with a model trained with the expert views used for SimCLR, comprised of many human-defined transformations targeting different kinds of invariances useful for image classification: cropping-and-resizing, blurring, horizontal flipping, color dropping, and shifts in brightness, contrast, saturation, and hue (Chen et al., 2020a).

heading: 4.1 TRANSFER RESULTS ON IMAGE CLASSIFICATION TASKS

text:
We evaluate our models on CIFAR-10, as well as eleven transfer tasks including MetaDataset (Triantafillou et al., 2019), MSCOCO (Lin et al., 2014), MNIST (LeCun et al., 1998), and FashionMNIST (Xiao et al., 2017). We use the standard linear evaluation protocol, which trains a logistic
regression on top of representations from a frozen model. We apply the same views as in pretraining, freezing the final viewmaker when using learned views; we apply no views during validation. Table 1 shows our results, indicating comparable overall performance with SimCLR and InstDisc, all without the use of human-crafted view functions. This performance is noteworthy as our `1 views cannot implement cropping-and-rescaling, which was shown to be the most important view function in Chen et al. (2020a). We speculate that the ability of the viewmaker to implement partial masking of an image may enable a similar kind of spatial information ablation as cropping.

heading: 4.1.1 COMPARISON TO RANDOM `1 NOISE

text: Is random noise sufficient to produce domain-agnostic views? To assess how important adversarial training is to the quality of the learned representations, we perform an ablation where we generate views by adding Gaussian noise normalized to the same ✏ = 0.05 budget as used in the previous section. Transfer accuracy on CIFAR-10 is significantly hurt by this ablation, reaching 52.01% for a SimCLR model trained with random noise views compared to 84.50% for our method, demonstrating the importance of adversarial training to our method.

heading: 4.1.2 THE IMPORTANCE OF INTER-PATCH MUTUAL INFORMATION AND CROPPING VIEWS

text:
Cropping-and-resizing has been identified as a crucial view function when pretraining on ImageNet (Chen et al., 2020a). However, what properties of a pretraining dataset make cropping useful? We hypothesize that such a dataset must have images whose patches have high mutual information. In other words, there must be some way for the model to identify that different patches of the same image come from the same image. While this may be true for many object or scene recognition datasets, it may be false for other important pretraining datasets, including medical or satellite imagery, where features of interest are isolated to particular parts of the image.
To investigate this hypothesis, we modify the CIFAR-10 dataset to reduce the inter-patch mutual information by replacing each 16x16 corner of the image with the corner from another image in the training dataset (see Figure 3 for an example). Thus, random crops on this dataset, which we call CIFAR-10-Corners, will often contain completely unrelated information. When pretrained on CIFAR-10-Corners, expert views achieve 63.3% linear evaluation accuracy on the original CIFAR10 dataset, while viewmaker views achieve 68.8%. This gap suggests that viewmaker views are less reliant on inter-patch mutual information than the expert views.

heading: 4.2 COMBINING VIEWMAKER AND HANDCRAFTED VIEWS

text:
Can viewmakers improve performance in cases where some useful handcrafted views have already been identified? Chen et al. (2020a) show that views produced through cropping are significantly improved by a suite of color-based augmentations, which they argue prevents the network from relying solely on color statistics to perform the contrastive task. Here, we show that viewmaker networks also enable strong gains when added on top of cropping and horizontal flipping views when pretraining on CIFAR-10—without any domain-specific knowledge. Alone, this subset of
handcrafted augmentations achieves 73.2% linear evaluation accuracy on CIFAR-10. Combining these views with learned viewmaker perturbations (✏ = 0.05) achieves 83.1%.1 This suggests that viewmakers can significantly improve representation learning even in cases where some domainspecific views have already been developed.

heading: 4.3 ROBUSTNESS TO COMMON CORRUPTIONS

text:
Image classification systems should behave robustly even when the data distribution is slightly different from that seen during training. Does using a viewmaker improve robustness against common types of corruptions not experienced at train time? To answer this, we evaluate both learned views, expert views, and their composition on the CIFAR-10-C dataset (Hendrycks & Dietterich, 2019), which assesses robustness to corruptions like snow, pixelation, and blurring. In this setting, corruptions are applied only at test time, evaluating whether the classification system is robust to some types of corruptions to which humans are robust.
When considering methods in isolation, SimCLR augmentations result in less of an accuracy drop from clean to corrupted data compared to our learned views, as shown in Table 4a. This gap is expected, as the expert views overlap significantly with the CIFAR-10-C corruptions: both include blurring, brightness, and contrast transformations. Interestingly, however, when we train a viewmaker network while also applying expert augmentations (“Combined,” Table 4a), we can further improve the robust accuracy, with notable gains on noise and glass blur corruptions (Figure 4b). This is noteworthy, as our learned views have no explicit overlap with the CIFAR-10-C corruptions, unlike the expert augmentations.2 In the Combined setting, we use a distortion budget of ✏ = 0.01,
1We did not see additional gains from using viewmakers on top of the full, well-optimized set of SimCLR augmentations.
2We do notice a smaller decline in contrast corruption accuracy, possibly due to interactions between changing pixel magnitudes and the `p constraint.
which we find works better than ✏ = 0.05, likely because combining the two augmentations at their full strength would make the learning task too difficult.
These results suggest that learned views are a promising avenue for improving robustness in selfsupervised learning models.

heading: 5 SPEECH

text: Representation learning on speech data is an emerging and important research area, given the large amount of available unlabeled data and the increasing prevalence of speech-based human-computer interaction (Latif et al., 2020). However, compared to images, there is considerably less work on self-supervised learning and data augmentations for speech data. Thus, it is a compelling setting to investigate whether viewmaker augmentations are broadly applicable across modalities.

heading: 5.1 SELF-SUPERVISED LEARNING SETUP

text: We adapt the contrastive learning setup from SimCLR (Chen et al., 2020a). Training proceeds largely the same as for images, but the inputs are 2D log mel spectrograms. We consider both view functions applied in the time-domain before the STFT, including noise, reverb, pitch shifts, and changes in loudness (Kharitonov et al., 2020), as well as spectral views, which involve masking or noising different parts of the spectrogram (Park et al., 2019). To generate learned views, we pass the spectrogram as input to the viewmaker. We normalize the spectrogram to mean zero and variance one before passing it through the viewmaker, and do not clamp the resulting perturbed spectrogram. See the Appendix for more details. We train on the Librispeech dataset (Panayotov et al., 2015) for 200 epochs, and display some examples of learned views in the Appendix.

heading: 5.2 SPEECH CLASSIFICATION RESULTS

text:
We evaluate on three speech classification datasets: Fluent Speech Commands (Lugosch et al., 2019), Google Speech Commands (Warden, 2018), and spoken digit classification (Becker et al., 2018), as well as speaker classification on VoxCeleb (Nagrani et al., 2017) and Librispeech (Panayotov et al., 2015), all using the linear evaluation protocol for 100 epochs. In Table 2, we report results with both the same distortion budget ✏ = 0.05 as in the image domain, as well as a larger ✏ = 0.1, for comparison. Both versions significantly outperform the preexisting waveform and spectral augmentations, with a +9 percentage point improvement on average for the ResNet-18 (✏ = 0.05) viewmaker model over the best expert views. The gains for real-world tasks such as command identification are compelling. One notable exception is the task of LibriSpeech speaker identification. Since LibriSpeech is the same dataset the model was pretrained on, and this effect is not replicated on VoxCeleb1, the other speaker classification dataset, we suspect the model may be picking up on dataset-specific artifacts (e.g. background noise, microphone type) which may make the speaker
ID task artificially easier. An interesting possibility is that the worse performance of viewmaker views may result from the model being able to identify and ablate such spurious correlations in the spectrograms.

heading: 6 WEARABLE SENSOR DATA

text: To further validate that our method for learning views is useful across different modalities, we consider time-series data from wearable sensors. Wearable sensor data has a broad range of applications, including health care, entertainment, and education (Lara & Labrador, 2012). We specifically consider whether viewmaker views improve representation learning for the task of human activity recognition (HAR), for example identifying whether a user is jumping rope, running, or cycling.

heading: 6.1 SELF-SUPERVISED LEARNING SETUP

text:
We consider the Pamap2 dataset (Reiss & Stricker, 2012), a dataset of 12 different activities performed by 9 participants. Each activity contains 52 different time series, including heart rate, accelerometer, gyroscope, and magnetometer data collected from sensors on the ankle, hand, and chest (all sampled at 100Hz, except heart rate, which is sampled at approximately 9Hz). We linearly interpolate missing data, then take random 10s windows from subject recordings, using the same train/validation/test splits as prior work (Moya Rueda et al., 2018). To create inputs for our model, we generate a multi-channel image composed of one 32x32 log spectrogram for each sensor timeseries window. Unlike speech data, we do not use the mel scale when generating the spectrogram. We then normalize the training and validation datasets by subtracting the mean and then dividing by the standard deviation of the training dataset.
We train with both our learned views and the spectral views (Park et al., 2019) that were most successful in the speech domain (for multi-channel spectral masking, we apply the same randomly chosen mask to all channels). We also compare against a variant of these views with spectrogram noise removed, which we find improves this baseline’s performance.

heading: 6.2 SENSOR-BASED ACTIVITY RECOGNITION RESULTS

text: We train a linear classifier on the frozen encoder representations for 50 epochs, reporting accuracy on the validation set. We sample 10k examples for each training epoch and 50k examples for validation. Our views significantly outperform spectral masking by 12.8 percentage points when using the same ✏ = 0.05 as image and speech, and by 16.7 points when using a larger ✏ = 0.5 (Table 3). We also find that a broad range of distortion budgets produces useful representations, although overly-aggressive budgets prevent learning (Table 3). These results provide further evidence that our method for learning views has broad applicability across different domains.

heading: 6.3 SEMI-SUPERVISED EXPERIMENTS

text:
An especially important setting for self-supervised learning is domains where labeled data is scarce or costly to acquire. Here, we show that our method can enable strong performance when labels for only a single participant (Participant 1) out of seven are available. We compare simple supervised learning on Participant 1’s labels against linear evaluation of our best pretrained model, which was
trained on unlabeled data from all 7 participants. The model architectures and training procedures are otherwise identical to the previous section. As Figure 4 shows, pretraining with our method on unlabeled data enables significant gains over pure supervised learning when data is scarce, and even slightly outperforms the hand-crafted views trained on all 7 participants (cf. Table 3).

heading: 7 CONCLUSION

text: We introduce a method for learning views for unsupervised learning, demonstrating its effectiveness through strong performance on image, speech, and wearable sensor modalities. Our novel generative model—viewmaker networks—enables us to efficiently learn views as part of the representation learning process, as opposed to relying on domain-specific knowledge or costly trial and error. There are many interesting avenues for future work. For example, while the `1 constraint is simple by design, there may be other kinds of constraints that enable richer spaces of views and better performance. In addition, viewmaker networks may find use in supervised learning, for the purposes of data augmentation or improving robustness. Finally, it is interesting to consider what happens as the viewmaker networks increase in size: do we see performance gains or robustnessaccuracy trade-offs (Raghunathan et al., 2019)? Ultimately, our work is a step towards more general self-supervised algorithms capable of pretraining on arbitrary data and domains.

heading: ACKNOWLEDGEMENTS

text: We would like to thank Dan Yamins, Chengxu Zhuang, Shyamal Buch, Jesse Mu, Jared Davis, Aditi Raghunathan, Pranav Rajpurkar, Margalit Glasgow, and Jesse Michel for useful discussions and comments on drafts. AT is supported by an Open Phil AI Fellowship. MW is supported by the Stanford Interdisciplinary Graduate Fellowship as the Karr Family Fellow.