Manifold Embeddings for Model-Based Reinforcement Learning of

Neurostimulation Policies

Keith Bush kbush@cs.mcgill.ca

Joelle Pineau jpineau@cs.mcgill.ca

Massimo Avoli massimo.avoli@mcgill.ca

Abstract

Real-world reinforcement learning problemsoften exhibit nonlinear, continuous-valued,noisy, partially-observable state-spaces thatare prohibitively expensive to explore. Theformal reinforcement learning framework, un-fortunately, has not been successfully demon-strated in a real-world domain having all ofthese constraints. We approach this domainwith a two-part solution. First, we overcomecontinuous-valued, partially observable state-spaces by constructing manifold embeddingsof the system’s underlying dynamics, whichsubstitute as a complete state-space represen-tation. We then define a generative modelover this manifold to learn a policy o!-line.The model-based approach is preferred be-cause it enables simplification of the learn-ing problem by domain knowledge. In thiswork we formally integrate manifold embed-dings into the reinforcement learning frame-work, summarize a spectral method for esti-mating embedding parameters, and demon-strate the model-based approach in a com-plex domain—adaptive seizure suppression ofan epileptic neural system.

1. Introduction

A driving force in reinforcement learning research isthe growing need for intelligent, autonomous controlstrategies that operate in real-world domains. Inter-esting real-world problems, however, often exhibit non-linear, continuous-valued state-spaces that are onlypartially observable through signals containing somedegree of noise. In a fully observable state represen-tation the first two domain characteristics may be ad-

Appearing in Proceedings of the ICML/UAI/COLT Work-shop on Abstraction in Reinforcement Learning, Montreal,Canada, 2009. Copyright 2009 by the author(s)/owner(s).

dressed by turning to a rich literature of function ap-proximation. The third characteristic, partial observ-ability, however, makes this class of problem arguablymore complex because the functional dependence ofthe state-space must first be guaranteed before approx-imation may proceed.

Partial observability is not unique to reinforcementlearning. Predictive modeling and dynamic systemsresearch have long identified and addressed partial ob-servability (Sauer et al., 1991) from a geometrical per-spective through the method of delayed embeddings,known as Takens Theorem (Takens, 1981). In this ap-proach, sequences of partial observations are grouped(i.e. embedded) onto a higher-dimensional embeddingto characterize the phase space of the underlying sys-tem. This method is the basis for many nonlinearnoise filtering and context-based predictive modelingapproaches (Parlitz & Merkwirth, 2000; Huke, March2006).

The appropriate embedding parameters of a system,unfortunately, are generally unknown a priori andmust be determined empirically. Assuming that fixed-policy data is available we can estimate these parame-ters via spectral analysis (Galka, 2000). Moreover, thedynamic preserving characteristics of manifold embed-dings allow us to take an additional step and directlymodel the system. This is an ideal domain for model-based reinforcement learning.

The primary contribution of this paper is to demon-strate reinforcement learning applied to a embedding-based generative model. For our example we choose achallenging problem in neuroscience research, adaptiveseizure suppression via neurostimulation of epilepticbrain tissue.

2. Background

Our research lies at the intersection of disparate re-search threads. To provide a single mathematical for-malism, in this section we first work through the math-

Manifold Embeddings for Model-Based Reinforcement Learning of Neurostimulation Policies

ematics of reinforcement learning, partial observabil-ity, and embedding theory. We then detail the spec-tral method used to construct manifold embeddings inpractice and motivate why this method is successful.

2.1. Reinforcement Learning

Reinforcement learning (RL) is a class of problemsin which an agent learns an optimal solution to amulti-step decision task by interaction with its envi-ronment (Sutton & Barto, 1998). We focus on theQ-learning formulation.

Consider a system (i.e. the environment), whichevolves according to nonlinear discrete dynamic sys-tem g,

s(t + 1) = g(s(t), a(t)). (1)

This function maps the state of the world, s(t), at thecurrent time, t, and action, a(t) (i.e. decision), ontothe state of world one time-step into the future. Theenvironment also includes a reward function r(t) =h(s(t)), which is a scalar measure of the goodness oftaking the previous action with respect to the goal ofthe multi-step decision task. The agent selects actionsaccording to the policy, !,

a(t) = !(s(t)). (2)

We define reinforcement learning as the process oflearning the policy function that maximizes the ex-pected sum of future rewards. The optimal sequenceof actions is the optimal policy, !!, and the maximumexpected sum of future rewards is termed the action-value function or Q-function, Q!(s(t),a(t)), defined as

Q!(s(t), a(t)) = r(t + 1) + "Q!(s(t + 1), a(t + 1)), (3)

where " is the discount factor. We can then specifythe optimal policy, !!, in terms of the Q-function,

!!(s(t)) = argmaxa

Q!(s(t), a). (4)

Equations 3 and 4 assume that Q! is known. Withouta priori knowledge of Q!, an approximation, Q, mustbe constructed iteratively, using temporal di!erenceerror (Sutton & Barto, 1998).

2.2. Embedding Foundations

The formulation of Q-learning presented above relieson an assumption of state observability. That is, thecomplete environment state, s, is an injective functionof the observable state, s̃. In real world domains thisis often untrue. A major topic in dynamic systems,

signal processing, and time-series analysis is the re-construction of complete state from incomplete obser-vations. One such reconstruction process, the methodof delayed embeddings, may be defined formally byapplying Takens Theorem (Takens, 1981). Here wewill present the key points utilizing the notation ofHuke (Huke, March 2006).

Consider the state-space, s, of the environment is anM -dimensional, real-valued vector space and a is areal-valued action input to the environment. We sub-stitute Equation 2 into Equation 1 and compose a newfunction, #,

s(t + 1) = g(s(t),!(s(t))),

= #(s(t)),

which specifies the discrete time evolution of our com-bined system of agent and environment. Assumethat this system is partially observable via observa-tion function, y, such that

s̃(t) = y(s(t)),

where y : RM ! R and y is noise free and representedwith infinite floating point precision. If # is invertible,and #, #"1, and y are di!erentiable we may applyTakens Theorem to reconstruct the dynamics of sys-tem #(s(t)) using the observables s̃(t). For each s̃(t),we construct a vector sE(t),

sE(t) = [s̃(t), s̃(t " 1), ..., s̃(t " E)]. (5)

If E > 2M then the vectors sE(t) lie on a subset of RE

which is an embedding of our original system (Huke,March 2006). This embedding forms a new dynamicsystem which preserves the structure of the originalsystem. Thus, not only does there exist a vector sE(t)for each observation s̃, but the dynamics governing theevolution of these vectors in time is preserved, suchthat there exists a function, $,

sE(t + 1) = $(sE(t)). (6)

In the context of reinforcement learning, the vectorssE(t) may be substituted into Equations 3 and 4 asreplacements for complete state s(t) without loss ofgenerality. In this same context, however, embeddingsexhibit a serious limitation. Because we roll the policyinto the definition of the state transition, each elementof the reconstructed state space is policy dependent.Therefore, an embedding is only explicitly valid for thepolicy ! under which the time-series was observed. Bychanging the policy slowly, and by carefully choosinga robust function approximation, these e!ects can beminimized.

Manifold Embeddings for Model-Based Reinforcement Learning of Neurostimulation Policies

2.3. Spectral Embedding Method

Embedding theory does not tell us how to select theembedding parameters. In general, the intrinsic di-mension of the system, M , is unknown and not eas-ily determined. Therefore, we need a reliable methodfor selecting parameters that yield a high-quality em-bedding. In practice we utilize a spectral method(Galka, 2000) employing the singular value decompo-sition (SVD).

A summary of this method follows. Consider apartially-observable, discretely-sampled time-series, s̃,of length S̃ that we desire to embed. We choose a suf-ficiently large fixed embedding dimension, E. Su"-ciently large refers to a cardinality of dimension whichis certain to be greater than the dimension in whichthe actual state-space resides.

The spectral method we propose works by manipulat-ing s̃ into a form usable in Equation 5. To do this wedefine the embedding window size, Tmin. We then uti-lize a desample rate, % = Tmin/(E " 1), to uniformlyselect elements of s̃ from the window Tmin, accordingto the rule,

sE(t) = [s̃(t), s̃(t " %), ..., s̃(t " (E " 1)%)]. (7)

We assume a range [T lowmin, Thigh

min ] and interval, #Tmin,over which to explore, Tmin(i) = T low

min + i#Tmin, i #1, 2, ..., Nmin where Nmin = Thigh

min /#Tmin. A good

upper bound, Thighmin , is the fundamental period of the

system if estimable.

For each i, we construct the vectors sE(t), t #

1, ..., NE(i), where NE(i) = S̃ " Tmin(i), and concate-nate them as rows of a matrix, SE(i), of size NE $E.We compute the SVD of SE(i),

SE(i) = U(i)!(i)VT (i)

where !(i) is a diagonal matrix of loadings, (i.e. theenergy contained in the dimensions of the new space)and VT (i) is the basis of the SVD coordinate space.We store the diagonal of each $(i) for analysis as theith row of matrix &, of size Nmin $ E.

Why is spectral method e!ective for identifying high-quality embedding parameters? In the SVD decompo-sition, the right singular vectors, VT are the Eigenvectors of the covariance matrix of the embedding,CE = ST

ESE , if matrix SE is mean centered. There-fore, columns of U define the principal components ofSE (Kirby, 2001).

In the limits of this embedding technique the covari-ance matrix, CE , will take on two forms. For verysmall Tmin the rows of SE are redundant, which drives

elements of CE toward a uniform constant value.Thus, CE is rank one and all of the variance is in thefirst component. This is the time-series itself. In thesecond case, for very long Tmin, the rows of SE becomeuncorrelated. Therefore, diagonal values of the covari-ance matrix, CE , will trend toward constant valueswhile o!-diagonal elements tend to zero. In this caseCE is full rank and is indistinguishable from noise.

Suitable embedding parameters can be found by analy-sis of the principal components at intermediate lengthsof Tmin. Good embedding parameters are found whena subset (preferably small) of the eigenvalues storedin the iopt row of & simultaneously exhibit a local op-tima. The value Tmin(iopt) identifies the appropriateembedding window length, which also determines thedesample rate %(iopt). The subset of eigenvalues withnon-trivial values defines the appropriate embeddingdimension, E#.

3. Methods

Our approach combines best practices from both non-linear dynamic analysis and reinforcement learning toidentify high-quality policies in partially observable,real-world domains that are prohibitively expensive toexplore. This practice incorporates the following steps:1) record a partially observable system under the con-trol of a random policy or some other policy or setof policies known to be near the desired optimal pol-icy; 2) perform spectral embedding; 3) identify goodcandidate parameters for embedding; 4) choose a lo-cal function approximator well-suited to the demandsof the embedding; 5) construct an integrable model ofthe system’s dynamics; and 6) learn a desired policyon this model via reinforcement learning.

4. Case Study: AdaptiveNeurostimulation

Epilepsy a%icts approximately 1% of the world’s pop-ulation (The Epilepsy Foundation, 2009). Of thosesu!ering from this disease 30% do not respond tocurrently available anticonvulsant treatments or arenot candidates for surgical resection. Development ofnew treatments, therefore, is a priority for epilepsy re-search.

Neurostimulation shows promise as an epilepsy treat-ment. In vitro studies indicate that fixed-frequencyexternal electrical stimulation applied to substruc-tures within the hippocampus can e!ectively suppressseizures (Durand & Bikson, 2001). Fixed-frequencypolicies, however, do have limitations. In vitro neu-rostimulation experiments suggest that the e"cacy of

Manifold Embeddings for Model-Based Reinforcement Learning of Neurostimulation Policies

fixed-frequency stimulation varies across epileptic neu-ral systems. The recency of this technology also raisesquestions about long-term impacts such as stimulationinduced tissue damage.

These limitations motivate the search for stimulationpolicies that satisfy additional constraints. Ideally, atreatment policy should adapt to optimally suppressseizures in each unique patient while minimizing thenumber of stimulations necessary to do so. By posingthe problem’s components in this way, an agent (im-plant) that learns a policy (treatment) that maximizesrewards (maximum suppression using minimum stim-ulation constraints) by interacting with the environ-ment (patient), we can recast neurostimulation treat-ment of epilepsy, formally, as a reinforcement learningproblem (Sutton & Barto, 1998).

Recasting as a learning problem requires that thesecomponents be mathematically well-defined. Thecomplex dynamics of neural systems, however, are typ-ically observable only through low-dimensional time-series corrupted by noise (e.g. extracellular record-ing electrodes). Therefore, the objectives of this casestudy are twofold: 1) identify low-dimensional state-space and transition model from field potential record-ings of neural systems that accurately reproduces ob-served neural dynamics and 2) learn a neurostimula-tion therapy in this model that minimizes both seizuresand stimulations.

To fulfill our first objective we construct a state-spaceand transition model from previously recorded dataunder fixed policies. Our dataset is comprised of fieldpotential recordings from five epileptic rat hippocam-pal slices were made under fixed-frequency stimula-tion policies of 0.2, 0.5, and 1.0 Hz as well as control(i.e., unstimulated). The dataset totals 4,639 secondsof recordings including 15 seizures (seizure labels arehand annotated).

From this dataset we desire to construct a high-qualitystate-space and transition model, using only the con-trol data. We measure quality as the predictive ac-curacy of a generative model built from these com-ponents. Using the spectral embedding method, pre-sented in Figure 1, we extract the manifold embed-ding parameters of the dataset (E# = 5, Tmin = 1.2seconds) and then embed the dataset, using Equa-tion 5. This defines the state-space of our system inRE!

. We define the transition model as the local time-derivative of the element of the dataset that is nearestthe current state (i.e., a nearest neighbors derivative).Because the stimulation events are not logged in thedataset, we must find a reasonable mapping of actionsinto our domain. To do this, we define an action as an

1

0

2 3

40

4

80 120 160

0.0

00.

21

0.4

20.

63

0.8

Tmin (s)

Tmin (s)

1.0

Mag

nitu

de o

f Eig

enva

lue

5th PC4th PC

3rd PC

2nd PC

Optimal Embedding Range

Detail of Inset

(b)

(see Detail of Inset)

3rd PC

2nd PC

1st PC

Principal Components Scalings

Mag

nitu

de o

f Eig

enva

lue

(a)

Figure 1. Selecting embedding parameters via PCA: (a)eigenvalues of the principal components of the embeddingas a function of embedding window length Tmin holdingembedding dimension constant at E! = 15; (b) detail ofprincipal components 2–15 from plot (a) for the embed-ding window range Tmin = [0.1, 4] seconds, highlightingthe local maxima of components 2–5 in this range.

interictal-like derivative which is added to the currentderivative of the system. We then define the genera-tive model of the system as numerical integration overthe embedding.

We simulate rat hippocampal dynamics under controland fixed-frequency stimulation policies, choosing theembedding parameters that yield the smallest simu-lation error with respect to analogous fixed-frequencypolicies in the dataset.

Using reinforcement learning, we find an optimal pol-icy which maximizes seizure suppression with minimalstimulations. We define the Q-function approximationfor each state in RE!

as the Q-value of the nearest el-ement of the model. Actions take one of two forms,either on or o!. We perform '-greedy SARSA on thegenerative model subject to the reward function (-1 for

Manifold Embeddings for Model-Based Reinforcement Learning of Neurostimulation Policies

Mean Frac. Mean Seizure Mean Seizureof Seiz. States Length (s) Interval (s)

model 0.19 ± 0.02 61.0 ± 4.6 251.7 ± 14.5data 0.19 ± 0.03 64.4 ± 39.4 271.6 ± 121.3

Table 1. Comparison of summary statistics between theoriginal dataset and 30 generative model simulations of40,000 seconds.

Mean Fraction of Seizure Statespolicy 0.2 Hz 0.5 Hz 1.0 Hz Adaptive"

model 0.15±0.02 0.14±0.02 0.12±0.02 0.02±0.01data 0.16±0.23 0.12±0.16 0.08±0.09 —

Table 2. Comparison of summary statistics between theoriginal dataset and 30 generative model simulations of40,000 seconds under fixed-frequency neurostimulationof 0.2, 0.5, and 1.0 Hz and adaptive neurostimulation("e!ective frequency = 0.02 Hz).

each stimulation and -20 for visiting an element of thedataset that is labeled seizure). We restart the stim-ulation from a random initial element of the datasetevery 1,500 steps (5 min. simulation time) until 90% ofthe model states have been visited during simulationmore than 20 times.

Performance validation of the simulation is summa-rized in Tables 1 and 2 for both control and fixed fre-quency stimulation policies. The model’s seizure dy-namics and suppression predictions fall within confi-dence intervals of the original dataset. Our learnedstimulation policy achieves 0.02 ± 0.01 fraction ofseizure states using an e!ective mean stimulation fre-quency of 0.02Hz, more than an order of magnitudebetter than the best fixed-frequency policy.

The learned policy also provides useful qualitativeknowledge of the domain. For each simulation time-step, we calculate the nearest dataset element of thecurrent state, the seizure label of this element, andthe action requested. From this data we construct theagent’s policy graph, Figure 2(a), projected onto thefirst two principle components. This graph is formedby 1) drawing edges between elements of the datasetthat request stimulation events less than 4 secondsapart and 2) scaling each element’s size by the pro-portion of stimulations requested by the agent at thatelement.

This graph unmasks two distinct classes in the learnedpolicy. The first class consists of pulse trains requestedwhen the simulation operates in the post-seizure regionof the manifold, Figure 2(b). These pulse trains are vi-sually identifiable as cliques in the policy graph formedby the edges. The second policy class consists of in-

dividual, well-timed stimulations requested when thesimulation operates in the dynamically normal regionof the manifold, Figure 2(c)—note the lack of edges,indicating individual stimulation requests.

Without on-line verification of this policy, these resultsare of limited value. Qualitatively, however, these re-sults agree well with neurostimulation results found inthe literature (Durand & Bikson, 2001). Prior stud-ies have observed 1) avoidance of full seizure devel-opment using well-timed stimulations immediately atseizure onset (i.e., at the interface between normal andseizure dynamics) (Durand & Warman, 1994) and 2)early termination of fully developed seizures via low-frequency stimulation (D’Arcangelo et al., 2005). Thelearned neurostimulation policy incorporates both ofthese policy classes, but applies them under di!er-ent dynamic regimes, seemingly to maximize the sup-pression benefits of each class. The literature doesnot claim that well-timed stimulations can abort fullydeveloped seizures. Moreover, there exists evidencesuggesting that fixed-frequency policies below 0.2 Hzexhibit little or no seizure suppression. Therefore, itmakes sense that a well-timed stimulation policy in thenormal regime, if possible, would be the best availablechoice.

5. Discussion

The reinforcement learning community has long beenaware of the need to find low-dimensional represen-tations to capture complex domains. Approaches fore"cient function approximation, basis function con-struction, and discovery of embeddings have been thetopic of significant investigations (Bowling et al., 2005;Keller et al., 2006; Smart, 2004; Mahadevan & Mag-gioni, 2007). However most of this work has beenlimited to the fully observable (MDP) case and hasnot been shown to extend to partially observable en-vironments. The question of state space representa-tion in partially observable domains was tackled underthe POMDP framework (McCallum, 1996) and morerecently in the PSR framework (Singh et al., 2003).While these methods tackle a similar problem theyhave primarily been limited to discrete action and ob-servation spaces.

The PSR framework was recently extended to contin-uous (nonlinear) domains (Wingate & Singh, 2007).This method is significantly di!erent from our work,both in terms of the class of representations it consid-ers and in the criteria used to select the appropriaterepresentation. Furthermore, it has not yet been ap-plied to real-world domains. Nonetheless it could po-tentially be used to tackle the problems presented in

Manifold Embeddings for Model-Based Reinforcement Learning of Neurostimulation Policies

−20

2

−10 −5 0

Principle Component 1Pr

incip

le C

ompo

nent

2(a)

(b) (c)

Figure 2. Policy analysis: (a) The learned policy graphplotted on the first two-principle components of the em-bedding manifold. (b) Detail of the policy graph in thepost-seizure portion of the attractor. (c) Detail of the pol-icy graph in the dynamically normal portion of the attrac-tor.

Section 4; an empirical comparison with our approachis left for future consideration.

The implications of this technique for modeling poorlyunderstood, partially observable domains are moreconcrete. We demonstrate a data-driven, generativemodel construction technique that accurately repro-duces the dynamics of fixed-frequency electrical stim-ulation applied to a slice of epileptic neural tissue. Us-ing reinforcement learning, we identify an optimal sup-pression policy that is, interestingly, two-class. Thislearned policy predicts a dependence between suppres-sion e"cacy of these two policy classes and the dy-namic regime in which these policy classes are applied,which is supported, in part, by previous in vitro ex-periments.

References

Bowling, M., Ghodsi, A., & Wilkinson, D. (2005). Actionrespecting embedding. Proceedings of ICML.

D’Arcangelo, G., Panuccio, G., Tancredi, V., & Avoli,M. (2005). Repetitive low-frequency stimulation re-duces epileptiform synchronization in limbic neuronalnetworks. Neurobiology of Disease, 19, 119–128.

Durand, D., & Warman, E. (1994). Desynchronization ofepileptiform activity by extracellular current pulses inrat hippocampal slices. Journal of Physiology, 71, 2033–2045.

Durand, D. M., & Bikson, M. (2001). Suppression andcontrol of epileptiform activity by electrical stimulation:A review. Proceedings of the IEEE, 89(7), 1065–1082.

Galka, A. (2000). Topics in nonlinear time series analysis:with implications for eeg analysis. World Scientific.

Huke, J. (March, 2006). Embedding nonlinear dynamicalsystems: A guide to takens’ theorem (Technical Report).Manchester Institute for Mathematical Sciences, Univer-sity of Manchester.

Keller, P., Mannor, S., & Precup, D. (2006). Automatic ba-sis function construction for approximate dynamic pro-gramming and reinforcement learning. Proceedings ofICML.

Kirby, M. (2001). Geometric data analysis: An empiricalapproach to the dimensionality reduction and the studyof patterns. Wiley and Sons.

Mahadevan, S., & Maggioni, M. (2007). Proto-value func-tions: A Laplacian framework for learning representa-tion and control in Markov decision processes. Journalof Machine Learning Research, 8.

McCallum, A. K. (1996). Reinforcement learning with se-lective perception and hidden state. Doctoral disserta-tion, University of Rochester.

Parlitz, U., & Merkwirth, C. (2000). Prediction of spa-tiotemporal time series based on reconstructed localstates. Physical Review Letters, 84(9), 1890–1893.

Sauer, T., Yorke, J. A., & Casdagli, M. (1991). Embedol-ogy. Journal of Statistical Physics, 65:3/4, 579–616.

Singh, S., Littman, M. L., Jong, N. K., Pardoe, D., &Stone, P. (2003). Learning predictive state representa-tions. Machine Learning: Proceedings of the 2003 Inter-national Conference (ICML) (pp. 712–719).

Smart, W. (2004). Explicit manifold representations forvalue-functions in reinforcement learning. Proceedingsof the Int. Sympo. on AI and Math.

Sutton, R. S., & Barto, A. G. (1998). Reinforcement learn-ing: An introduction. Cambridge, MA: The MIT Press.

Takens, F. (1981). Detecting strange attractors in turbu-lence, dynamical systems and turbulence. In D. A. R.. L. S. Young (Ed.), Lecture notes in mathematics, vol.898, 366–381. Springer.

The Epilepsy Foundation (2009). Epilepsy andseizure statistics. Available at http://www.epilepsyfoundation.org/about/statistics.cfm.

Wingate, D., & Singh, S. (2007). On discovery and learn-ing of models with predictive state representations ofstate for agents with continuous actions and observa-tions. Proceedings of AAMAS.