FLEX: an Adaptive Exploration Algorithm for Nonlinear Systems

An important class of models is the class of functions with an affine dependence on the parameters. In the remainder of this work, we use the word “linear” for an affine dependence. Note that the term “linear” refers to the parameter dependence of a model, but the dependence in the state-action pair $z$ is still assumed to be nonlinear in general.

The class of linear models is critically important for several reasons. First, the dynamics can very often be formulated as an affine function of some well-defined scalar parameters, hence allowing for efficient estimation by ordinary least squares (Tedrake, 2022). Besides, a natural mathematical exploration objective from optimal experimental design theory can be derived for linear models (see Section 3.2). Finally, the computations allowed for linear models can be generalized to nonlinear models, as we will see in Section 5.

An important case of linear models is when the parameters are separated in a matrix form according to row-wise dependence, as follows.

Let $y$ be an observation whose distribution $p(y|\pi,\theta)$ depends on a parameter $\theta\in\mathbb{R}^{n}$ and a decision variable (or design) denoted $\pi$. Experimental design theory provides a quantitative answer to the question: how informative are the observations $y$ for estimating $\theta$ when the decision $\pi$ is taken? The information is defined as a scalar functions of the Fisher information matrix (Fedorov, 2010) as follows.

This information gain quantifies the information about $\theta$ provided by the observations with design $\pi$, and it may be interpreted as the volume of the confidence ellipsoid for the parameter vector $\theta$. Several other functionals can be used instead of $\log\det$, leading to other optimality criteria. The D-optimality criterion benefits from a property of scale invariance (Pukelsheim, 2006) and is suitable for our online setting because it provides a simple rank-one update formula, as we will see in Section 4. A D-optimal design maximizes the information gain:

From a Bayesian perspective, D-optimality minimizes the entropy of the expected posterior on $\theta$, or equivalently the mutual information between the current prior and the expected posterior (Chaloner & Verdinelli, 1995).

In this section, we apply the optimal experimental design framework to our dynamical setting, where the Gaussian noise assumption and the linear structure of the model allow for exact computations. The observations are the trajectory ${y=(x_{0:t},u_{0:t-1})}$ and the design is the policy $\pi$ generating the inputs $u_{0:t-1}$.

Note that $I$ does not depend explicitly on $\theta$ but only on the observations, because the model is linear in the parameters.

It follows from (3.6) and Proposition 1 that, in our setting, D-optimal inputs solve the following optimal control problem:

where we neglect the noise in the dynamics for simplicity, and recall that $V_{t}=V(z_{t})$.

Linear models can be learned online with the recursive least squares formula for the estimator (3.8). Assuming for simplicity $d=1$, $c=0$ and denoting ${v_{t}={V(z_{t})}^{\top}\in\mathbb{R}^{n}}$, online learning takes the form

with the squared error loss  $\ell_{t}(\theta)=\frac{1}{2}\times\|{v_{t}}^{\top}\theta-x_{t+1}\|^{2}_{2}$ and the matrix step size $H_{t}:=M_{t}^{-1}$. Equation (3.11b) makes it transparent that recursive least squares is an online gradient descent step. The memory cost of learning is $\mathcal{O}(n^{2})$ for the storage of $M_{t}$ and $\theta_{t}$. For $d>1$, there is one update for each row, hence $d$ updates per time step $t$.

Since we seek minimal complexity and adaptivity, we choose to devote the computational effort at time $t$ to the choice of the next input $u_{t}$ only. We want to define an informativeness measure $F_{t}$ for input $u_{t}$ and solve a sequence of problems of the form

As stated before, we attach great importance to the computational time of solving (4.1).

The function $F_{t}$ should quantify the information brought by $u_{t}$, for which we derived a mathematical expression in our information-theoretic considerations of Section 3. Therefore, we want the problem sequence (4.1) to be an approximation of the problem (3.10).

A greedy approximation can be derived as follows. At time $t$, the past trajectory $z_{0:t}$ is known and the choice of $u_{t}$ immediately determines the next state $x_{t+1}$. We define $F_{t}$ as the predicted information gain truncated at $t+1$.

When $V$ is of rank one, a simpler formula can be derived.

We adopt the rank-one approximation of Lemma 4.1 in the remainder of this work. Note that the index $k$ of row $v$ can be either drawn randomly or chosen using prior knowledge: the $k$-th row of $V$ is informative if the $k$-th component of the model is sensitive with respect to the parameters. One could also design a numerical criterion for this choice.

For linear systems, recent approaches have shown that a greedy D-optimal policy yields good exploration performance (Blanke & Lelarge, 2022). However for nonlinear systems, the nonlinearity of $v(z)$ makes the maximization problem (4.1) challenging. To obtain a simpler optimization problem, we linearize the model with respect to the state. Intuitively, since we are planning between $t$ and $t+\mathrm{d}t$, the corresponding change in the state is small so it is reasonable to use a linear approximation to the mapping $z\mapsto v(z)$.

Our exploration policy is summarized in Algorithm 2. Note that since it optimizes a greedy objective at each time step, it is adaptive by nature and it does not require the knowledge of the time horizon $T$. The following result shows that solving (4.5) can be achieved at low cost, ensuring that our policy is computationally efficient.

The developments of Section 4 are based on the linear dependence of the model on the parameters. For nonlinear models, a natural idea is to make a linear expansion of the model: assuming that the parameter vector is close to a convergence value $\theta_{\star}$, we can linearize ${f}$ to first order in ${\theta-\theta_{\star}}$:

In the limiting regime where this linear approximation would hold, ${f}$ would be a linear model with features

Considering this analogy, we can generalize D-optimal experimental design (MacKay, 1992) and we want to extend the results of Section 4 to nonlinear models.

In our dynamical framework, we expect the approximation (5.1) to be increasingly accurate as more observations are collected, hence motivating the generalization of the exploration strategy developed in Section 4 to nonlinear models. The features of the linearized model (5.1) are unknown because the Jacobian (5.2) depends on the unknown parameter $\theta_{\star}$ in general. However, we can approximate $\theta_{\star}$ by the current estimate $\theta_{t}$ at each time step along the trajectory. By defining

we extend the notion of the Gram matrix in Definition 3.4 as well as the optimal control problem (3.10) to nonlinear models. Note that since $V_{t}$ depends only on $(z_{t},\theta_{t})$, the Gram matrix can still be computed online along the trajectory. Similarly, we extend the approach developed in Section 4 by defining

with $1\leq k\leq d$ chosen as in Lemma 4.1. With these quantities defined, Algorithm 2 is extended to arbitrary models.

In Algorithm 2, we need to compute $D=\partial v/\partial x$, which amounts to computing the derivatives of $f$ in both $\theta$ and $z$:

For neural networks, this matrix can be computed by automatic differentiation. The cost of solving (4.5) does not depend on $n$ so the computation of $D$ becomes the computational bottleneck for large models. The complexity of the latter operation is $\mathcal{O}(nd)$ with automatic differentiation.

We train nonlinear models using online gradient descent:

which extends (3.11b). For neural networks, online learning with an adaptive learning rate (and scalar $H_{t})$ is known to be effective (Bottou, 2012; Kingma & Ba, 2015). The gradient step can be averaged over a batch for smoother learning, at the cost of storing a small amount of data points.

We first test our policy on various nonlinear environments from classical control, covering different values for $d$ and $m$. We compare its performance in terms of sample efficiency to that of various baselines. The agents have the same learning model, but different exploration policies. Random exploration draws inputs at random. For pendulum-like environments, a periodic oracle baseline excites the system at an eigenmode, yielding resonant trajectories of large amplitude. A baseline called “uniform” maximizes the distance of the trajectory points in the state space: it optimizes objective (4.2) with ${{G(z)=\frac{1}{2}\sum_{s=0}^{t}\|{x}-x_{s}\|^{2}}}$.

In real systems, dynamics may vary over time and an adaptive exploration is crucial for accommodating to changes in the environment. We test the adaptivity of our method by exploring a time-varying system, and compare to an episodic agent. The system is a central, repulsive force field centered on a star, moving around a circle uniformly at a period $T$. The agent is a spaceship and the control variable is the acceleration. The force field varies significantly only in the vicinity of the star: the spaceship learns information about the dynamics only when it is close, and hence needs to track the star.

The goal of exploration is ultimately to obtain an accurate model for the system for model-based control. In order to validate the relevance of our approach to this framework, we evaluate the model learned during exploration on model-based control tasks. We compare it to the recent active exploration algorithms RHC and MAX (Schultheis et al., 2020; Shyam et al., 2019) and to the model-free reinforcement learning algorithm SAC (Haarnoja et al., 2018).