Learning Sparse Representations in Reinforcement Learning with Sparse Coding

1 Introduction

For tasks with large state or action spaces, where tabular representations are not feasible, reinforcement learning algorithms typically rely on function approximation. Whether they are learning the value function, policy or models, the success of function approximation techniques hinges on the quality of the representation. Typically, representations are hand-crafted, with some common representations including tile-coding, radial basis functions, polynomial basis functions and Fourier basis functions [Sutton, 1996; Konidaris et al., 2011]. Automating feature discovery, however, alleviates this burden and has the potential to significantly improve learning.

Representation learning techniques in reinforcement learning have typically drawn on the large literature in unsupervised and supervised learning. Common approaches include feature selection, including ℓ1 regularization on the value function parameters [Loth et al., 2007; Kolter and Ng, 2009; Nguyen et al., 2013] and matching pursuit [Parr et al., 2008; Painter-Wakefield and Parr, 2012]; basis-function adaptation approaches [Menache et al., 2005; Whiteson et al., 2007];

∗These authors contributed equally.

Martha White

Dept. of Computer Science Indiana University Bloomington, IN, USA martha@indiana.edu

instance-based approaches, such as locally weighted regression [Atkeson and Morimoto, 2003], sparse distributed memories [Ratitch and Precup, 2004], proto-value functions [Mahadevan and Maggioni, 2007] and manifold learning techniques [Mahadevan, 2009]; and neural network approaches, including more standard feedforward neural networks [Coulom, 2002; Riedmiller, 2005; Mnih et al., 2015] as well as random representations [Sutton andWhitehead, 1993], linear threshold unit search [Sutton and Barto, 2013], and evolutionary algorithms like NEAT [Stanley and Miikkulainen, 2002].

Surprisingly, however, there has been little investigation into using sparse coding for reinforcement learning. Sparse coding approaches have been developed to learn MDP models for transfer learning [Ammar et al., 2012]; outside this work, however, little has been explored. Nonetheless, such sparse coding representations have several advantages, including that they naturally enable local models, are computationally efficient to use, are much simpler to train than more complicated models such as neural networks and are biologically motivated by the observed representation in the mammalian cortex [Olshausen and Field, 1997].

In this work, we develop a principled sparse coding objective for policy evaluation. In particular, we formulate a joint optimization over the basis and the value function parameters, to provide a supervised sparse coding objective where the basis is informed by its utility for prediction. We highlight the importance of using the Bellman error or mean-squared return error for this objective, and discuss how the projected Bellman error is not suitable. We then show that, despite being a nonconvex objective, all local minima are global minima, under minimal conditions. We avoid the need for careful initialization strategies needed for previous optimality results for sparse coding [Agarwal et al., 2014; Arora et al., 2015], using recent results for more general dictionary learning settings [Haeffele and Vidal, 2015; Le and White, 2017], particularly by extending beyond smooth regularizers using Γ-convergence. Using this insight, we provide a simple alternating proximal gradient algorithm and demonstrate the utility of learning supervised sparse coding representations versus unsupervised sparse coding and a variety of tile-coding representations.

2 Background

In reinforcement learning, an agent interacts with its environment, receiving observations and selecting actions to maximize a scalar reward signal provided by the environment. This interaction is usually modeled by a Markov decision process (MDP). An MDP consists of (S, A, P, R) where S is the set of states; A is a finite set of actions; P : S ×A× S → [0, 1], the transition function, which describes the probability of reaching a state s′ from a given state and action (s, a); and finally the reward function R : S × A × S → R, which returns a scalar value for transitioning from state-action (s, a) to state s′ . The state of the environment is said to be Markov if Pr(st+1|st, at) = Pr(st+1|st, at, . . . , s0, a0).

One important goal in reinforcement learning is policy evaluation: learning the value function for a policy. A value function Vπ : S → R approximates the expected return. The return Gt from a state st is the total discounted future reward, discounted by γ ∈ [0, 1), for following policy π : S×A → [0, 1]

where Vπ (st) is the expectation of this return from state st. This value function can also be thought of as a vector of values Vπ ∈ R|S| satisfying the Bellman equation

Given the reward function and transition probabilities, the solution can be analytically obtained: Vπ = (I - γPπ )-1 rπ .

In practice, however, we likely have a prohibitively large state space. The typical strategy in this setting is to use function approximation to learn Vπ (s) from a trajectory of samples: a sequence of states, actions, and rewards s0, a0, r0, s1, a1, r1, s2, r2, a2 . . ., where s0 is drawn from the start-state distribution, st+1 ∼ P(·|st, at) and at ∼ π(·|st). Commonly, a linear function is assumed, Vπ (s) ≈ φ(s)⊤w for w ∈ Rk a parameter vector and φ : S → Rk a feature function describing states. With this approximation, however, typically we can no longer satisfy the Bellman equation in (1), because there may not exist a w such that Φw equals rπ + γPπΦw for Φ ∈ R|S|×k . Instead, we focus on minimizing the error to the true value function.

Reinforcement learning algorithms, such as temporal difference learning and residual gradient, therefore focus on finding an approximate solution to the Bellman equation, despite this representation issue. The quality of the representation is critical to accurately approximating Vπ with Φw, but also balancing compactness of the representation and speed of learning. Sparse coding, and sparse representations, have proven successful in machine learning and in reinforcement learning, particularly as fixed bases, such as tile coding, radial basis functions and other kernel representations. A natural goal, therefore, and the one we explore in this work, is to investigate learning these sparse representations automatically.

3 Sparse Coding for Reinforcement Learning

In this section, we formalize sparse coding for reinforcement learning as a joint optimization over the value function parameters and the representation. We introduce the true objective over all states, and then move to the sampled objective for the algorithm in the next section.

We begin by formalizing the representation learning component. Many unsupervised representation learning approaches consist of factorizing input observations 1 X∈R|S|×d into a basis dictionary B∈Rk×d and new representation Φ ∈ R|S|×k . The rows of B form a set of bases, with columns in Φ weighting amongst those bases for each observation (column) in X. Though simple, this approach encompasses a broad range of models, including PCA, CCA, ISOMAP, locally linear embeddings and sparse coding [Singh and Gordon, 2008; Le and White, 2017]. The (unsupervised) sparse coding objective is [Aharon et al., 2006]

where ‖Y‖2F = ∑ij Y2 ij is the squared Frobenius norm; B ∈ Rk×d is a learned basis dictionary; βB, βφ > 0 determine the magnitudes of the regularizers; D ∈ [0, 1]|S|×|S| is a diagonal matrix giving a distribution over states, corresponding to the stationary distribution of the policy dπ : S → [0, 1]; and ||z||2D = z⊤Dz is a weighted norm. The reconstruction error

is weighted by the stationary distribution dπ because states are observed with frequency indicated by dπ . The weighted ℓ1

promotes sparsity on the entries of Φ, preferring entries in Φ to be entirely pushed to zero rather than spreading magnitude across all of Φ. The Frobenius norm regularizer on B ensures that B does not become too large. Without this regularizer, all magnitude can be shifted to B, producing the same ΦB, but pushing ‖Φ‖D,1 to zero and nullifying the utility of its regularizer. Optimizing this sparse coding objective would select a sparse representation φ for each observation x such that φB approximately reconstructs x.

Further, however, we would like to learn a new representation that is also optimized towards approximating the value function. Towards this aim, we need to jointly learn Φ and w, where Φw provides the approximate value function. In this way, the optimization must balance between accurately recreating X and approximating the value function Φw. For this, we must choose an objective for learning w.

We consider two types of objectives: fixed-point objectives and squared-error objectives. Two common fixed-point objectives are the mean-squared Bellman error (MSBE), also called the Bellman residual [Baird, 1995]

and mean-squared projected BE (MSPBE) [Sutton et al., 2009]

1 This variable X can also be a base set of features, on which the agent can improve or which the agent can sparsify.

where D ∈ [0, 1]|S|×|S| is a diagonal matrix giving a distribution over states, corresponding to the stationary distribution of the policy; ||z||2D = z⊤Dz is a weighted norm; and the projection matrix for linear value functions is Π = Φ(Φ⊤DΦ)-1Φ⊤D. The family of TD algorithms converge to the minimum of the MSPBE, whereas residual gradient algorithms typically use the MSBE (see [Sun and Bagnell, 2015] for an overview). Both have useful properties [Scherrer, 2010], though arguably the MSPBE is more widely used.

There are also two alternative squared-error objectives, that do not correspond to fixed-point equations: the mean-squared return error (MSRE) and the Bellman error (BE). For a trajectory of samples {(xi, ri+1, xi+1)}t-1 i=0 , BE is defined as

and the MSRE as t-1∑ i=0 ‖gi+1 -φ⊤i w‖22

where gi+1 = ∑t-1 j=i γ j-i rj+1 is a sample return. In expectation, these objectives are, respectively

where the expectation is w.r.t. the transition probabilities and taking actions according to policy π.

These differ from the fixed-point objectives because of the placement of the expectation. To see why, consider the MSBE and BE. The expected value of the BE is the expected squared error between the prediction from this state and the reward plus the value from a possible next state. The MSBE, on the other hand, is the squared error between the prediction from this state and the expected reward plus the expected value for the next state. Though the MSPBE and MSBE constitute the most common objectives chosen for reinforcement learning, these squared-error objectives have also been shown to be useful particularly for learning online [Sun and Bagnell, 2015].

For sparse coding, however, the MSPBE is not a suitable choice-compared to the MSBE, BE and MSRE-for two reasons. First, the MSBE, BE and MSRE are all convex in Φ, whereas the MSPBE is not. Second, because of the projection onto the space spanned by the features, the MSPBE can be solved with zero error for any features Φ. Therefore, because it does not inform the choice of Φ, the MSPBE produces a two stage approach 2 , where features are learned in a completely unsupervised way and prediction performance does not influence Φ.

2 This problem seems to have been overlooked in two approaches for basis adaptation based on the MSPBE: adaptive bases algorithm for the projected Bellman error (ABPBE) [Di Castro and Mannor, 2010, Algorithm 9] and mirror descent Q(λ) with basis adaptation [Mahadevan et al., 2013]. For example, for ABPBE, it is not immediately obvious this would be a problem, because a stochastic approximation approach is taken. However, if written as a minimization over the basis parameters and the weights, one would obtain a

The final objective for loss L(Φ,w) set to either MSBE, BE or MSRE is

4 Algorithm for Sparse Coding

We now derive the algorithm for sparse coding for policy evaluation: SCoPE. We generically consider either the BE or MSRE. For a trajectory of samples {(xi, ri+1, xi+1)}t-1 i=0 , the objective is

for BE, yi = ri+1 and ¯ γ = γ and for MSRE, yi = ∑tj=i γ j-i rj+1 and ¯ γ = 0. We consider two possible powers for the ℓ1 norm p = 1 or 2, where the theory relies on using p = 2, but in practice we find they perform equivalently and p = 1 provides a slightly simpler optimization. The loss is averaged by t, to obtain a sample average, which in the limit converges to the expected value under dπ . This averaged loss is also more scale-invariant-in terms of the numbers of samples-to the choice of regularization parameters.

SCoPE consists of alternating amongst these three variables, B,w and Φ, with a proximal gradient update for the nondifferentiable ℓ1 norm. The loss in terms of B and w is differentiable; to solve for B (or w) with the other variables fixed, we can simply used gradient descent. To solve for Φ with the B and w fixed, however, we cannot use a standard gradient descent update because the ℓ1 regularizer is nondifferentiable. The proximal update consists of stepping in the direction of the gradient for the smooth component of the objective-which is differentiable-and then projecting back to a sparse solution using the proximal operator: a soft thresholding operator. The convergence of this alternating minimization follows from results on block coordinate descent for non-smooth regularizers [Xu and Yin, 2013].

To apply the standard proximal operator for the ℓ1 regularizer, we need to compute an upper bound on the Lipschitz constant for this objective. The upper bound is 2(1+ ¯ γ 2 )‖w‖22 +2‖B‖2sp , computed by finding the maximum singular value of the Hessian of the objective w.r.t. φi for each

minimum error solution (i.e., error zero) immediately for any basis parameters. The basis parameters are considered to change on a slow timescale, and the weights on a fast timescale, which is a reflection of this type of separate minimization. Menache et al. [2005] avoided this problem by explicitly using a two-stage approach, using MSPBE approaches for learning the parameters and using other score function, such as the squared Bellman error, to update the bases. This basis learning approach, however, is unsupervised.

Representation learning strategies for the MSPBE have been developed, by using local projections [Yu and Bertsekas, 2009; Bhatnagar et al., 2009]. These strategies, however, do not use sparse coding.

i. We will provide additional details for this calculation, and implementation details, in a supplement.

4.1 Local Minima Are Global Minima

In this section, we show that despite nonconvexity, the objective for SCoPE has the nice property that all local minima are in fact global minima. Consequently, though there may be many different local minima, they are in fact equivalent in terms of the objective. This result justifies a simple alternating minimization scheme, where convergence to local minima ensures an optimal solution is obtained.

We need the following technical assumption. It is guaranteed to be true for a sufficiently large k ≤ t (see Haeffele and Vidal [2015]; Le and White [2017]).

Assumption 1 For the given k ≥ d, the following function is convex in Z ∈ Rt+1×d+1

Theorem 1 (Landscape of the SCoPE objective). For the objective in equation (3) with p = 2,

1. under Assumption 1, all full-rank local minima are global minima; and

2. if a local minimum (Φ,B,w) has Φ:i = 0 (i.e., a zero column) and wi = 0, Bi: = 0 (i.e., a zero row) for some 1 ≤ i ≤ k, then it is a global minimum.

Proof. For the first statement, we construct a limit of twicedifferentiable functions fn that Γ-converge to the SCoPE objective f. With this, we can then show that all minimizers of the sequence converge to minimizers of f, and vice-versa [Braides, 2013]. Because all local minimizers of the twicedifferentiable functions fn are global minimizers from [Le and White, 2017, Theorem 10], we can conclude that all corresponding minimizers of f are global minimizers.

We use the pseudo-Huber loss [Fountoulakis and Gondzio, 2013], which is twice-differentiable approximation to the absolute value: |x|µ = √µ2 +x2 -µ. Let θ = (Φ,B,w). The sequence of functions fn are defined with µn = 1/n, as

where L(θ) equals the equation in (3), but without the ℓ21 regularizer on Φ.

Part 1: All local minima of fn for all n are global minima. To show this, we show each fn satisfies the conditions of [Le and White, 2017, Theorem 10 and Proposition 11].

Part 1.1 We can rewrite the loss in terms of Φ[B w ]

where I1:t ∈ Rt×t+1 a diagonal matrix of all ones with the first diagonal entry set to zero, and I0:t-1 with the last diagonal entry set to zero. This loss is convex in the joint variable Φ[B w ] because equation (4) is the composition of a convex function (squared norm) and an affine function (multiplication by ¯ γI1:t - I0:t-1 and addition of y).

Part 1.2 The regularizer on [B w ] must be a weighted Frobenius norm, with weightings on each column; here, we have weighting using regularization parameters βB for the first d columns (corresponding to B) and regularization parameter βw for the last column (corresponding to w).

Part 1.3 The inner dimension k > d, which is true by assumption and the common setting for sparse coding.

Part 1.4 The pseudo-Huber loss, on the columns of Φ, is convex, centered and twice-differentiable.

Part 2: The sequence fn converges uniformly to f. To see why, recall the definition of uniform convergence. A sequence of functions {fn} is uniformly convergent with limit f if for every ϵ > 0, there exists N ∈ N such that for all θ ∈ Θ all n ≥ N, |fn(θ) - f(θ)| < ϵ. Further recall that for any complete metric space, if fn is uniformly Cauchy, then it is uniformly convergent. The sequence is uniformly Cauchy if for all n, m ≥ N, |fn(θ) - fm(θ)| < ϵ. Take any ϵ > 0 and let N = ⌈ 4k(t+1)βφ tϵ ⌉. Then

The upper bound of the first component is maximized when Φij = 0, and so we get

Part 3: Asymptotic equivalence of minimizers of fn and f. Because f is continuous, and so lower semi-continuous, and fn uniformly converges to f, we know that fn Γ-converges to f: fn Γ→ f Braides [2013].

By the fundamental theorem of Γ-convergence, if the {fn} is an equi-coercive family of functions, then the minimizers of fn converge to minimizers of f. A sequence of functions {fn} is equi-coercive iff there exists a lower semi-continuous coercive function ψ : Θ → R ∪ {-∞, ∞} such that fn ≥ ψ on Θ for every n ∈ N [Dal Maso, 2012, Proposition 7.7]. A function is coercive if ψ(θ) → ∞ as ‖θ‖ → ∞. For ψ(θ) = L(θ), it is clear that ψ is coercive, as well as lower semicontinuous (since it is continuous). Further, fn(θ) ≥ L(θ) = ψ(θ), because the regularizer on Φ is non-negative. Therefore, the family {fn} is equi-coercive, and so the minimizers of fn converge to minimizers of f .

For the other direction, if a local minimum θ of f is an isolated local minimum, then there exists a sequence θn → θ with θn a local minimizer of fn for µn sufficiently small [Braides, 2013, Theorem 5.1]. Because we have Frobenius norm regularizers on B,w, which are strongly convex, the objective is strictly convex with respect to B,w. Further, because X is full rank, ‖ΦB-X‖2F is a strictly convex function with respect to Φ. Therefore, locally the objective is strictly convex with respect to θ. We therefore know that local minima of f are isolated, and so there exists an N such that for all n > N, θn are local minimizers of fn. Since these local minimizers are global minimizers, and they converge to θ, this means θ is a global minimum of f .

For the second statement, we use [Haeffele and Vidal, 2015, Theorem 15]. Because we already showed above that our loss can be cast as factorization, it is clear our loss and regularizers are positively homogenous, of order 2. A minimum is guaranteed to exist for our objective, because the loss function is continuous, bounded below (by zero) and goes to infinity as the parameters go to ±∞.

5 Experimental Results

We aim to address the question: can we learn useful representations using SCoPE? We therefore tackle the setting where the representation is first learned, and then used, to avoid conflating incremental estimation and the utility of the representation. We particularly aim to evaluate estimation accuracy, as well as qualitatively understanding the types of sparse representations learned by SCoPE.

Domains. We conducted experiments in three benchmark RL domains - Mountain Car, Puddle World and Acrobot [Sutton, 1996]. All domains are episodic, with discount set to 1 until termination. The data in Mountain Car is generated using the standard energy-pumping policy policy with 10% randomness. The data in Puddle World is generated by a policy that chooses to go North with 50% probability, and East with 50% probability on each step, with the starting position in the lower-left corner of the grid, and the goal in the top-right corner. The data in Acrobot is generated by a near-optimal policy.

Evaluation. We measure value function estimation accuracy using mean absolute percentage value error (MAPVE), with rollouts to compute the true value estimates. MAPVE = 1 ttest ∑s∈Xtest | ˆ V (s)-V ∗ (s)| |V ∗ (s)| , where Xtest is the set of test states, ttest = 5000 is the number of samples in the test set,

ˆ V (s) is the estimated value of state s and V ∗ (s) is the true value of state s computed using extensive rollouts. Errors are averaged over 50 runs.

Algorithms. We compare to using several fixed tile-coding (TC) representations. TC uses overlapping grids on the observation space. It is a sparse representation that is well known to perform well for Mountain Car, Puddle World, and Acrobot. We varied the granularity of the grid-size N and number of tilings D, where D is the number of active features for each observation. The grid is either N×N for Mountain Car and Puddle World or N4 for Acrobot. We explore (D=4, N=4), (D=4,N=8), (D=16,N=4), (D=16,N=8), (D=32,N=4), (D=32,N=8); a grid size of 16 performed poorly, and so is omitted. For Mountain Car and Puddle World the number of features respectively are 64, 256, 256, 1024, 512, 2048, then hashed to 1024 dimensions; for Acrobot, the number of features are 1024, 16384, 4096, 65536, 8192, 131072, then hashed to 4096. Both of these hashed sizes are much larger than our chosen k = 100.

For consistency, once the SCoPE representation is learned, we use the same batch gradient descent update on the MSRE for all the algorithms, with line search to select step-sizes. The regularization weights βB are chosen from {1-5 , . . . , 1-1 , 0}, based on lowest cumulative error. For convenience, βw is fixed to be the same as βB. For learning the SCoPE representations, regularization parameters were chosen using 5-fold cross-validation on 5000 training samples, with βφ = 0.1 fixed to give a reasonable level of sparsity. This data is only used to learn the representation; for the learning curves, the weights are learned from scratch in the same way they are learned for TC. The dimension k = 100 is set to be smaller than for tile coding, to investigate if SCoPE can learn a more compact sparse representation. We tested unsupervised sparse coding, but the error was poor (approximately 10× worse). We discuss the differences between the representations learned by supervised and unsupervised sparse coding below.

Learning curves. We first demonstrate learning with increasing number of samples, in Figure 1. The weights are recom- Features puted using the entire batch up to the given number of samples.

Figure 2: Learned representations Φ for 5000 samples and k = 100, respectively for SCoPE, unsupervised sparse coding, and non-negative unsupervised sparse coding, in Puddle World. The representations learned for Mountain Car and Acrobot have similar structure. The non-negative sparse coding has the additional constraint that each entry in Φ is non-negative. The goal for this addition was to determine if further constraints could improve prediction accuracy for unsupervised sparse coding; though the representation qualitatively looks more reasonable, prediction performance remained poor.

Across domains, SCoPE results in faster learning and, in Mountain Car and Acrobot, obtains lowest final error. Matching the performance of TC is meaningful, as TC is wellunderstood and optimized for these domains. For Acrobot, it is clear a larger TC is needed resulting in relatively poor performance, whereas SCoPE can still perform well with a compact, learned sparse representation. These learning curves provide some insight that we can learn effective sparse representations with SCoPE, but also raise some questions. One issue is that SCoPE is not as effective in Puddle World as some of the TC representations, namely 4-4 and 16-4. The reason for this appears to be that we optimize MSRE to obtain the representation, which is a surrogate for the MAPVE. When measuring MSRE instead of MAPVE on the test data, SCoPE consistently outperforms TC. Optimizing both the representation and weights according to MSRE may have overfitting issues; extensions to MSBE or BE, or improvements in selecting regularization parameters, may alleviate this issue.

Learned representations. We also examine the learned representations, both for unsupervised sparse coding and SCoPE, shown in Figure 2. We draw two conclusions from these results: the structure in the observations is not sufficient for unsupervised sparse coding, and the combination of supervised and unsupervised losses sufficiently constrain the space to obtain discriminative representations. For these two-dimensional and four-dimensional observations, it is relatively easy to reconstruct the observations by using only a small subset of dictionary atoms (row vectors of B in equation (2)). The unsupervised representations, even with additional non-negativity constraints to narrow the search space, are less distributed, with darker and thicker blocks, and more frequently pick less features. For the supervised sparse coding representation, however, the sparsity pattern is smoother and more distributed: more features are selected by at least one sample, but the level of sparsity is similar. We further verified the utility of supervised sparse coding, by only optimizing the supervised loss (MSRE), without including the unsupervised loss; the resulting representations looked similar to the purely unsupervised representations. The combination of the two losses, therefore, much more effectively constrains or regularizes the space of feasible representations and improves discriminative power.

The learning demonstrated for SCoPE here is under ideal conditions. This was intentionally chosen to focus on the question: can we learn effective sparse representations using the SCoPE objective? With the promising results here, future work needs to investigate the utility of jointly estimating the representation and learning the value function, as well as providing incremental algorithms for learning the representations and setting the regularization parameters.

6 Conclusion

In this work, we investigated sparse coding for policy evaluation in reinforcement learning. We proposed a supervised sparse coding objective, for joint estimation of the dictionary, sparse representation and value function weights. We provided a simple algorithm that uses alternating minimization on these variables, and proved that this simple and easy-touse approach is principled. We finally demonstrate results on three benchmark domains, Mountain Car, Puddle World and Acrobot, against a variety of configurations for tile coding.

This paper provides a new view of using dictionary learning techniques frommachine learning in reinforcement learning. It lays a theoretical and empirical foundation for further investigating sparse coding, and other dictionary learning approaches, for policy evaluation and suggests that they show some promise. Formalizing representation learning as a dictionary learning problem facilitates extending recent and upcoming advances in unsupervised learning to the reinforcement learning setting. For example, though we considered a batch gradient descent approach for this first investigation, the sparse coding objective is amenable to incremental estimation, with several works investigating effective stochastic gradient descent algorithms [Mairal et al., 2009, 2010; Le and White, 2017]. The generality of the approach and easy to understand optimization make it a promising direction for representation learning in reinforcement learning.