Year: 2019

Abstract:
Deep reinforcement learning requires a heavy price in terms of sample efficiency and overparameterization in the neural networks used for function approximation. In this work, we employ tensor factorization in order to learn more compact representations for reinforcement learning policies. We show empirically that in the low-data regime, it is possible to learn online policies with 2 to 10 times less total coefficients, with little to no loss of performance. We also leverage progress in second order optimization, and use the theory of wavelet scattering to further reduce the number of learned coefficients, by foregoing learning the topmost convolutional layer filters altogether. We evaluate our results on the Atari suite against recent baseline algorithms that represent the state-of-the-art in data efficiency, and get comparable results with an order of magnitude gain in weight parsimony.

text: Deep reinforcement learning requires a heavy price in terms of sample efficiency and overparameterization in the neural networks used for function approximation. In this work, we employ tensor factorization in order to learn more compact representations for reinforcement learning policies. We show empirically that in the low-data regime, it is possible to learn online policies with 2 to 10 times less total coefficients, with little to no loss of performance. We also leverage progress in second order optimization, and use the theory of wavelet scattering to further reduce the number of learned coefficients, by foregoing learning the topmost convolutional layer filters altogether. We evaluate our results on the Atari suite against recent baseline algorithms that represent the state-of-the-art in data efficiency, and get comparable results with an order of magnitude gain in weight parsimony.

heading: 1 INTRODUCTION
text: The successes of reinforcement learning (thereafter ’RL’), and specifically deep RL, come at a heavy computational price. It is well known that achieving human-level performance in domains such as Atari (Sutton & Barto, 2018; Mnih et al., 2013; Hessel et al., 2017) requires hundreds of millions of frames of environment interaction. As such, the problem of sample efficiency in RL is of critical importance. Several tracks of concurrent research are being investigated, and have reduced by orders of magnitude the number of environment interactions required for good performance beyond the previous benchmark of biologically-inspired episodic control methods (Blundell et al., 2016; Pritzel et al., 2017) to a couple hours of human gameplay time (van Hasselt et al., 2019; Kaiser et al., 2019).
However, while the data-efficiency of RL methods has seen recent drastic performance gains, the function approximators they use still require millions of learned weights, potentially still leaving them heavily overparameterized. Independently motivated by biological facts like the behavioural readiness of newborn animals, several authors (Gaier & Ha, 2019; Cuccu et al., 2018; Wang et al., 2019) have recently looked at doing away with learning so many weights for RL tasks. Smaller networks not only train faster, but may yet offer another avenue for gains in the form of better generalization (Zhang et al., 2016). Recent work from Gaier & Ha (2019) studies the effect of inductive bias of neural architectures in RL ; they forego training altogether, but transfer networks that only obtain ’better than chance performance on MNIST’. In similar fashion, Wang et al. (2019) investigate the effect of random projections in the restricted setting of imitation learning. Finally, Cuccu et al. (2018) manage human-level performance on the Atari suite using a separate dictionary learning procedure for their features, bypassing the usual end-to-end learning paradigm. The perspective of neural architecture search applied to RL appears difficult, if not computationally inextricable.
Concurrently, the study of biologically-inspired models of learning has exhibited two mathematical characterizations that might be critical in explaining how biological learning takes place so efficiently. First, the low-rank properties of learned perceptual manifolds (Chung et al., 2018; 2016) are giving rise to a rich theory borrowing from statistical physics. Second, another well known line of work has identified Gabor filters, and more generally wavelet filter-like structures, in the actual visual cortex of animals (Jones & Palmer, 1987), and linked those to sparsity-promoting methods and dictionary learning (Olshausen & Field, 1996; 1997; Hyvärinen & Hoyer, 2001). But these breakthroughs have not, so far, been reflected as inductive priors in the shape of modifications in deep RL neural networks architectures, which remain fairly fixed on the Atari domain.
Therefore the following questions remain: how parsimonious do function approximators in RL need to be, in order to maintain good performance? And can we be at once sample-efficient and weightefficient ? In this work, we turn to the mathematical theories of tensor factorization (Cichocki et al., 2009), second-order optimization (Amari, 1998; Martens & Grosse, 2015) and wavelet scattering (Mallat, 2011) to answer this question positively and empirically, in a model-free setting.
We propose to use these methods in order to save weights and therefore favour convergence of policies:
• We replace dense, fully-connected layers with tensor regression layers. • Optionally, we replace the topmost layer in the convolutional architecture with a scattering
layer; the deeper convolutional layers are left untouched.
• The (positive) impact of second-order optimization is also evaluated.
To the best of our knowledge, this is the first time those fields have been combined together in this context, and that tensor factorization is applied to deep RL.

heading: 2 BACKGROUND & RELATED WORK

heading: 2.1 DEEP REINFORCEMENT LEARNING
text: We consider the standard Markov Decision Process framework as in Sutton & Barto (2018). This setting is characterised by a tuple 〈S,A, T,R, γ〉, where S is a set of states, A a set of actions, R a reward function that is the immediate, intrinsic desirability of a certain state, T a transition dynamics and γ ∈ [0, 1] a discount factor. The purpose of the RL problem is to to find a policy π, which represents a mapping from states to a probability distribution over actions, that is optimal, i.e., that maximizes the expected cumulative discounted return ∑∞ k=0 γ
kRt+k+1 at each state st ∈ S. In Q-learning, the policy is given implicitly by acting greedily or -greedily with respect to learned action-value functions qπ(s, a), that are learned following the Bellman equation. In deep Q-learning, qθ becomes parameterized by the weights θ of a neural network and one minimizes the expected Bellman loss :
E ( Rt+1 + γt+1 max
a′ qθ (St+1, a
′)− qθ (St, At) )2
In practice, this is implemented stochastically via uniform sampling of transitions in an experience replay buffer, as is done in the seminal paper Mnih et al. (2013). Several algorithmic refinements to that approach exist. First, Double Q-learning (van Hasselt et al., 2015) proposes to decouple learning between two networks in order to alleviate the Q-value overestimation problem. Second, dueling Q-networks (Wang et al., 2015) explicitly decompose the learning of an action-value function qθ(s, a) as the sum of an action-independent state-value, much like what is traditionally done in policy gradient methods (Sutton & Barto, 2018), implemented via a two-headed neural network architecture. Finally, prioritized RL (Schaul et al., 2015) proposes to replace the uniform sampling of transitions in the experience replay buffer with importance sampling, by prioritizing those transitions that present the most Bellman error (those transitions that are deemed the most ’surprising’ by the agent). Fortunato et al. (2017) use extra weights to learn the variance of the exploration noise in a granular fashion, while Bellemare et al. (2017) propose to learn a full distribution of action-values for each action and state. Combined, those methods form the basis of the Rainbow algorithm in Hessel et al. (2017).

heading: 2.2 TENSOR FACTORIZATION
text: Here we introduce notations and concepts from the tensor factorization literature. An intuition is that the two main decompositions below, CP and Tucker decompositions, can be understood as multilinear algebra analogues of SVD or eigendecomposition.
CP decomposition. A tensor X ∈ RI1×I2×···×IN , can be decomposed into a sum of R rank-1 tensors, known as the Canonical-Polyadic decomposition, where R is the rank of the decomposition. Its purpose is to find vectors u(1)k ,u (2) k , · · · ,u (N) k , for k = [1 . . . R], as well as a vector of weights
λ ∈ RR such that:
X = R∑ k=1 λku (1) k ◦ u (2) k ◦ · · · ◦ u (N) k︸ ︷︷ ︸
rank-1 components
Tucker decomposition. A tensor X ∈ RI1×I2×···×IN , can be decomposed into a low rank approximation including a core G ∈ RR1×R2×···×RN and a set of projection factors ( U(0), · · · ,U(N−1) ) ,
with U(k) ∈ RRk,Îk , k ∈ (0, · · · , N − 1) that, when projected along the corresponding dimension of the core, reconstruct the full tensor X . The tensor in its decomposed form can be written: X = G ×1 U(1) ×2 U(2) × · · · ×N U(N) = î G;U(1), · · · ,U(N) ó Tensor regression layer. For two tensors X ∈ RK1×···×Kx×I1×···×IN and Y ∈ RI1×···×IN×L1×···×Ly , we denote by 〈X ,Y〉N ∈ RK1×···×Kx×L1×···×Ly the contraction of X by Y along their N last modes; their generalized inner product is
〈X ,Y〉N = I1∑ i1=1 I2∑ i2=1 · · · IN∑ in=1 X...,i1,i2,...,inYi1,i2,...,in,...
This makes it possible to define a tensor regression layer (Kossaifi et al., 2017b) that is differentiable and learnable end-to-end by gradient descent. Let us denote by X ∈ RI1×I2×···×IN the input activation tensor for a sample and y ∈ RIN the label vector. A tensor regression layer estimates the regression weight tensorW ∈ RI1×I2×···×IN under a low-rank decomposition. In the case of a Tucker decomposition (as per our experiments) with ranks (R1, · · · , RN ), we have :
y = 〈X ,W〉N + b withW = G ×1 U(1) ×2 U(2) · · · ×N U(N)
as G ∈ RR1×···×RN , U(k) ∈ RIk×Rk for each k in [1 . . . N ] and U(N) ∈ R1×RN .

heading: 2.3 WAVELET SCATTERING
text: The wavelet scattering transform was originally introduced by Mallat (2011) and Bruna & Mallat (2012) as a non-linear extension to the classical wavelet filter bank decomposition. Its principle is as follows. Denoting by x ~ y[n] the 2-dimensional, circular convolution of two signals x[n] and y[n], let us assume that we have pre-defined two wavelet filter banks available ¶ ψ
(1) λ1 [n] © λ1∈Λ1¶
ψ (2) λ2 [n] © λ2∈Λ2 , with λ1 and λ2 two frequency indices. These wavelet filters correspond to high frequencies, so we also give ourselves the data of a lowpass filter φJ [n]. Finally, and by opposition to traditional linear wavelet transforms, we also assume a given nonlinearity ρ(t). Then the scattering transform is given by coefficients of order 0,1, and 2, respectively :
S0x[n] = x~ φJ [n] S1x [n, λ1] = ρ Ä x~ ψ(1)λ1 ä ~ φJ [n] λ1 ∈ Λ1
S2x [n, λ1, λ2] = ρ Ä ρ Ä x~ ψ(1)λ1 ä ~ ψ(2)λ2 ä ~ φJ [n] λ1 ∈ Λ1, λ2 ∈ Λ2 (λ1)
This can effectively be understood and implemented as a two-layer convolutional neural network whose weights are not learned but rather frozen and given by the coefficients of wavelets ψ and φ (with Gabor filters as a special case (Mallat, 1998)). The difference with traditional filter banks comes from the iterated modulus/nonlinear activation function applied at each stage, much like in traditional deep learning convolutional neural networks. The generic mathematical definition involves order n iterated scatterings, in the vein of Six above, but sometimes restricts nonlinearity ρ to be a modulus function | · |. In practice, the potential of scattering transforms to accelerate learning by providing ready-made convolutional layers has been investigated in Oyallon et al. (2013) and Oyallon et al. (2018) and is a subject of active ongoing research. Scattering will be the second of our weight-saving methods.

heading: 2.4 SECOND ORDER OPTIMIZATION WITH K-FAC
text: While stochastic gradient descent is usually performed purely from gradient observations derived from auto-differentiation, faster, second order optimization methods first multiply the weights’ θ gradient vector∇θ by a preconditioning matrix, yielding the weight update θn+1 ← θn− ηG−1n ∇θ, with η a step size. In the case of second order methods, the matrixG−1n is chosen to act as a tractable iterative approximation to the inverse Hessian or Empirical Fisher Information Matrix (Amari, 1998) of the neural network model in question. Kronecker-factored approximate curvature or K-FAC (Martens & Grosse, 2015) enforces a Kronecker decomposition of the type G = A ⊗ B, with A and B being smaller, architecture-dependent matrices. Unlike the above methods, K-FAC has been applied as a plug-in in the deep RL literature and shown to promote both anytime convergence properties as well as terminal accuracies (Wu et al., 2017).

heading: 3 OBSERVATIONS AND METHODS

heading: 3.1 EXPLORING TRAINED AGENTS
text: Stability of trained dense layers eigenvalues. In order to assess experimentally if tensor factorization can make sense in RL, we investigate the eigenvalues of the dense layers of a deep RL agent. Unlike the traditional supervised learning setting, the input data distribution to RL function approximators shifts as the agent explores its environment; as such, concentration properties of the eigenvalues of the linear layers cannot be guaranteed all the way throughout training. Since conditioning techniques such as batch normalization (Ioffe & Szegedy, 2015) are rarely used in deep RL, this is all the more important. Our experiments (see figure 1) show that the distribution of eigenvalues does not seem to widen significantly, at least during the initial phases of training we care about.
Furthermore and interestingly, it does not seem to deviate significantly from the one observed at initialization. All together, this suggests there might be some merit in learning low-rank policies.

heading: 3.2 ARCHITECTURAL MODIFICATIONS
text: Baseline. We then proceed to build upon the well-known Rainbow algorithm (Hessel et al., 2017). Rainbow uses a fairly standard shallow convolutional architecture like the seminal DQN paper of Mnih et al. (2013), and is an oft-cited baseline on the Atari suite. In spite of its performance, Rainbow often requires dozens of millions of a single game’s frames in order to perform well. Very recently, a ’data-efficient’ efficient version of Rainbow has been proposed by van Hasselt et al. (2019), with a view to match or beat the latest state-of-the-art results achieved by model-based RL methods. This is achieved with no major change in network architecture, but via a selection of mildly handtuned hyperparameters favouring efficiency against wall-clock running time (see appendix). We do take this as a baseline.
Changes. We modify the architecture of the neural network function approximators used, in accordance with the principles described above, combining them to reflect inductive biases promoting fewer learnable parameters:
• We replace the fully-connected, linear layers used in Rainbow and data-efficient Rainbow with tensor regression layers (Kossaifi et al., 2017b) in order to learn low-rank policies (ranks in appendix).
• We use either the K-FAC second order stochastic optimizer, or the standard ADAM optimizer (Kingma & Ba, 2014). Optimization with K-FAC yields better results ceteris paribus and therefore works to counter performance loss due to using fewer weights.
• We combine the two methods with various rank (and therefore weight compression) ratios, targetting sensible compression ratios guided by deep learning intuition; and evaluate those on the same subset of Atari games as both van Hasselt et al. (2019); Kaiser et al. (2019).
• When possible, we replace the first convolutional layer in the approximating neural network with a scattering layer for further gains in terms of learnable weights. In that way, we investigate the impact of not actually learning one of the convolutional layer weights. Deep RL, paradoxically, tends to use shallow network architectures with two or three convolutional layers only, which makes it very fit for scattering methods. But since learned convolutional features discriminate well between high and low rewards, the promise of fully unsupervised scattering layers seems remote, without resorting to further ad-hoc methods like dictionary learning. Therefore we limit ourselves to one-step (one single layer) scatterings.
This is illustrated in figure 2. We then proceed to evaluate the merit of these changes.

heading: 4 EXPERIMENTAL RESULTS

heading: 4.1 PRIORITIZED TENSORIZED DQN
text: Proof of concept. We begin with showing proof of concept on the simple Pong Atari game. Our experimental setup consists in our own implementation of prioritized double DQN as a baseline (Schaul et al., 2015; van Hasselt et al., 2015). We replaced the densely connected layer of the original DQN architecture with a tensor regression layer implementing Tucker decomposition for different Tucker ranks, yielding different network compression factors.
Qualitative behaviour. First results, both in terms of learning performance and compression factor, can be seen in figure 3. Our two main findings are that first, the final performance of the agent remains unaffected by the tensor factorization, even with high compression rates - with respect to all network weights - of five times. In line with intuition, larger compression rates do however cause more delays in learning. Second, tensor factorization negatively affects stability during training - in tough compression regimes, the plateauing phases of learning curves feature occasional noisy drawdowns, illustrating the increased difficulty of learning, as seen in figure 4. Interestingly, approximation errors incurred by tensor regression noise do sometimes have poor consequences illustrated by those drawdowns, but overall seem to behave as additional exploration noise.

heading: 4.2 DATA-EFFICIENT RAINBOW ON ATARI
text: Evaluation protocol. For all our Atari experiments, we used OpenAI Gym (Brockman et al., 2016), and a combination of PyTorch (Paszke et al., 2017), TensorLy (Kossaifi et al., 2016) and Kymatio (Andreux et al., 2018) for auto-differentiation. We evaluated our agents in the low-data regime of
100,000 steps, on half the games, with 3 different random seeds for reproducibility (Henderson et al., 2017), taking the data-efficient Rainbow agent (van Hasselt et al., 2019) as our baseline. Our specific hyperparameters are described in appendix. We report our results in tables 1 and 2.
Table 1 shows proof of concept of the online learning of low-rank policies, with a loss of final performance varying in proportion to the compression in the low-rank linear layers, very much like in the deep learning literature (Kossaifi et al., 2017a;b). The number of coefficients in the original data-efficient Rainbow is of the order of magnitude of 1M and varies depending on the environment and its action-space size. The corresponding tensor regression layer ranks are in appendix, and chosen to target 400k, 200k and 100k coefficients respectively. While individual game results tend to decrease monotonously with increasing compression, we observe that they are noisy due to the nature of exploration in RL, and average scores reported correspond to the intuition that performance seems to decrease fast after a certain overparameterization threshold is crossed. To take this noisy character into account, we take care to be conservative and report the average of the final three episodes of the learned policy after 80k, 90k and 100k steps, respectively.
Denoised baseline. So as to not muddy the discussion and provide fair baselines, we do report on the NoisyNet (Fortunato et al., 2017) ablation of Rainbow (’Denoised’ columns), as the NoisyLinear layer doubles up the number of coefficients required and actually performs worse in our experiments. Its principle is that the variance of exploration noise represents a criticial tradeoff for performance (too little and one stalls, too much and one risks catastrophic updates), so it is sensible to treat it as a parameter to learn. However, in order to decouple both factors of overparametrization and exploration in the discussion of deep RL performance, we simply use a fixed exploration schedule. This denoised exploration baseline is an ablation we can then compare our tensorized methods to.
Second-order optimization. We then proceed to assess the impact of second-order optimization in architecture by substituting ADAM optimization for K-FAC, and introducing scattering, in table 2. In spite of our conservative reporting, the efficiency boost from using a second order scheme more than makes up for low-rank approximation error (109% performance) with five times less coefficients than van Hasselt et al. (2019), and learning with a full order of magnitude less coefficients (98% performance) is made possible by the combination of K-FAC and TRL. The results however do show a sharp drop in average performance when scattering is added. Interestingly enough some games perform relatively very well with that method, simultaneously showing proof of viability and of additional work required.

heading: 5 CONCLUSION
text: We have demonstrated that in the low-data regime, it is possible to leverage biologically plausible characterizations of experience data (namely low-rank properties and wavelet scattering separability) to exhibit architectures that learn policies with an order of magnitude less weights than current state-of-the-art baselines, essentially without loss of performance, and in an online fashion. In particular, this provides a compelling alternative to methods like policy distillation (Rusu et al., 2015; Czarnecki et al., 2019). We do hope that this will lead to even further progress towards sample efficiency and speedy exploration methods. Further work will, first, focus on thorough evaluation and research of scattering architectures in order to achieve further gains, and second investigate additional, orthogonal biologically-friendly research directions such as promoting sparsity via, for instance, L1 regularization. Finally, we are excited by the potential of tensor factorization to offer shared core tensors for policies in multi-task and meta-learning.

heading: A APPENDIX
text: HYPERPARAMETERS AND REPRODUCIBILITY
Hyperparameters are as follows. First, our specific architecture-modified hyperparameters:
Furthermore, we mirror the Data-Efficient Rainbow van Hasselt et al. (2019) baseline:
Our codebase is available on request.

Review summaries:
This paper investigates a list of methods to reduce the number of weights for deep RL architecture under the low-data regime. These methods include tensor regression, wavelet scattering, as well as second-order optimization (K-FAC). The experiments on the Atari games shows that by using tensor regression to replace the dense layer of the neural nets and using K-FAC for the optimization, one can reduce around 10 times of parameters without losing too much of performance.