# Optimistix: modular optimisation in JAX and Equinox

Jason Rader  Terry Lyons  Patrick Kidger

###### Abstract

We introduce Optimistix: a nonlinear optimisation library built in JAX and
Equinox. Optimistix introduces a novel, modular approach for its minimisers
and least-squares solvers. This modularity relies on new practical
abstractions for optimisation which we call search and descent, and which
generalise classical notions of line search, trust-region, and learning-rate
algorithms. It provides high-level APIs and solvers for minimisation,
nonlinear least-squares, root-finding, and fixed-point iteration. Optimistix
is available at https://github.com/patrick-kidger/optimistix.

Machine Learning

##  1 Introduction

JAX is a Python autodifferentiation framework popular for scientific computing
and machine learning (Bradbury et al., 2018; Kidger, 2021). Equinox (Kidger &
Garcia, 2021) extends many JAX core transformations and concepts, and adds
additional functionality for parameterised functions. Equinox has become a
popular choice for machine learning (Hall et al., 2023a, b; Singh, 2022; Wang,
2023) and scientific machine learning ‘sciML’ (Kidger, 2021; Rader et al.,
2023; Pastrana et al., 2023; Pastrana, 2023) in JAX.

We introduce Optimistix, a nonlinear optimisation library built in JAX +
Equinox. Optimistix targets differentiable scientific computing and sciML
tasks. The sciML ecosystem in JAX is large and growing, and already includes
packages for differentiable rigid-body physics simulation (Freeman et al.,
2021), computational fluid dynamics (Dresdner et al., 2022; Bezgin et al.,
2022), protein structure prediction (Jumper et al., 2021), linear solves and
least-squares (Rader et al., 2023), ordinary and stochastic differential
equations (Kidger, 2021), general-purpose optimisation (Blondel et al., 2021),
structural design (Pastrana et al., 2023), Bayesian optimisation (Song et al.,
2022; Golovin et al., 2017), and probabilistic modeling (Stanojević & Sartran,
2023).

###  1.1 Contributions

We introduce Optimistix, a JAX nonlinear optimisation library with the
following features:

  1. 1.

Modular optimisers.

  2. 2.

Fast compile times and run times.

  3. 3.

Support for general PyTrees, 111JAX data structures consisting of arbitrarily
nested container types, containing other JAX/Python types. The containers
(tuples/dictionaries/lists/custom types) are referred to as ‘nodes’, and the
data types they hold as ‘leaves’. and use of PyTrees for solver state.

We highlight the first point as the central contribution of this work. This
allows users to define custom optimisers for a specific problem by swapping
components of the optimiser.

To achieve this, Optimistix introduces two practical abstractions: search and
descent. These abstractions generalise classical line search, trust-region,
and learning-rate algorithms. To the best of our knowledge, Optimistix is the
first optimisation software modularised in this way.

Optimistix includes APIs for four types of optimisation tasks: minimisation,
nonlinear least-squares, root-finding, and fixed-point iteration. Each has a
high-level API, and support automatic conversion of appropriate problem types.

Optimistix is already seeing adoption in sciML. For example, in software
libraries for solving differential equations (Kidger, 2021) and probabilistic
inference (Carroll, 2023).

Example usage on the Rosenbrock problem (Rosenbrock, 1960) using BFGS and
automatically converting least-squares to minimisation is:

⬇

1import jax

2import optimistix as optx

3

4def loss(x, scaling):

5 res1 = scaling * (x[1:] - x[:-1]**2)

6 return (res1, 1 - x)

7

8init = jax.numpy.zeros(100)

9scaling = 10

10solver = optx.BFGS(rtol=1e-5, atol=1e-6)

11minimum = optx.least_squares(

12 loss, solver, init, args=scaling

13)

##  2 Background: Preconditioned Gradient Methods

The dominant approach to differentiable optimisation is to locally approximate
the objective function $f:\mathbb{R}^{N}\to\mathbb{R}$ at an iterate $x_{k}$
using a quadratic model function:

$m_{k}(p)=f(x_{k})+\nabla f(x_{k})^{T}p+\frac{1}{2}p^{T}H_{k}p$ (1)

where $p\in\mathbb{R}^{N}$, and $H_{k}\in\mathbb{R}^{N\times N}$ is an
approximation to the Hessian $\nabla^{2}f(x_{k})$ (Nocedal & Wright,
2006)[sections 1-6], (Bonnans et al., 2006)[section 4], (Conn et al., 2000).
This quadratic is used to find the next step in an iterative algorithm to
minimise $f$. The minimum of (1) is found at $p=-H_{k}^{-1}\nabla f(x_{k})$.
As such, these are sometimes called preconditioned gradient methods (Gupta et
al., 2018) (not to be confused with preconditioned conjugate gradient methods
(Nocedal & Wright, 2006)[section 5].)

Most first-order methods common in the machine learning literature, such as
Adam (Kingma & Ba, 2017) and Adagrad (Duchi et al., 2011), are also
preconditioned gradient methods. In these algorithms, $H^{-1}$ is usually
stored and updated directly, rather than computed from $H$. Moreover, $H^{-1}$
is often highly structured or diagonal to reduce memory cost.

The class of preconditioned gradient methods is extremely large, and includes
Newton’s method, quasi-Newton methods, and Gauss-Newton methods (Bonnans et
al., 2006)[section 4], (Nocedal & Wright, 2006)[section 6], some nonlinear
conjugate gradient methods (Sherali & Ulular, 1990), gradient descent, and
adaptive gradient methods such as Adam (Kingma & Ba, 2017) and Adagrad (Duchi
et al., 2011). In many of these algorithms, the quadratic approximation is
implicit.

Preconditioned gradient methods are central in Optimistix, and were an
important motivator of the descent and search abstractions we now introduce in
section 3.

##  3 Modularity in Optimistix

This section introduces the key advancement of Optimistix: modularity. To
attain this, Optimistix uses a generalised approach to line searches, trust-
regions, and learning-rates through the abstractions of search and descent.
These concepts offer a precise formalism of ideas already present in the
optimisation literature. This approach makes it easier both to advance theory
(in Section 3.1 we demonstrate how this may be used to create a novel
minimiser), and serve as a practical approach to modularity. To the best of
our knowledge this approach is new here, and is not present in any other open-
source optimisation package.

Consider a scalar function to minimise: $f:\mathbb{R}^{n}\to\mathbb{R}$.
Searches consume local information about $f$, such as its value, gradient,
and/or Hessian at a point, and return a scalar. This scalar corresponds to the
distance along a line search, the trust-region radius of a trust-region
method, the value of a learning-rate, etc. Searches are a generalisation of
these.

A descent consumes this scalar, as well as the same local information about
$f$, and returns the step the optimiser should take. For example, gradient
descent with a fixed learning-rate $\alpha$ is implemented as a search which
returns a fixed value $\alpha\in\mathbb{R}$, and a descent which returns
$\alpha\operatorname{\nabla}f(x_{k})$.

Searches and descents are modular components in Optimistix, and they can be
easily swapped by the user.

###  3.1 Creating a Novel Optimiser with Ease: an Example

Optimistix provides a new ”mix-and-match” API which allows users to easily
create new optimisers. For example, here we define a custom optimiser and use
this to solve a toy linear regression problem:

⬇

1from collections.abc import Callable

2import jax.numpy as jnp

3import jax.random as jr

4import optimistix as optx

5

6class HybridMinimiser(optx.AbstractBFGS):

7 rtol: float

8 atol: float

9 norm: Callable = optx.max_norm

10 use_inverse: bool = False

11 descent: optx.AbstractDescent = (

12 optx.DoglegDescent()

13 )

14 search: optx.AbstractSearch = (

15 optx.LearningRate(0.1)

16 )

17

18solver = HybridMinimiser(

19 rtol=1e-8, atol=1e-9

20)

21

22def loss(weights, data):

23 x, y = data

24 return weights.T @ x - y

25

26key1, key2 = jr.split(jr.PRNGKey(0))

27noise = 0.1 * jr.normal(key1, shape=(99,))

28

29true_weights = jnp.array([3.14, -7, 2.71])

30xs = jr.normal(key2, shape=(3, 99))

31ys = true_weights.T @ xs + noise

32init = jnp.array([0.0, 0.0, 0.0])

33

34# Convert least-squares to minimisation

35# then solve.

36minimum = optx.least_squares(

37 loss, solver, init, args=(xs, ys)

38)

39# ‘minimum.value‘:

40# Array([3.129, -6.983, 2.732]

`HybridMinimiser` defines an optimiser which, at each step $k$, uses the
gradient and BFGS quasi-Newton approximation
$B(x_{k})\approx\nabla^{2}f(x_{k})$ (Nocedal & Wright, 2006)[section 6.1] to
construct a quadratic model function as discussed in section 2. A ‘dogleg’
descent path is then built, which interpolates between $x_{k}$,
$\operatorname{\nabla}f(x_{k})$, and
$B^{-1}(x_{k})\operatorname{\nabla}f(x_{k})$ with a piecewise linear curve.
Finally, a constant step of length `0.1` is taken along this descent path, to
form the next iterate $x_{k+1}$.

`HybridMinimiser` is not an off-the-shelf optimiser, and would require a
custom implementation in most optimisation packages. Implementing a custom,
performant algorithm may require significant experience in optimisation, and
hundreds of lines of technical code. In Optimistix, a performant
implementation of this novel optimiser takes fewer than 10 lines of code.

Custom optimisers allow users to choose optimisation methods appropriate for
the problem at hand. For example, the novel optimiser above can solve the
poorly-scaled Biggs EXP6 function (Moré et al., 1981), whereas the standard
BFGS algorithm with backtracking line search fails to solve the problem to an
acceptable accuracy.

###  3.2 Search

For most nonlinear functions, the quadratic approximation $m_{k}$ in (1) is
only a reasonable approximation in a small neighborhood of $x_{k}$. The full
preconditioned gradient step $-H_{k}^{-1}\operatorname{\nabla}f(x_{k})$ often
overshoots the region where the approximation is good, slowing the
optimisation process. Line searches, trust-regions, and learning-rates are all
methods to keep steps within a region where the approximation is good.

We begin by discussing the classical approach to line searches and trust-
regions.

Line searches

Line searches (Nocedal & Wright, 2006)[section 3] move in the direction of the
preconditioned gradient, but only move a certain amount
$\alpha_{k}\in\mathbb{R}^{+}$. ie.

$x_{k+1}=x_{k}-\alpha_{k}H_{k}^{-1}\operatorname{\nabla}f(x_{k}).$

Algorithms for choosing $\alpha_{k}$ often seek to satisfy conditions of
sufficient decrease and curvature – keeping $\alpha_{k}$ from growing too
large or shrinking too small. Popular conditions are the Armijo conditions,
Wolfe conditions, and Goldstein conditions with various relaxations of each
(Nocedal & Wright, 2006)[section 3], (Moré & Thuente, 1994), (Bonnans et al.,
2006).

Trust-region methods

Trust-region methods (Nocedal & Wright, 2006)[section 4] are another popular
class of algorithms which seek to approximately solve the constrained
optimisation problem

$\displaystyle p^{*}$
$\displaystyle=\operatorname*{arg\,min}_{p}m_{k}(p)\text{ subject to
}\left\lVert p\right\rVert\leq\Delta_{k}$ (2) $\displaystyle x_{k+1}$
$\displaystyle=x_{k}+p^{*}$ (3)

for some norm $\left\lVert\cdot\right\rVert$ and trust-region radius
$\Delta_{k}\in\mathbb{R}^{+}$. The trust-region radius is chosen at each step
based upon how well $m_{k}$ approximated $f$ at the previous step, see (Conn
et al., 2000)[sections 6.1 and 10.5] for details.

The minimum in (2) is found using a variety of approximate methods (Nocedal &
Wright, 2006)[section 4], (Steihaug, 1983), and when $\left\lVert-
H_{k}^{-1}\operatorname{\nabla}f(x_{k})\right\rVert\leq\left\lVert\Delta_{k}\right\rVert$,
then by construction $p^{*}=-H_{k}^{-1}\operatorname{\nabla}f(x_{k})$. For
this reason, trust-region methods are often interpreted as a class of methods
for interpolating between $x_{k}$ and $-H^{-1}\operatorname{\nabla}f(x_{k})$
via the scalar parameter
$\Delta_{k}\in\left[0,\left\lVert-H^{-1}\operatorname{\nabla}f(x_{k})\right\rVert\right].$

One advantage line search algorithms have over trust-region algorithms is the
value $H_{k}^{-1}\operatorname{\nabla}f(x_{k})$ can be computed once, cached,
and reused as $\alpha_{k}$ varies. This is not always true for trust-region
algorithms, which may require significant effort to recompute $p^{*}$ as
$\Delta_{k}$ varies.

Searches in Optimistix

A _search_ is a new abstraction introduced in Optimistix to generalise line
searches, trust-region methods, and learning-rates. Searches are defined as
functions taking local information and an internal state, and producing a
scalar

$s:\mathcal{D}\times\mathcal{S}\to\mathbb{R}.$ (4)

For an example of a line search, consider the backtracking Armijo update. The
Armijo backtracking algorithm uses local data
$d_{k}=(d_{k}^{(1)},d_{k}^{(2)})\in\mathcal{D}=\mathbb{R}\times\mathbb{R}^{N}$,
where $d_{k}^{(1)}$ represents the objective function value $f(x_{k})$, and
$d_{k}^{(2)}$ it’s gradient $\operatorname{\nabla}f(x_{k})$. The Armijo search
state is
$\sigma_{k}=(\sigma_{k}^{(1)},\sigma_{k}^{(2)})\in\mathcal{S}=(0,1]\times\\{0,1\\}$,
where $\sigma_{k}^{(1)}$ represents the current step-size and
$\sigma_{k}^{(2)}$ represents whether to shrink or reset the step size,
depending on whether the last step satisfied the Armijo condition. For a
decrease factor $c\in(0,1]$, the Armijo search takes the step

$s(d_{k},\sigma_{k})=\begin{cases}c\sigma_{k}^{(1)}&\sigma_{k}^{(2)}=0\\\
1&\sigma_{k}^{(2)}=1,\end{cases}$ (5)

and updates its state $\sigma_{k}$ via

$\sigma_{k+1}^{(1)}=s(d_{k},\sigma_{k})$

and $\sigma_{k+1}^{(2)}=1$ when the Armijo condition

$f(x_{k}+\delta_{k})\leq d_{k}^{(1)}+\eta\delta_{k}^{T}d_{k}^{(2)}$ (6)

is satisfied, or $\sigma_{k+1}^{(2)}=0$ otherwise. Here, $0<\eta<1$ is a
hyperparameter which determines how much a step must decrease to be accepted
(larger means more decrease is required) and $\delta_{k}$ is the proposed
step.

There are a number of different line search algorithms, but trust-region
methods typically use the same trust-region radius selection algorithm. This
algorithm is represented by the search

$s_{\text{TR}}(d_{k},\sigma_{k})=\begin{cases}c_{2}\sigma_{k}^{(1)}&\sigma_{k}^{(2)}>C_{2},\\\
\sigma_{k}^{(1)}&C_{1}<\sigma_{k}^{(2)}<C_{2}\\\
c_{1}\sigma_{k}^{(1)}&\sigma_{k}^{(2)}<C_{1}\\\ \end{cases}$ (7)

where $0<c_{1}<1$ is a decrease amount, $c_{2}>1$ an increase amount, $C_{1}$
a low cutoff (close to $0$) and $C_{2}$ is a high cutoff (close to $1$.) The
state is updated via

$\sigma_{k+1}^{(1)}=s(d_{k},\sigma_{k})$

and

$\sigma_{k}^{(2)}=\frac{f(x_{k})-f(x_{k}+\delta_{k})}{m_{k}(0)+m_{k}(\delta_{k})}$
(8)

where $m_{k}$ is the quadratic model function (1) and $\delta_{k}$ is again
the proposed step $x_{k+1}=x_{k}+\delta_{k}$. The second component of the
state, $\sigma_{k}^{(2)}$, is referred to as the ‘trust-region ratio’, and
roughly indicates how well the model function $m_{k}$ predicted the decrease
in $f$ that would come from taking the step $\delta_{k}$.

Current searches implemented in Optimistix are: learning rate, backtracking
Armijo line search (Nocedal & Wright, 2006)[section 3.1], the classical trust-
region ratio update (Conn et al., 2000)[section 6.1], and a trust-region
update using a linear local approximation for first-order methods (Conn et
al., 2000).

###  3.3 Descent

A _descent_ is another new abstraction in Optimistix. It is a function taking
a scalar, the same local data, and an internal state and producing a vector:

$\delta:\mathbb{R}\times\mathcal{D}\times\mathcal{R}\to\mathbb{R}^{N}.$ (9)

The value produced by $\delta$ is used to update the iterates $x_{k}$

$x_{k+1}=x_{k}+\delta(\alpha_{k},d_{k},\gamma_{k})$

with $\alpha_{k}\in\mathbb{R},\ d_{k}\in\mathcal{D},\
\gamma_{k}\in\mathcal{R}$.

The main function of the descent is to map the scalar produced by the search
$\alpha_{k}$ into a meaningful optimisation step.

For example, we can represent the preconditioned gradient step used in a line
search as a descent via

$\delta_{\text{LS}}(\alpha_{k},d_{k},\gamma_{k})=-\alpha_{k}(d^{(2)})^{-1}(d^{(1)})$
(10)

here,
$d_{k}=(d^{(1)}_{k},d^{(2)}_{k})\in\mathcal{D}=\mathbb{R}^{N}\times\mathbb{R}^{N\times
N}$. Breaking this descent down, consider applying (10) to a minimisation
problem. If $d_{k}=(\operatorname{\nabla}f(x_{k}),H_{k})$, then (10) is
equivalent to the standard preconditioned gradient descent
$\delta_{\text{TR}}(\alpha_{k},d_{k},\gamma_{k})=-\alpha_{k}H_{k}^{-1}\operatorname{\nabla}f(x_{k})$,
with a step-size $\alpha_{k}$. If $d_{k}=(r_{k},J_{k})$, the residual vector
$r_{k}$ and Jacobian $J_{k}$ at step $k$, then (10) is the Gauss-Newton step
with step-size $\alpha_{k}$ for a least-squares problem.

As another example, we can represent the trust-region subproblem as a descent
via

$\delta_{\text{TR}}(\alpha_{k},d_{k},\gamma_{k})\approx\operatorname*{arg\,min}_{p,\
\|p\|\leq\alpha_{k}}d_{k}^{(1)}+\operatorname{\nabla}(d_{k}^{(2)})^{T}p+\frac{1}{2}p^{T}d_{k}^{(3)}p$
(11)

for
$d_{k}\in\mathcal{D}=\mathbb{R}\times\mathbb{R}^{N}\times\mathbb{R}^{N\times
N}$. For a minimisation problem this is
$d_{k}=(f(x_{k}),\operatorname{\nabla}f(x_{k}),H_{k})$. Here, the output of
the search is mapped not to a line search, but to the trust-region radius of a
trust-region algorithm.

$\alpha_{k}$ has a similar meaning in both of these algorithms. It’s used to
restrict the size of the update at step $k$ when confidence in the
approximation $m_{k}$ is low. By abstracting the map from $\alpha_{k}$ to the
optimiser update $\delta_{k}$, descents allow a search to be used in a variety
of different algorithms without knowing exactly how the step-size will affect
the optimiser update. This decoupling is novel to optimistix, and is one of
its most powerful features.

Optimistix descents, like classical line searches, allow certain values used
in their computation to be cached and reused. For example, in
$\delta_{\text{LS}}$, the value $(d^{(2)})^{-1}d^{(1)}$ is only computed once
as $\alpha_{k}$ varies.

Current descents implemented in Optimistix are: steepest descent, nonlinear
conjugate gradient descent (Shewchuk, 1994), direct and indirect damped Newton
(used in the Levenberg-Marquardt algorithm (Moré, 1978; Nocedal & Wright,
2006)), and dogleg (Nocedal & Wright, 2006)[section 4.1]. The indirect damped
Newton and dogleg descents both approximately solve (2) (see (Nocedal &
Wright, 2006)[sections 4.1 & 4.3]) and are capable of using any user-provided
function which maps a PyTree to a scalar as the norm $\|\cdot\|$.

Flattening the bilevel optimisation problem

The classical approach to implementing a line search is as a bilevel
optimisation problem. The inner loop to obtain the solution to the line
search, the outer loop over the iterates $x_{k}$.

The termination condition of the line-search in the inner loop, such as the
Armijo or Wolfe conditions, may evaluate $f(x_{k+1})$ when determining whether
to halt the inner loop. This can result in an extra function evaluation or
additional logic at each step of the outer loop, as $f(x_{k+1})$ is usually
computed in each step of the outer loop as well (eg. when computing the
gradient of $f$.)

Not every search has termination conditions which require computing
$f(x_{k+1})$, so we cannot avoid an extra function evaluation by requiring the
inner loop to evaluate $f(x_{k+1})$ and pass this to the outer loop.

Optimistix gets around this by using step-rejection: a step $x_{k+1}$ is
proposed, and the outer loop computes $f(x_{k+1})$, and at the next step the
search decides whether to accept or reject $x_{k+1}$.

##  4 Converting Between Problem Types

Optimistix handles four types of nonlinear optimisation task, each with a
simple user API:

  * •

Minimisation:

`optimistix.minimise`.

  * •

Nonlinear least-squares solves:

`optimistix.least_squares`.

  * •

Root-finding:

`optimistix.root_find`.

  * •

Fixed-point iteration:

`optimistix.fixed_point`.

For example, to find a root of the function `fn` we call:

⬇

solver = optimistix.Newton(

rtol=1e-3, atol=1e-3

)

optimistix.root_find(fn, solver, y0)

where `y0` is an initial guess and `optimistix.Newton` is the optimistix
solver used to find the root. `rtol` and `atol` are relative and absolute
tolerances used to deteremine whether the solver should terminate, as
described in section 6.

Optimistix can convert between these problem types when appropriate. A fixed-
point iteration $f(x)=x$ is converted to a root-find problem by solving
$f(x)-x=0$. A root-find problem $f(x)=0$ is converted to a nonlinear least-
squares problem by using entries of $f(x)$ as the residuals:
$\min_{x}\sum_{i}f(x)_{i}^{2}$. A least-squares problem is already a special
case of minimisation problem, so we take the objective function
$g(x)=\sum_{i}f(x)_{i}^{2}$ and provide this directly to a minimiser.

Converting between problem types is not just a neat trick, but a common
pattern in optimisation, especially for root-finding and fixed-point
iteration. While standard methods for root-finding and fixed-point iteration,
such as Newton, chord, bisection, and fixed-point iteration (Nocedal & Wright,
2006; Bonnans et al., 2006) are available in Optimistix, these methods tend
not to work as well for complex, highly nonlinear problems. In this case,
common practice is to convert the fixed-point or root-find problem into a
least-squares or minimisation problem, and solve using a minimiser or
nonlinear least-squares solver (Nocedal & Wright, 2006)[section 11].

The conversion between problem types outlined above is done automatically: the
solver is checked, and the problem automatically converted and lowered to a
solve of that type. For example,

⬇

solver = optimistix.BFGS(

rtol=1e-3, atol=1e-3

)

optimistix.root_find(fn, solver, y0)

will convert the root-find problem to a minimisation problem and use the BFGS
algorithm (Nocedal & Wright, 2006)[section 6.1] to find a root.

##  5 Automatic Differentiation

All solvers in Optimistix are iterative, and provide two methods for automatic
differentiation:

  * •

Differentiation via the implicit function theorem (Blondel et al., 2021)

  * •

Differentiation via online treeverse (Stumm & Walther, 2010; Wang et al.,
2009).

The former is the default for both forward-mode autodiff and backpropagation
in Optimistix. Both are known automatic differentiation techniques; however,
to the best of our knowledge, Optimistix is the first nonlinear optimisation
library in JAX to feature online treeverse. We now discuss each of these in
turn.

###  5.1 Forward-Autodiff and Backpropagation Using the Implicit Function
Theorem

Consider a continuously differentiable function
$F:\mathbb{R}^{n}\times\mathbb{R}^{m}\to\mathbb{R}$ and values
$x_{0}\in\mathbb{R}^{n}$ and $\theta_{0}\in\mathbb{R}^{m}$ satisfying

$\displaystyle F(x_{0},\theta_{0})$ $\displaystyle=0$ (12)
$\displaystyle\det\left(\frac{\mathrm{d}F}{\mathrm{d}x}(x_{0},\theta_{0})\right)$
$\displaystyle\neq 0.$ (13)

The implicit function theorem (IFT) states there exists neighborhoods
$x_{0}\in N_{x_{0}}$ and $\theta_{0}\in N_{\theta_{0}}$ and a differentiable
function $x^{*}:N_{\theta_{0}}\to N_{x_{0}}$ such that

$\displaystyle x^{*}(\theta_{0})$ $\displaystyle=x_{0}$ $\displaystyle
F(x^{*}(\theta),\theta)$ $\displaystyle=0\ \ \forall x\in N_{\theta_{0}}.$

and

$\frac{\mathrm{d}x^{*}}{\mathrm{d}\theta}=-\left(\frac{\mathrm{d}F}{\mathrm{d}x}(x_{0},\theta_{0})\right)^{-1}\frac{\mathrm{d}F}{\mathrm{d}\theta}(x_{0},\theta_{0})$
(14)

The implicit function theorem gives means of calculating the derivative
$\frac{\mathrm{d}x^{*}}{\mathrm{d}\theta}$, despite the fact that $x^{*}$ is
only defined implicitly, and is the solution to a nonlinear problem.

Every optimisation task in Optimistix (minimisation, …) may be rewritten in
the form of equation (12), and thus the derivative of the solution (the
argminimum, …) with respect to any parameters $\theta$ may be found using
equation (14).

For root finding, equations (12-13) apply directly. For fixed-point iteration,
the system of nonlinear equations is $F(x,\theta)=f(x,\theta)-x$. For
minimisation and least-squares,
$F(x,\theta)=\frac{\mathrm{d}f}{\mathrm{d}x}(x,\theta)$, as the minimum is
found at a critical points of $f$.

We highly recommend (Blondel et al., 2021) for many more details on using the
implicit function theorem to differentiate nonlinear optimisation routines.

###  5.2 Backpropagation using Online Treeverse

An optimisation routine can also be differentiated naively, by backpropagating
through the components of the optimiser at each iteration $k$. This method
works even when the assumptions (12-13) of the implicit function theorem are
not satisfied, or the full Jacobian $\frac{dF}{dx}$ cannot be constructed due
to memory constraints. Further, it can be faster in some cases, especially
when the dimensionality of the domain of the optimisation problem is much
larger than the number of iterations $K$. See (Ablin et al., 2020) for more
details comparing the efficiency of these methods.

The main downside of the naive backpropagation algorithm for nonlinear
optimisation is memory cost. It costs $\mathcal{O}(KM)$ memory to
backpropagate through an optimisation, were $M$ is the memory cost of
backpropagating through a single iteration of the optimiser, and $K$ is the
total number of iterations used. The classical way around this limitation is
checkpointing, which chooses a fixed number $N\in\mathbb{N}$ of checkpoint
iterations, and stores the values of the optimiser only at those checkpointed
iterations, reducing the memory cost to $\mathcal{O}(NM)$. Iterations which
are not checkpointed are recomputed from the previous checkpoint, trading off
memory cost for extra computation time.

Online treeverse is a checkpointing backpropagation scheme designed for
performing backpropagation over loops where the number of iterations is
a-priori unknown. Online treeverse dynamically updates its checkpoint
locations as the loop continues, keeping optimal computation times for a fixed
number of checkpoints $c$. For further details see (Stumm & Walther, 2010) and
(Wang et al., 2009).

##  6 Termination Condition

Optimistix introduces a Cauchy convergence criteria, where iterations are
stopped if

$\left\lvert\frac{f(x_{k+1})-f(x_{k})}{\varepsilon_{a}+\varepsilon_{r}f(x_{k})}\right\rvert<1\text{
and
}\left\lVert\frac{x_{k+1}-x_{k}}{\varepsilon_{a}+\varepsilon_{r}x_{k}}\right\rVert<1$

for $\varepsilon_{a}$ an absolute tolerance set by the user and
$\varepsilon_{r}$ a relative tolerance set by the user, and division is
performed elementwise. By default, the norm $\left\lVert\cdot\right\rVert$ is
$\left\lVert\cdot\right\rVert_{\infty}$, ie. the maximum over the elementwise
absolute differences in $x_{k+1}$ and $x_{k}$. However, users can set this
themselves to be any function from a PyTree to a scalar.

To justify this choice, we note that the optimisation literature is not
consistent with the termination criteria used (Nocedal & Wright, 2006),
(Bonnans et al., 2006), (Conn et al., 2000), (Moré et al., 1980). For example,
(SAS Institute Inc, 2004) outlines 20 different convergence criteria currently
in use for nonlinear optimisation.

Without an established convention, we choose this termination condition to
match the better established convention for the literature on solving
numerical differential equations (Hairer et al., 2008; Hairer & Wanner, 2002).
This is reasonable choice for optimisation, resembling (but not exactly
matching) a combination of the ‘X-convergence’ and ‘F-convergence’ criteria in
the highly successful MINPACK (Moré et al., 1980) optimisation suite.
Alternatively, it is roughly a combination of all four of the ‘FTOL’,
‘ABSFTOL’, ‘XTOL’, and ‘ABSXTOL’ criteria found in (SAS Institute Inc, 2004).
This also ensures a consistent approach between Optimistix and existing JAX +
Equinox libraries, such as Diffrax (Kidger, 2021).

##  7 Experiments

###  7.1 Summary of Experiments

We demonstrate the fast compilation and run times of Optimistix by
benchmarking our solvers for both runtime and compile time against the state-
of-the-art optimisation algorithms in JAXopt (Blondel et al., 2021) and SciPy
(Virtanen et al., 2020). For SciPy, we only benchmark runtimes, since SciPy
does not have a notion of compile times (however, we still compile the
function passed to the SciPy solver. The overhead of this function compilation
is excluded.)

Table 1: Average minimum runtime comparison Average minimum runtime
(milliseconds)
---
Optimiser |  Optx |  JAXopt |  SciPy
BFGS |  0.836 |  16.4 |  250
Nonlinear CG |  0.466 |  6.27 |  208
LM1 |  - |  1.70 |  -
LM2 |  144 |  - |  188
Gauss-Newton |  4.54 |  3.08 |  N/A
Table 2: Average compile time comparison Average compile time (milliseconds)
---
Optimiser |  Optx |  JAXopt
BFGS |  251 |  854
Nonlinear CG |  130 |  849
LM1 |  - |  554
LM2 |  296 |  -
Gauss-Newton |  288 |  416

Tables 1-2 present the average minimum runtime and average compile time for a
set of over 100 test problems. Each problem is solved 10 times to a fixed
tolerance with a max budget of 2000 iterations, and the minimum over those 10
iterations is taken to be the runtime of the solver for that problem. The
reported values are averaged over all the problems in the test set. The full
methodology and additional comparisons are described in appendix A.

The average compile times for Optimistix (Optx) are up to 6.5 times faster
than the average compile times for JAXopt. The runtimes include all operations
for a full solve; notably, it includes the difference in the number of
iterations arising from convergence criteria and step-sizes in Optimistix and
JAXopt. The times for a single iteration are presented in appendix A, where
JAXopt is marginally faster for a single iteration than Optimistix.

The difference between Levenberg-Marquardt (LM) algorithms LM1 and LM2 relates
to a technical detail in their implementation. At each iteration, LM performs
a linear solve of the equation $(J_{k}^{T}J_{k}+\lambda I)p=-J^{T}r_{k}$,
where $J_{k}$ is the Jacobian and $r_{k}$ is the residual of the objective
function at step $k$. This linear system can be solved directly with a
Cholesky or conjugate gradient solver, but the condition number of $J_{k}$ is
squared in $J_{k}^{T}J_{k}$. This is the approach used in LM1, which JAXopt
implements. Alternatively, this linear system can be transformed into an
equivalent least-squares problem which does not square the condition number,
but which requires a slower linear solve. This is the approach used in LM2,
which Optimistix and SciPy use. We accept the slower solve in return for
increased robustness and accuracy. This is detailed in (Moré, 1978) and
(Nocedal & Wright, 2006)[section 10]. A similar tradeoff is made for Gauss-
Newton. The difference in robustness is noticeable: Optimistix LM failed to
solve 3 of the test problems to the specified accuracy within 2000 iterations,
SciPy failed to solve 7, and JAXopt failed to solve 13.

Compile times in the table are 3-4 orders of magnitude longer than runtimes.
For solves with fewer than many thousands of iterations, the total time is
dominated by compilation, where Optimistix is faster.

###  7.2 Benchmarking with Performance Profiles

Performance profiles are an established technique for benchmarking
optimisation software (Dolan & Moré, 2001; Beiranvand et al., 2017).
Performance profiles are defined in terms of solver performance ratios.

Formally, let $\mathcal{S}$ be a collection of optimisers, and $\mathcal{P}$ a
collection of problems. Letting

$\displaystyle R_{s,p}$ $\displaystyle=\text{the time for optimiser }s\text{
to solve problem }p,$ excluding compilation time. $\displaystyle C_{s,p}$
$\displaystyle=\text{the time for optimiser }s\text{ to compile problem }p,$

the runtime performance ratio and compile time performance ratios are:

$\displaystyle r^{R}_{s,p}$ $\displaystyle=\frac{R_{s,p}}{\min\\{R_{s,p}\colon
s\in\mathcal{S}\\}}$ $\displaystyle r^{C}_{s,p}$
$\displaystyle=\frac{C_{s,p}}{\min\\{C_{s,p}\colon s\in\mathcal{S}\\}}$

The runtime and compile time performance profiles are then:

$\displaystyle\rho^{R}_{s}(\tau)=\frac{1}{|\mathcal{P}|}|\\{p\in\mathcal{P}\colon
r^{R}_{s,p}\leq\tau\\}|$
$\displaystyle\rho^{C}_{s}(\tau)=\frac{1}{|\mathcal{P}|}|\\{p\in\mathcal{P}\colon
r^{C}_{s,p}\leq\tau\\}|$

ie. $\rho^{R}_{s}(\tau)$ is the proportion of problems solver $s$ solved
within a factor $\tau$ of the minimum runtime attained by any solver in
$\mathcal{S}$. The value $\rho^{R}_{s}(1)$ is the proportion of problems where
solver $s$ had the best runtime performance in $\mathcal{S}$. For large values
of $\tau$, $\rho^{R}_{s}(\tau)$ represents the proportion of problems solver
$s$ accurately solved. Similarly, $\rho^{C}_{s}(\tau)$ is the proportion of
problems solved within a factor of $\tau$ of the best compile time.

Figure 1: BFGS runtime comparison Figure 2: Levenberg-Marquardt runtime
comparison Figure 3: BFGS compile time comparison Figure 4: Levenberg-
Marquardt compile time comparison

In figures 1-2 we see the runtime performance profiles of BFGS, and Levenberg-
Marquardt with Optimsitix in red, JAXopt in blue, and SciPy in purple (no
distinction is made between LM1 and LM2 in the plot.) The plots are log scale
on the x-axis (representing $\tau$ in the performance profile,) and range from
$1$ to $2^{7}$. Optimistix performs the best of all solvers for BFGS. For
Levenberg-Marquardt it is generally slower than JAXopt but solves more
problems in total.

It may appear as though SciPy has not solved many of the test problems;
however, it is just more than 250 times slower and therefore does not show on
the plots.

In figures 3-4 we see a compile time comparison of BFGS and Levenberg-
Marquardt. Optimistix outperforms JAXopt on every problem for both optimisers.

##  8 Related Work

Gradient-based minimisation  Optax (Babuschkin et al., 2020) implements
algorithms for gradient-based minimisation in JAX. Optimistix has a broader
scope than Optax, handling other nonlinear optimisation tasks and second-order
optimisers. Optax and Optimistix are compatible libraries, and Optax
minimisers can be used within Optimistix via `optimistix.OptaxMinimiser`.

General-purpose optimisation  JAXopt (Blondel et al., 2021) is an excellent
differentiable linear and nonlinear optimisation library in JAX which includes
many optimisation tasks, including ones out-of-scope for Optimistix such as
quadratic programming and non-smooth optimisation.

We see Optimistix and JAXopt as having fundamentally different scopes and core
abstractions. JAXopt has a larger scope than Optimistix, but its focus is not
on modularity. While specific choices of line search algorithms can be
interchanged in some cases, for the most part introducing a new optimiser
requires writing the algorithm in its entirety. We recommend JAXopt for
general optimisation tasks.

SciPy (Virtanen et al., 2020) also offers general-purpose nonlinear
optimisation routines, many of which call into MINPACK (Moré et al., 1980), a
nonlinear optimisation library written in Fortran. SciPy includes many
optimisation tasks, including minimisation, least-squares, root-finding, and
global optimisation. However, SciPy implementations are not differentiable,
and are generally difficult to extend.

##  9 Conclusion

We introduced Optimistix, a nonlinear optimisation software for minimisation,
least-squares, root-finding, and fixed-point iteration with a novel, modular
approach.

##  10 Impact Statement

This paper presents work whose goal is to advance the field of Machine
Learning. There are many potential societal consequences of our work, none
which we feel must be specifically highlighted here.

## References

  * Ablin et al. (2020) Ablin, P., Peyré, G., and Moreau, T.  Super-efficiency of automatic differentiation for functions defined as a minimum, 2020.
  * Ali et al. (2005) Ali, M. M., Khompatraporn, C., and Zabinsky, Z. B.  A numerical evaluation of several stochastic algorithms on selected continuous global optimization test problems.  _Journal of Global Optimization_ , 2005.
  * Andrei (2008) Andrei, N.  An unconstrained optimization test functions collection.  In _Advanced Modeling and Optimization, Volume 10_ , 2008.  URL https://api.semanticscholar.org/CorpusID:63504217.
  * Averick et al. (1992) Averick, B., Carter, R., Moré, J., and Xue, G.-L.  The minpack test problem collection.  Technical report, Office of Energy Research, US Department of Energy, 1992\.
  * Babuschkin et al. (2020) Babuschkin, I., Baumli, K., Bell, A., Bhupatiraju, S., Bruce, J., Buchlovsky, P., Budden, D., Cai, T., Clark, A., Danihelka, I., Dedieu, A., Fantacci, C., Godwin, J., Jones, C., Hemsley, R., Hennigan, T., Hessel, M., Hou, S., Kapturowski, S., Keck, T., Kemaev, I., King, M., Kunesch, M., Martens, L., Merzic, H., Mikulik, V., Norman, T., Papamakarios, G., Quan, J., Ring, R., Ruiz, F., Sanchez, A., Sartran, L., Schneider, R., Sezener, E., Spencer, S., Srinivasan, S., Stanojevi´c, M., Stokowiec, W., Wang, L., Zhou, G., and Viola, F.  The deepmind jax ecosystem, 2020.  URL http://github.com/deepmind.
  * Beiranvand et al. (2017) Beiranvand, V., Hare, W., and Lucet, Y.  Best practices for comparing optimization algorithms.  _Optimization and Engineering_ , 18(4):815–848, 2017.
  * Bezgin et al. (2022) Bezgin, D. A., Buhendwa, A. B., and Adams, N. A.  Jax-fluids: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows.  _Computer Physics Communications_ , pp. 108527, 2022.
  * Blondel et al. (2021) Blondel, M., Berthet, Q., Cuturi, M., Frostig, R., Hoyer, S., Llinares-López, F., FabianPedregosa, and Vert, J.-P.  Efficient and modular implicit differentiation.  _arXiv preprint arXiv:2105.15183_ , 2021.
  * Bonnans et al. (2006) Bonnans, F., Gilbert, C., Lemaréchal, C., and Sagastizábal, C.  _Numerical Optimization: Theoretical and Practical Aspects_.  Springer Berlin Heidelberg, 2006.
  * Bradbury et al. (2018) Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q.  Jax: composable transformations of python+numpy programs, 2018.  URL http://github.com/google/jax.
  * Carroll (2023) Carroll, C.  Bayeux.  Accessed 2024, 2023.  URL https://github.com/jax-ml/bayeux.
  * Conn et al. (2000) Conn, A. R., Gould, N. I. M., and Toint, P. L.  _Trust Region Methods_.  Society for Industrial and Applied Mathematics, 2000.
  * Dolan & Moré (2001) Dolan, E. D. and Moré, J. J.  Benchmarking optimization software with performance profiles.  _CoRR_ , cs.MS/0102001, 2001.
  * Dresdner et al. (2022) Dresdner, G., Kochkov, D., Norgaard, P., Zepeda-Núñez, L., Smith, J. A., Brenner, M. P., and Hoyer, S.  Learning to correct spectral methods for simulating turbulent flows.  _arXiv_ , 2022.  doi: 10.48550/ARXIV.2207.00556.  URL https://arxiv.org/abs/2207.00556.
  * Duchi et al. (2011) Duchi, J., Hazan, E., and Singer, Y.  Adaptive subgradient methods for online learning and stochastic optimization.  _Journal of Machine Learning Research_ , 12:2121–2159, 2011\.
  * Freeman et al. (2021) Freeman, C. D., Frey, E., Raichuk, A., Girgin, S., Mordatch, I., and Bachem, O.  Brax - a differentiable physics engine for large scale rigid body simulation, 2021.  URL http://github.com/google/brax.
  * Golovin et al. (2017) Golovin, D., Solnik, B., Moitra, S., Kochanski, G., Karro, J., and Sculley, D.  Google vizier: A service for black-box optimization.  In _Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, Halifax, NS, Canada, August 13 - 17, 2017_ , pp. 1487–1495. ACM, 2017.  URL https://doi.org/10.1145/3097983.3098043.
  * Gupta et al. (2018) Gupta, V., Koren, T., and Singer, Y.  Shampoo: Preconditioned stochastic tensor optimization, 2018.
  * Hairer & Wanner (2002) Hairer, E. and Wanner, G.  _Solving Ordinary Differential Equations II Stiff and D ifferential-Algebraic Problems_.  Springer, Berlin, second revised edition edition, 2002.
  * Hairer et al. (2008) Hairer, E., Nørsett, S., and Wanner, G.  _Solving Ordinary Differential Equations I Nonstiff P roblems_.  Springer, Berlin, second revised edition edition, 2008.
  * Hall et al. (2023a) Hall, D., Zhou, I., and Liang, P.  Haliax.  Accessed 2023, 2023a.  URL https://github.com/stanford-crfm/haliax.
  * Hall et al. (2023b) Hall, D., Zhou, I., and Liang, P.  Levanter — legible, scalable, reproducible foundation models with jax.  Accessed 2023, 2023b.  URL https://github.com/stanford-crfm/levanter.
  * Jamil & Yang (2013) Jamil, M. and Yang, X. S.  A literature survey of benchmark functions for global optimisation problems.  _International Journal of Mathematical Modelling and Numerical Optimisation_ , 4(2):150, 2013.
  * Jumper et al. (2021) Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Žídek, A., Potapenko, A., Bridgland, A., Meyer, C., Kohl, S. A. A., Ballard, A. J., Cowie, A., Romera-Paredes, B., Nikolov, S., Jain, R., Adler, J., Back, T., Petersen, S., Reiman, D., Clancy, E., Zielinski, M., Steinegger, M., Pacholska, M., Berghammer, T., Bodenstein, S., Silver, D., Vinyals, O., Senior, A. W., Kavukcuoglu, K., Kohli, P., and Hassabis, D.  Highly accurate protein structure prediction with AlphaFold.  _Nature_ , 596(7873):583–589, 2021.  doi: 10.1038/s41586-021-03819-2.
  * Kidger (2021) Kidger, P.  _On Neural Differential Equations_.  PhD thesis, University of Oxford, 2021.
  * Kidger & Garcia (2021) Kidger, P. and Garcia, C.  Equinox: neural networks in jax via callable pytrees and filtered transformations.  _Differentiable Programming workshop at Neural Information Processing Systems 2021_ , 2021.
  * Kingma & Ba (2017) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization, 2017.
  * Moré (1978) Moré, J. J.  The levenberg-marquardt algorithm: Implementation and theory.  In _Numerical Analysis_ , pp. 105–116, Berlin, Heidelberg, 1978\. Springer Berlin Heidelberg.
  * Moré et al. (1981) Moré, J., Garbow, B., and Hillstrom, K.  Testing unconstrained optimistation software.  Technical report, US Department of Energy, 1981.
  * Moré & Thuente (1994) Moré, J. J. and Thuente, D. J.  Line search algorithms with guaranteed sufficient decrease.  _ACM Transactions on Mathematical Software_ , 20:286–307, 1994.  ISSN 0098-3500.
  * Moré et al. (1980) Moré, J. J., Garbow, B. S., and Hillstrom, K. E.  User guide for minpack-1.  Technical report, Argonne National Lab. (ANL), 1980.
  * Nocedal & Wright (2006) Nocedal, J. and Wright, S.  _Numerical Optimization (Second Edition)_.  Springer New York, 2006.
  * Pastrana (2023) Pastrana, R.  flowmc.  Accessed 2023, 2023.  URL https://github.com/kazewong/flowMC.
  * Pastrana et al. (2023) Pastrana, R., Oktay, D., Adams, R. P., and Adriaenssens, S.  JAX FDM: A differentiable solver for inverse form-finding.  In _ICML 2023 Workshop on Differentiable Almost Everything: Differentiable Relaxations, Algorithms, Operators, and Simulators_ , 2023.  URL https://openreview.net/forum?id=Uu9OPgh24d.
  * Rader et al. (2023) Rader, J., Lyons, T., and Kidger, P.  Lineax: unified linear solves and linear least-squares in jax and equinox.  _AI for science workshop at Neural Information Processing Systems 2023, arXiv:2311.17283_ , 2023.
  * Rosenbrock (1960) Rosenbrock, H. H.  An Automatic Method for Finding the Greatest or Least Value of a Function.  _The Computer Journal_ , 3(3):175–184, 01 1960\.
  * SAS Institute Inc (2004) SAS Institute Inc.  _SAS/IML 9.1 User’s Guide_.  SAS Institute Inc, 2004.
  * Sherali & Ulular (1990) Sherali, H. D. and Ulular, O.  Conjugate gradient methods using quasi-newton updates with inexact line searches.  _Journal of Mathematical Analysis and Applications_ , 150:359–377, 8 1990.  ISSN 0022247X.
  * Shewchuk (1994) Shewchuk, J. R.  An introduction to the conjugate gradient method without the agonizing pain.  Technical report, Carnegie Mellon University, USA, 1994.
  * Singh (2022) Singh, A.  Eqxvision.  Accessed 2023, 2022.  URL https://github.com/paganpasta/eqxvision.
  * Song et al. (2022) Song, X., Perel, S., Lee, C., Kochanski, G., and Golovin, D.  Open source vizier: Distributed infrastructure and api for reliable and flexible black-box optimization.  In _Automated Machine Learning Conference, Systems Track (AutoML-Conf Systems)_ , 2022.
  * Stanojević & Sartran (2023) Stanojević, M. and Sartran, L.  SynJax: Structured Probability Distributions for JAX.  _arXiv preprint arXiv:2308.03291_ , 2023.
  * Steihaug (1983) Steihaug, T.  The conjugate gradient method and trust regions in large scale optimization.  _SIAM Journal on Numerical Analysis_ , 20(3):626–637, 1983.
  * Stumm & Walther (2010) Stumm, P. and Walther, A.  New algorithms for optimal online checkpointing.  _SIAM Journal on Scientific Computing_ , 32(2):836–854, 2010.  doi: 10.1137/080742439.
  * Virtanen et al. (2020) Virtanen, P., Gommers, R., Oliphant, T. E., Haberland, M., Reddy, T., Cournapeau, D., Burovski, E., Peterson, P., Weckesser, W., Bright, J., van der Walt, S. J., Brett, M., Wilson, J., Millman, K. J., Mayorov, N., Nelson, A. R. J., Jones, E., Kern, R., Larson, E., Carey, C. J., Polat, İ., Feng, Y., Moore, E. W., VanderPlas, J., Laxalde, D., Perktold, J., Cimrman, R., Henriksen, I., Quintero, E. A., Harris, C. R., Archibald, A. M., Ribeiro, A. n. H., Pedregosa, F., van Mulbregt, P., and SciPy 1.0 Contributors.  SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python.  _Nature Methods_ , 17:261–272, 2020.
  * Wang (2023) Wang, P.  Palm - jax.  Accessed 2023, 2023.  URL https://github.com/lucidrains/PaLM-jax.
  * Wang et al. (2009) Wang, Q., Moin, P., and Iaccarino, G.  Minimal repetition dynamic checkpointing algorithm for unsteady adjoint calculation.  _SIAM Journal on Scientific Computing_ , 31(4):2549–2567, 2009.  doi: 10.1137/080727890.

##  Appendix A Experiment details

###  A.1 Methodology

We choose as our comparison set BFGS, nonlinear CG, Levenberg-Marquardt, and
Gauss-Newton, as these are the four nontrivial minimisation/least-squares
algorithms shared by Optimistix and JAXopt (ie. gradient descent is excluded
intentionally.) Where solver implementations differ, we try and match them as
closely as possible.

Specifically, where algorithms differ from the base implementations of the
algorithms is:

  * •

BFGS uses backtracking Armijo line search in both implementations (default for
JAXopt is zoom.)

  * •

Gauss-Newton uses a conjugate gradient solver on the normal equations in both
implementations (default for Optimisix is the polyalgorithm
`AutoLinearSolver(well_posed=None)` from (Rader et al., 2023).

The set of test problems consists of 104 minimisation problems, 63 of which
are least-squares problems, taken from the test collections (Jamil & Yang,
2013), (Ali et al., 2005), (Andrei, 2008), (Averick et al., 1992), and (Moré
et al., 1981). Levenberg-Marquardt and Gauss-Newton are only ran on the least-
squares problems. Solvers are initialised at canonical initialisations when
available (see (Moré et al., 1981) and (Andrei, 2008)).

Runtime is a noisy measurement. During an experiment, the computer an
experiment is running on may be running background processes which we cannot
control. To mitigate this, when assessing runtime we run each problem 10 times
and take the minimum over these repeats. This indicates roughly what the ”best
we can expect” from a given solver is.

The global minima are not known for all test problems. On problems where
minima are known, we require that

$\frac{|f(x_{N})-f(x^{*})|}{\epsilon_{a}+\epsilon_{r}f(x^{*})}<1$

for an absolute tolerance $\epsilon_{a}$ and relative tolerance
$\epsilon_{r}$, and $x_{N}$ the argmin found by the solver. If this condition
fails, the runtime is set to `jnp.inf`. This information is automatically
incorporated in to the performance profile; however, it is not included in the
tables of average runtime performance. This is because there is no obvious way
to penalise nonconvergence when comparing runtimes. To get around this, we
assure that in all cases the number of failures for an Optimistix solver was
less than or equal to the number of failures for a JAXopt solver, as to not
give an unfair advantage to Optimistix.

Though this convergence criteria looks similar to the Cauchy termination
condition for Optimistix, it is applied to $f(x_{N})-f(x^{*})$ and not
$f(x_{k+1})-f(x_{k})$. $f(x^{*})$ is an unknown quantity to both solvers, so
there is no advantage provided by the convergence criterion in Optimistix.

Finally, while the general formulation of performance profile allows for
$\mathcal{S}$ to contain more than two optimisers, using performance profiles
with more than two optimisers requires careful interpretation. For example,
when comparing three optimisers, it is possible that the method which appears
to be second best performs worse than the method which appears to be third
best when these two methods are compared directly.

For this reason, we provide all the pairwise comparisons of Optimistix vs
JAXopt and SciPy in A.2, with the Optimsitix solver in red and the comparison
solver (JAXopt or SciPy) in blue.

###  A.2 Further experiments

For simplicity, we included a number of performance profiles with one-on-one
comparisons of Optimistix vs JAXopt and Optimistix vs SciPy. Also included is
a table of average minimum runtimes for a single optimiser run. Ultimately, no
user will use only a single optimiser run, so the aggregate information
presented in section 7 should be taken as more informative.

Average minimum runtime for a single optimiser run (milliseconds)
---
Optimiser |  Optimistix |  JAXopt |  SciPy
BFGS |  0.151 |  0.138 |  2.20
Nonlinear CG |  0.123 |  0.0645 |  2.61
LM |  0.515 |  0.0917 |  3.31
Gauss-Newton |  0.393 |  0.0578 |  N/A

Figure 5: Optimistix vs JAXopt BFGS

Figure 6: Optimistix vs JAXopt nonlinear CG

Figure 7: Optimistix vs JAXopt Gauss-Newton

Figure 8: Optimistix vs SciPy BFGS (left) and nonlinear CG (right) Figure 9:
Optimistix vs SciPy Levenberg-Marquardt runtimes


# S-RL Toolbox: Environments, Datasets and Evaluation Metrics for State
Representation Learning

Antonin Raffin<EMAIL_ADDRESS>Ashley Hill
<EMAIL_ADDRESS>René Traoré<EMAIL_ADDRESS>Timothée Lesort<EMAIL_ADDRESS>Natalia Díaz-Rodríguez
<EMAIL_ADDRESS>David Filliat<EMAIL_ADDRESS>
U2IS, ENSTA ParisTech / INRIA FLOWERS Team http://flowers.inria.fr
Palaiseau, France

###### Abstract

State representation learning aims at learning compact representations from
raw observations in robotics and control applications. Approaches used for
this objective are auto-encoders, learning forward models, inverse dynamics or
learning using generic priors on the state characteristics. However, the
diversity in applications and methods makes the field lack standard evaluation
datasets, metrics and tasks. This paper provides a set of environments, data
generators, robotic control tasks, metrics and tools to facilitate iterative
state representation learning and evaluation in reinforcement learning
settings.

Keywords: Deep learning, reinforcement learning, state representation
learning, robotic priors

##  1 Introduction

Robotics control relies on compact and expressive representations of sensor
data, as the task objectives are often expressed in much smaller dimensions
than the sensor space dimension (e.g., the position of an object versus the
size of an image). These representations are usually hand crafted by human
experts, but deep-learning now makes it possible to avoid this feature
engineering by using end-to-end learning (e.g., learning a policy from raw
pixels). However, such approach is mostly possible in simulation because it
requires a huge amount of training data –usually millions of samples– that
makes it impractical in the real world. To overcome this issue, State
Representation Learning (SRL) methods (Jonschkowski, 2018) can be used to
create an intermediate representation that should contain only useful
information to control a robot and thus simplify the policy learning task.

Many different SRL approaches have been proposed (see (Lesort et al., 2018)
for a review), but comparing their performances is challenging. A common
approach to evaluate learned representations is to compare performance in a
Reinforcement Learning (RL) setting. However, because of the instability of RL
algorithms and their cost, this should not be the only method used to assess
learned states. Moreover, this approach gives no mean to interpret a state
representation, making it difficult to understand which information is encoded
in this representation.

While Reinforcement Learning has well established benchmarks, SRL has no
metrics nor universal criterion to compare the different approaches. With that
in mind, we propose a set of environments with an increasing difficulty,
designed for comparing State Representation Learning (SRL) methods for robotic
control. We also introduce qualitative and quantitative metrics along with
visualization tools to facilitate the development and the comparison of SRL
algorithms. The proposed framework allows fast iteration and eases research of
new SRL methods by making it easy to produce statistically relevant results:
the simulated environments run at 250 FPS on a 8-core machine that allows to
train a RL agent on 1 Million steps in only 1h (or to generate 20k samples in
less than 2 min)111Environments, code and data are available at
https://github.com/araffin/robotics-rl-srl.

In this paper, we first present quickly the reinforcement learning framework
and the main state representation learning approaches that are implemented,
before presenting the SRL Toolbox environments and datasets, the qualitative
and quantitative evaluation methods, and a set of experiments illustrating the
performances of the implemented approaches.

##  2 Reinforcement Learning and State Representation Learning

In this section, we introduce Reinforcement Learning (RL), as well as the
State Representation Learning (SRL) approaches we integrated in our framework.

###  2.1 Reinforcement Learning

In RL, an agent must learn to select the best action to maximize a reward it
will receive. More formally, in a given state $s_{t}$ ($t$ denotes the current
time-step), an agent performs an action $a_{t}$ and receives a reward $r_{t}$.
The learned behaviour, that should maximize the long-term discounted reward by
mapping states to actions, is called a policy: $a_{t}=\pi(s_{t})$.

In the most common settings, the state $s_{t}$ is either a low dimension
representation given by an expert human, or corresponds to the raw observation
(end-to-end learning). In order to differentiate these cases, we introduce the
observation $o_{t}$ that correspond to the raw sensor data, and use the term
state only to refer to low dimensional representations that can be provided by
humans, or can be learned using SRL.

In this paper, we work in the context of Markov Decision Processes (MDP),
where the next state of the system only depends on the previous state and the
taken action. As described above, the term observations refers to raw sensor
data (mostly images), but do not imply partial observability as assumed in
Partially Observable Markov Decision Processes (POMDP).

###  2.2 State Representation Learning

SRL (Lesort et al., 2018) aims at learning compact representations from raw
observations (e.g., learn a position $(x,y)$ directly from raw pixels) without
explicit supervision. Most of the time, the goal is to use that representation
to solve a task with RL. The idea is that a low-dimensional representation
should only keep the useful information and reduce the search space, thus
contributing to address two main challenges of RL: sample inefficiency and
instability. Moreover, a state representation learned for a particular task
may be transferred to related tasks and therefore speed up learning in
multiple task settings.

Using RL notations, SRL corresponds to learning a transformation222In
practice, the learned transformation is a neural network $\varphi$ from the
observation space $\mathcal{O}$ to the state space $\mathcal{S}$. Then, a
policy $\pi$, that takes a state $s_{t}\in\mathcal{S}$ as input and outputs
action $a_{t}$, is learned to solve the task:

$o_{t}\xrightarrow[SRL]{\varphi}s_{t}\xrightarrow[RL]{\pi}a_{t}$ (1)

In the next sections, we present approaches of SRL that are implemented in our
toolbox. Each method is not mutually exclusive and can be combined to create
new models.

####  2.2.1 Auto-encoders (AE, VAE)

A first approach to learn a state representation is to compress the
observation into a low dimensional state that is sufficient to reconstruct the
observation. This approach does not take advantage of the robotic context
because it ignores the possible actions, therefore it is often associated with
different objectives (e.g. forward model), and provides a performance
baseline.

We integrated Auto-Encoders (Baldi, 2012) and Variational Auto-encoders (VAE)
(Kingma and Welling, 2013), i.e., auto-encoders that enforce the latent
variables to follow a given distribution.

####  2.2.2 Robotic Priors

To build a relevant representation of states, one can use prior knowledge
about the dynamics or physics of the world. This knowledge can account for
temporal continuity or causality principles that reflect the interactions of
an agent with its environment (Jonschkowski and Brock, 2015; Jonschkowski et
al., 2017). Robotic Priors are defined as objective functions that constrain
the state representation.

####  2.2.3 Forward and Inverse Models

The dynamics of the world can be integrated by learning a forward model that
predicts state $s_{t+1}$ given state $s_{t}$ and action $a_{t}$. Constraints
on the state representation can be added by constraining the forward model,
for instance, by enforcing the system to follow linear dynamics (Watter et
al., 2015).

Another approach is to learn an inverse model (Shelhamer et al., 2017; Pathak
et al., 2017), which predicts the taken action $a_{t}$ given two successive
states $s_{t}$ and $s_{t+1}$. This enforces the states to encode information
about the dynamics, in order to recover the action needed for such transition.

####  2.2.4 Combining Approaches

Auto-encoders tend to reconstruct everything (that is salient enough in the
observation), including static objects and irrelevant features for the task
(distractors), whereas forward and inverse models focus on the dynamics,
usually encoding the position of the controlled robot in the representation,
but not a goal that is not controlled by the actions.

However, these approaches are not mutually exclusive and can be combined to
create improved state representations. For instance, (Pathak et al., 2017)
combine both inverse and forward models and (Zhang et al., 2018) additionally
integrate an auto-encoder. In the same vein, Ha and Schmidhuber (2018) use a
VAE with a recurrent forward model to learn a state representation.

##  3 Datasets and Environments

In this section, we describe a set of environments with incremental
difficulty, designed to assess SRL algorithms for robotic control. They all
follow the interface defined by OpenAI Gym (Brockman et al., 2016), which
makes integration with RL algorithms easy.

Mobile Navigation | Robotic Arm | Real Robot
---|---|---
|  |
Figure 1: Environments and datasets for state representation learning.

###  3.1 Environments Details

The settings we propose (Fig. 1) are variations of two environments: a 2D
environment with a mobile robot and a 3D environment with a robotic arm. In
all settings, there is a controlled robot and one or more targets (that can be
static, randomly initialized or moving). Each environment can either have a
continuous or discrete action space, and the reward can be sparse or shaped,
allowing us to cover many different situations.

Static & random target mobile navigation: This setting simulates a navigation
task using a small car resembling the task of (Jonschkowski and Brock, 2014),
with either a cylinder or a horizontal band on the ground as a goal, which can
be fixed or moving from episode to episode. The car can move in four
directions (forward, backward, left, right) and will get a +1 reward when
reaching the target, -1 when hitting walls, and 0 otherwise.

Static & random target robotic arm: This setting simulates a robotic arm
(Kuka), fixed on a table, with the task of pushing a button that may move or
not in between episodes. The arm can be controlled either in the $x$, $y$ and
$z$ position using inverse kinematics, or directly controlling the joints. The
robot will get +1 reward when the arm pushes the button, -1 when it hits the
table (this will end the episode), and 0 otherwise. A variant adds
distractors, i.e., moving objects on the table irrelevant for solving the
task.

Simulated & real robotic arm: We used a real Baxter robot arm (Gazebo in
simulation) to perform the same button pushing task as the previous task, with
the same actions and the same rewards. The goal is to test the different
methods in a real world setup.

A ground truth state is defined in each scenario: the absolute robot position
in static scenarios and the relative position (w.r.t. the target) in moving
goal scenarios. Note that apart from providing all described environments, we
also provide the corresponding datasets used in our evaluations: images are
224x224 pixels, navigation datasets use 4 discrete actions (right, left,
forward, backward); robot arms use one more (down) action.

In section A.2, we provide baselines results (Ground Truth, Auto-Encoder, Raw
Pixels) for each environment.

###  3.2 Motivation of the Goal-Based Robotics Tasks

The environments proposed have several characteristics that make them suitable
for research and benchmarking.

Designed for Robotics: The proposed environments cover basic goal-based
robotics tasks: navigation for a mobile robot and reaching a desired position
for a robotic arm.

Designed for State Representation Learning: The simplicity of the environments
makes the extracted features easier to interpret (correlation can be computed
between learned states and position of relevant objects). It is also clear
what a good state representation should encode because of the small number of
important elements: there is only the controllable robot and the target.

Designed for Research: The environments have incremental difficulty, the
minimal number of variables for describing each environment (minimal state
dimension for solving the task with RL) is increasing from 2 (mobile robot
with static target) to 6 (robotic arm with random target). Having simple
environments of gradual difficulty is really important when developing new
methods. The proposed environments are also easily customizable so that they
cover all possibilities: reward can be sparse/dense, actions can be
continuous/discrete. Finally, our benchmark is completely free and fast (it
runs at 250 FPS on a 8-core machine with one GPU).

##  4 Evaluation of Learned State Representations

The most practical evaluation of SRL is assessing if the learned states can be
used for solving the task in RL. However, algorithm development can benefit
from other metrics and visually assessing the validity of the representation
being learned, for faster iteration and interpretation of the state embedding
space. We provide tools and metrics for this.

###  4.1 Qualitative Evaluation

Real-time SRL | Interactive scatter | Latent visualization
---|---|---
|  |
Figure 2: Visual tools for analysing SRL; Left: Live trajectory of the robot
in the state space. Centre: 3D scatter plot of a state space; clicking on any
point displays the corresponding observation. Right: reconstruction of the
point in the state space defined by the sliders. See complementary material
for videos.

Qualitative evaluation in our case is the perceived utility of the state
representation using visualization tools. The perceived utility depends on the
task at hand. For example, the state representation of the robotic arm dataset
is expected to have a continuous and correlated change with respect to the arm
tip position. Three tools are proposed (Fig. 2):

Real-time SRL: This tool is used in conjunction with a graphical interface of
the simulated environment. It shows the correspondence between observation and
state by plotting current position of the observation in the state
representation.

Interactive scatter: A clickable plot of the state representation, with the
reward defining the colour for each point. Here, we expect to visualize
structure in the state representation.

Latent visualization: It allows to navigate in the latent space by projecting
the state to the observation space. This is achieved either by reconstructing
the output (for AE and VAE), or by using a nearest neighbour approach for the
models lacking a reconstruction.

To deal with state dimensions larger than three, PCA is used to visually
assess the learned representations in a qualitative manner.

These tools allow to have better insights of the learned representation over
different settings of state representation learning. This is especially useful
in situations where the state dimension is greater than 3. This way, we are
able to validate different configurations before running more exhaustive and
time consuming methods.

###  4.2 Metrics

####  4.2.1 KNN-MSE

We use an assessment of the representation’s quality based on a Nearest-
Neighbours approach (as in (Sermanet et al., 2017)). While the nearest
neighbour coherence can be assessed visually, KNN-MSE (Lesort et al., 2017)
derives a quantitative metric from this information.

For a given observation $o$, we find the nearest neighbours of its associated
state $s$ in the learned state space, and project them in the ground truth
state space. Then we compute the average distance to its neighbours in the
latter space:

$\textrm{KNN-MSE}(s)=\frac{1}{k}\sum_{s^{\prime}\in
KNN(s,k)}||\tilde{s}-\tilde{s}^{\prime}||^{2}$ (2)

where $\textrm{KNN}(s,k)$ returns the $k$ nearest neighbours of $s$ (chosen
with the Euclidean distance) in the learned state space $\mathcal{S}$,
$\tilde{s}$ is the ground truth state associated to $s$, and
$\tilde{s}^{\prime}$ is the one associated to $s^{\prime}$. A low KNN-MSE
means that a neighbour in the ground truth is still a neighbour in the learned
representation, and thus, local coherence is preserved.

####  4.2.2 Correlation

A Pearson $\rho$ correlation coefficient’s matrix is computed for each
dimension pair ($s$, $\tilde{s}$), where $\tilde{s}$ is the ground truth (GT)
state, $s$ the learned state, and $\mu_{s}$ and $\sigma_{s}$ are the mean and
standard deviation, respectively, of state $s$:

$\rho_{s,\tilde{s}}=\frac{\mathbb{E}[(s-\mu_{s})*(\tilde{s}-\mu_{\tilde{s}})]}{\sigma_{s}*\sigma_{\tilde{s}}}$
(3)

We can visualize the correlation matrix to quantitatively assess the ability
of a model to encode relevant information in the states learned. For instance,
the correlation matrix in Figure 3 shows degrees of correlation between the
mobile robot position and the learned states. The plot illustrates that for
each dimension $i$ of the predicted states $s$, there is a correlation close
to 1 (in absolute value) with at least one dimension $j$ of the agent’s real
position $\tilde{s}_{j}$. Therefore, this gives measurable evidences that the
model was able to encode the position of the mobile robot.

Figure 3: Correlation matrix for mobile robot navigation dataset (static
target), between each dimension $s_{i}$ of predicted states $s$ and the ground
truth $\tilde{s}_{j}$. We consider the ground truth to be the agent’s real
position. The states (dimension=2) are learned by combining a forward and an
inverse model.

This visualization tool is quite useful for low-dimensional spaces. However,
for state spaces with a high number of dimensions, looking at the correlation
matrix becomes impractical. Therefore, we introduce the following measure,
named GTC for Ground Truth Correlation, that allows to compare the models
ability to encode relevant information:

$GTC_{(i)}=\max\limits_{j}|\rho_{s,\tilde{s}}(i,j)|\in[0,1]$ (4)

with $i\in\llbracket 0,|\tilde{s}|\rrbracket$, $j\in\llbracket
0,|s|\rrbracket$, $\tilde{s}=[\tilde{s}_{1};...;\tilde{s}_{n}]$, and
$\tilde{s}_{k}$ being the $k^{th}$ dimension of the ground truth state vector.

For instance, in the Mobile Robot environment with random target, the ground
truth state is composed of the 2D robot position and 2D target position. That
is to say, the ground truth states have a dimension of 4: $|\tilde{s}|=4$.

The vector GTC gives for each component $i$ of the ground truth states
$\tilde{s}$, the maximum absolute correlation value between $\tilde{s}_{i}$
and any component of the predicted states $s$. Therefore, GTC measures the
similarity per component, between the learned states $s$ and the ground truth
states $\tilde{s}$.

We also introduce a metric, which is the mean of GTC, that allows to compare
learned states using one scalar value:

$GTC_{mean}=\mathbb{E}[GTC]$ (5)

###  4.3 Quantitative Evaluation With Reinforcement Learning

Comparing the performance of RL algorithms, using the learned state
representations, is the most relevant approach to evaluate the SRL methods. To
do so, our framework integrates 8 algorithms (A2C, ACKTR, ACER, DQN, DDPG,
PPO1, PPO2, TRPO) from Stable-Baselines (Hill et al., 2018) (a fork of OpenAI
baselines (Dhariwal et al., 2017)), Augmented Random Search (ARS) (Mania et
al., 2018), Covariance Matrix Adaptation Evolutionary Strategy (CMA-ES)
(Hansen et al., 2003) and Soft Actor Critic (SAC) (Haarnoja et al., 2018).

##  5 Experiments

We perform experiments on the proposed datasets using states learned with the
approaches described in Section 2.2 along with ground truth (GT). Here, we
report results obtained with PPO.

Ground Truth States | Learned States | RL Performance
---|---|---
|  |
Table 1: Ground truth states (left), states learned (Inverse and Forward)
(centre), and RL performance evaluation (PPO) (right) in the mobile robot
environment. Colour denotes the reward, red for positive, blue for negative
and grey for null reward (left and center). The full resolution plot of the RL
performance evaluation can be found in appendix Fig. 4 (right)

Table 1 illustrate the qualitative evaluation of a state space learned by
combining forward and inverse models on the mobile robot environment. It also
shows the performance of PPO based on the states learned by several
approaches.

Dataset | Mobile-robot (1) | Robotic-arm (3) | Robotic-arm-real (4)
---|---|---|---
Static/Random Target | Static | Random | Static | Random | Static
Ground Truth | 0.0099 | 0.0164 | 0.0025 | 0.0025 | 0.0105
AE | 0.0168 | 0.7213 | 0.00336 | 0.0027 | 0.0179
VAE | 0.0161 | 0.1295 | 0.0032 | 0.0027 | 0.0177
Robotic Priors | 0.0200 | 0.0900 | 0.0029 | 0.0027 | 0.0213
Forward | 0.5111 | 1.1557 | 0.1425 | 0.2564 | 0.0796
Inverse | 0.0191 | 0.7703 | 0.0182 | 0.2705 | 0.0521
Fwd+Inv | 0.0164 | 0.7467 | 0.0176 | 0.2705 | 0.0368
Table 2: KNN-MSE results for each SRL method. Ground Truth (robot absolute or
relative position with respect to target object), auto-encoders, forward
(fwd), inverse (inv) models and combinations.

Table 2 shows the KNN-MSE for the different SRL approaches on the implemented
environments.

The table 3 gives the $GTC_{mean}$ metrics for several approaches and the
associated RL performance using PPO. It shows that $GTC_{mean}$ is a good
indicator for the performance than can be obtained in RL: a good
disentanglement with a good correlation with the ground truth state (i.e. a
high $GTC_{mean}$) lead to a higher mean reward in RL.

During the development of S-RL Toolbox, we learned some useful insights on
SRL. We observed that auto-encoders do not reconstruct objects if they are too
small. This is an issue if the object is relevant for the task. An inverse
model is usually sufficient to learn a coherent state representation (in our
case, the position of the controllable object). Continuity in the state space
is important to perform well in RL. Using real-time SRL tools, we note
discontinuities in the state space learned by AE, even if reconstruction error
is low, and that hinders RL.

GroundTruthCorrelation | $x_{robot}$ | $y_{robot}$ | $x_{target}$ | $y_{target}$ | Mean | RL
---|---|---|---|---|---|---
Robotic Priors | 0.2 | 0.2 | 0.41 | 0.66 | 0.37 | 5.4 $\pm$ 3.1
Random | 0.68 | 0.65 | 0.34 | 0.31 | 0.50 | 163.4 $\pm$ 10.0
Supervised | 0.69 | 0.73 | 0.70 | 0.72 | 0.71 | 213.3 $\pm$ 6.0
Auto-Encoder | 0.52 | 0.51 | 0.24 | 0.23 | 0.38 | 138.5 $\pm$ 12.3
GT | 1 | 1 | 1 | 1 | 1 | 229.7 +- 2.7
Table 3: GTC and $GTC_{mean}$, and mean reward performance in RL (using PPO)
per episode after 2 millions steps, with standard error for each SRL method in
mobile robot with random target environment.

##  6 Related Work

In the RL literature, several classic benchmarks are used to compare
algorithms. For discrete actions, Atari Games from OpenAI Gym suite (Brockman
et al., 2016) are often adopted, whereas for continuous actions, MuJoCo
locomotion tasks (Todorov et al., 2012) is favoured. OpenAI recently open-
sourced robotics goal-based tasks (Plappert et al., 2018). These environments
are close to what we propose; however, they use a non-free physics engine,
observations are not images and customization is not easy. Also, our tools
were designed with SRL in mind and offer a gradual difficulty.

Regarding the SRL literature, very diverse environments are used, without
common metrics or visualizations. The SRL methods are usually only compared
with learning from raw observations. Jonschkowski and Brock (2015) assess the
quality of learned representations in a slot car scenario (similar to Lange et
al. (2012)) and a mobile navigation task (following Sprague (2009); Boots et
al. (2011)). More recently, Ha and Schmidhuber (2018) present their results on
Vizdoom and RaceCar environments. They show interesting insights of what was
learned by navigating in the latent space and projecting states back to the
pixel space (similar to Fig. 2). Zhang et al. (2018) use a subset of MuJoCo
tasks, with joints as input, and a binary maze environment. They perform an
insightful ablation study of their SRL model.

Our contribution is two-fold: we provide a framework integrated with RL
algorithms, environments and tools for SRL, along with implementations of the
main SRL methods.

##  7 Discussion and Conclusions

This paper presents a set of environments on which to perform SRL benchmarks
of incremental difficulty to solve tasks in RL, specifically, in robotics
control. Our proposed toolbox facilitates fast iterations, interpretability
and reproducibility, with a set of qualitative and quantitative metrics and
interactive visualization tools. We believe such a framework is needed to have
fair comparisons, focused on robotics control, among SRL methods.

Acknowledgments

This work is supported by the DREAM project333http://www.robotsthatdream.eu
through the European Union Horizon 2020 FET research and innovation program
under grant agreement No 640891.

##  A Implementation Details

###  A.1 Datasets Details

In this section, we provide the parameters used to generate datasets, for the
results presented in Table 2. The simulated datasets were created with a
random policy, using PyBullet (Coumans et al., 2018) (generating up to 30k
samples per minute on a 8-core machine 444CPU: Intel Core i7-7700K GPU:Nvidia
GeForce GTX 1080 Ti). The real Baxter dataset was recorded using ROS.

We used 10k samples of each dataset to learn a state representation. Reward is
sparse (see Section 3 for details) and actions are discrete (encoded as
integers) for all datasets 555Datasets can also be generated with continuous
actions and dense (shaped) reward.

Each target object can be fixed or randomly positioned between episodes; what
we call "absolute" or "relative" dataset, refers to the fact of learning a
relative or absolute position of the robot with respect to the target object.
The images dimension in every dataset is 224x224 pixels.

Table 4: Dataset details. Distractors can both be static or moving Dataset | Reward | Actions | Distractors
---|---|---|---
Mobile-robot | Sparse | 4, Discrete ($\Delta$(X,Y) pos.) | No
Robotic-arm | Sparse | 5, Discrete ($\Delta$(X,Y,Z) pos.) | No
Robotic-arm-real | Sparse | 5, Discrete ($\Delta$(X,Y,Z) pos.) | Yes

###  A.2 Baselines Results

In this section, we provide baselines results for each environment.

Budget (in timesteps) | 1 Million | 2 Million
---|---|---
Ground Truth | 198.0 $\pm$ 16.1 | 211.6 $\pm$ 14.0
Raw Pixels | 177.9 $\pm$ 15.6 | 215.7 $\pm$ 9.6
Auto-Encoder | 159.8 $\pm$ 16.1 | 188.8 $\pm$ 13.5
Table 5: Mean reward performance in RL (using PPO) per episode (average on 100 episodes) for different budgets, with standard error in Navigation 1D target environment. Budget (in timesteps) | 1 Million | 2 Million | 3 Million | 5 Million
---|---|---|---|---
Ground Truth | 227.8 $\pm$ 2.8 | 229.7 $\pm$ 2.7 | 231.5 $\pm$ 1.9 | 234.4 $\pm$ 1.3
Raw Pixels | 136.3 $\pm$ 11.5 | 188.2 $\pm$ 9.4 | 214.0 $\pm$ 5.9 | 231.5 $\pm$ 3.1
Auto-Encoder | 97.0 $\pm$ 12.3 | 138.5 $\pm$ 12.3 | 167.7 $\pm$ 11.1 | 192.6 $\pm$ 8.9
Table 6: Mean reward performance in RL (using PPO) per episode (average on 100 episodes) for different budgets, with standard error in Navigation 2D random target environment. Budget (in timesteps) | 1 Million | 2 Million | 3 Million | 5 Million
---|---|---|---|---
Ground Truth | 4.1 $\pm$ 0.5 | 4.1 $\pm$ 0.6 | 4.1 $\pm$ 0.6 | 4.2 $\pm$ 0.5
Raw Pixels | 0.6 $\pm$ 0.3 | 0.8 $\pm$ 0.3 | 1.2 $\pm$ 0.3 | 2.6 $\pm$ 0.3
Auto-Encoder | 0.92 $\pm$ 0.3 | 1.6 $\pm$ 0.3 | 2.2 $\pm$ 0.3 | 3.4 $\pm$ 0.3
Table 7: Mean reward performance in RL (using PPO) per episode (average on 100 episodes) for different budgets, with standard error in robotic arm with random target environment. Budget (in timesteps) | 1 Million | 2 Million | 3 Million | 5 Million
---|---|---|---|---
Ground Truth | 4.3 $\pm$ 0.3 | 4.4 $\pm$ 0.2 | 4.4 $\pm$ 0.2 | 4.6 $\pm$ 0.2
Raw Pixels | 0.8 $\pm$ 0.3 | 1.0 $\pm$ 0.3 | 1.2 $\pm$ 0.3 | 2.0 $\pm$ 0.3
Auto-Encoder | 1.17 $\pm$ 0.3 | 1.5 $\pm$ 0.3 | 1.9 $\pm$ 0.4 | 3.0 $\pm$ 0.4
Table 8: Mean reward performance in RL (using PPO) per episode (average on 100
episodes) for different budgets, with standard error in robotic arm with
moving target environment.

###  A.3 Reinforcement Learning Experiments

####  A.3.1 Reinforcement Learning

We use the implementations from Stable-Baselines (Hill et al., 2018) for the
RL experiments (except for CMA-ES, ARS and SAC which we implemented
ourselves). Where we used the default hyperparameters set by OpenAI Baselines.
With 10 random seeds in order to have quantitative results.

As for the network of the policies, the same architecture is used in all
different methods. For the SRL and ground truth approaches, it is a 2-layers
MLP, whereas for learning from raw pixels, it is the CNN from (Mnih et al.,
2015) implemented in OpenAI baselines.

All CNN policies normalize the input image by dividing it by 255. Observations
are not stacked. When learning from SRL, the states are normalized using a
running mean/std average. Reinforcement learning metrics reported are the
average returned rewards over 5 policies, independently trained using the same
RL algorithm with a different seed.

####  A.3.2 State Learning Representation

For the SRL methods, we used an altered form of ResNet (Table 9), which is
more compact for our needs. However, for forward and inverse models, we used a
linear model in addition to the CNN for learning a state representation. We
are using the Adam (Kingma and Ba, 2014) optimizer in all our models, with
learning rates that vary from $10^{-2}$ to $10^{-3}$. The batch size of
forward and inverse is 128, 256 for robotic priors, and 32 for the other
models.

Table 9: CNN architecture for SRL methods. The first layers are inspired by ResNet architecture (He et al., 2015).  Layer | Architecture
---|---
1 | Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) + BN + ReLU
2 | MaxPool2d(kernel_size=3, stride=2, padding=1)
3 | Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False) + BN + ReLU
4 | MaxPool2d(kernel_size=3, stride=2)
5 | Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False) + BN + ReLU
6 | MaxPool2d(kernel_size=3, stride=2)
7 | Linear(6 * 6 * 64, state_dim)

####  A.3.3 Results

Fig. 4 shows the average reward during learning of the PPO algorithm.

Figure 4: Performance (mean and standard error for 10 runs) for PPO algorithm
for different state representations learned in mobile-robot-navigation (random
target) environment Figure 5: Performance (mean and standard error) on RL
algorithms using ground truth states with mobile robot (random target)
environment

We observed during the development of the _S-RL Toolbox_ , that PPO was one of
the best RL algorithm for SRL benchmarking. It allows us to obtain good
performance and was consistent without needing to change any hyperparameters.
Figure 5 illustrates this, in addition to the variability of the RL
algorithms.

##  B Supplementary material

The visualization and interactive state space exploration tools are
demonstrated in the following videos:

  * •

S-RL Toolbox Showcase: https://youtu.be/qNsHMkIsqJc

  * •

S-RL Toolbox Environments: https://tinyurl.com/y973vhfy

  * •

Kuka robot arm: RL running PPO (SRL trained with VAE):
https://tinyurl.com/yarpbs2c

  * •

Kuka robot arm environment: State Representation Learning Benchmark running
PPO on ground truth states: https://tinyurl.com/yd549l7o

## References

  * Baldi (2012) Pierre Baldi.  Autoencoders, unsupervised learning, and deep architectures.  In _Proceedings of ICML Workshop on Unsupervised and Transfer Learning_ , pages 37–49, 2012.
  * Boots et al. (2011) Byron Boots, Sajid M Siddiqi, and Geoffrey J Gordon.  Closing the learning-planning loop with predictive state representations.  _The International Journal of Robotics Research_ , 30(7):954–966, 2011.
  * Brockman et al. (2016) Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba.  Openai gym.  _CoRR_ , abs/1606.01540, 2016.  URL http://arxiv.org/abs/1606.01540.
  * Coumans et al. (2018) E. Coumans, Y. Bai, and J. Hsu.  Pybullet physics engine.  http://pybullet.org/, 2018.
  * Dhariwal et al. (2017) Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, and Yuhuai Wu.  Openai baselines.  https://github.com/openai/baselines, 2017.
  * Ha and Schmidhuber (2018) D. Ha and J. Schmidhuber.  World Models.  _ArXiv e-prints_ , March 2018.
  * Haarnoja et al. (2018) Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  _arXiv preprint arXiv:1801.01290_ , 2018.
  * Hansen et al. (2003) Nikolaus Hansen, Sibylle D Müller, and Petros Koumoutsakos.  Reducing the time complexity of the derandomized evolution strategy with covariance matrix adaptation (cma-es).  _Evolutionary computation_ , 11(1):1–18, 2003\.
  * He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  _CoRR_ , abs/1512.03385, 2015.  URL http://arxiv.org/abs/1512.03385.
  * Hill et al. (2018) Ashley Hill, Antonin Raffin, Rene Traore, Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, and Yuhuai Wu.  Stable baselines.  https://github.com/hill-a/stable-baselines, 2018.
  * Jonschkowski (2018) Rico Jonschkowski.  Learning robotic perception through prior knowledge.  2018\.
  * Jonschkowski and Brock (2014) Rico Jonschkowski and Oliver Brock.  State representation learning in robotics: Using prior knowledge about physical interaction.  In _Proceedings of Robotics: Science and Systems_ , July 2014.
  * Jonschkowski and Brock (2015) Rico Jonschkowski and Oliver Brock.  Learning state representations with robotic priors.  _Autonomous Robots_ , 39(3):407–428, 2015.  ISSN 0929-5593.
  * Jonschkowski et al. (2017) Rico Jonschkowski, Roland Hafner, Jonathan Scholz, and Martin A. Riedmiller.  PVEs: Position-Velocity Encoders for Unsupervised Learning of Structured State Representations.  _CoRR_ , abs/1705.09805, 2017.  URL http://arxiv.org/abs/1705.09805.
  * Kingma and Welling (2013) D. P Kingma and M. Welling.  Auto-Encoding Variational Bayes.  _ArXiv e-prints_ , December 2013.
  * Kingma and Ba (2014) Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _CoRR_ , abs/1412.6980, 2014.  URL http://arxiv.org/abs/1412.6980.
  * Lange et al. (2012) Sascha Lange, Martin Riedmiller, and Arne Voigtlander.  Autonomous reinforcement learning on raw visual input data in a real world application.  In _Neural Networks (IJCNN), The 2012 International Joint Conference on_ , pages 1–8. IEEE, 2012.
  * Lesort et al. (2017) Timothée Lesort, Mathieu Seurin, Xinrui Li, Natalia Díaz-Rodríguez, and David Filliat.  Unsupervised state representation learning with robotic priors: a robustness benchmark.  _CoRR_ , abs/1709.05185, 2017.  URL http://arxiv.org/abs/1709.05185.
  * Lesort et al. (2018) Timothée Lesort, Natalia Díaz-Rodríguez, Jean-François Goudou, and David Filliat.  State representation learning for control: An overview.  _Neural Networks_ , 2018.  ISSN 0893-6080.  doi: https://doi.org/10.1016/j.neunet.2018.07.006.  URL http://www.sciencedirect.com/science/article/pii/S0893608018302053.
  * Mania et al. (2018) Horia Mania, Aurelia Guy, and Benjamin Recht.  Simple random search provides a competitive approach to reinforcement learning.  _arXiv preprint arXiv:1803.07055_ , 2018.
  * Mnih et al. (2015) Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al.  Human-level control through deep reinforcement learning.  _Nature_ , 518(7540):529–533, 2015.
  * Pathak et al. (2017) Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell.  Curiosity-driven exploration by self-supervised prediction.  In _ICML_ , 2017.
  * Plappert et al. (2018) Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, Vikash Kumar, and Wojciech Zaremba.  Multi-goal reinforcement learning: Challenging robotics environments and request for research.  _CoRR_ , abs/1802.09464, 2018.  URL http://arxiv.org/abs/1802.09464.
  * Sermanet et al. (2017) Pierre Sermanet, Corey Lynch, Jasmine Hsu, and Sergey Levine.  Time-contrastive networks: Self-supervised learning from multi-view observation.  _CoRR_ , abs/1704.06888, 2017.  URL http://arxiv.org/abs/1704.06888.
  * Shelhamer et al. (2017) Evan Shelhamer, Parsa Mahmoudieh, Max Argus, and Trevor Darrell.  Loss is its own reward: Self-supervision for reinforcement learning.  _arXiv preprint arXiv:1612.07307_ , 2017.
  * Sprague (2009) Nathan Sprague.  Predictive projections.  In _IJCAI_ , pages 1223–1229, 2009.
  * Todorov et al. (2012) Emanuel Todorov, Tom Erez, and Yuval Tassa.  Mujoco: A physics engine for model-based control.  In _Intelligent Robots and Systems (IROS), 2012 IEEE/RSJ International Conference on_ , pages 5026–5033. IEEE, 2012.
  * Watter et al. (2015) Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller.  Embed to control: A locally linear latent dynamics model for control from raw images.  In C. Cortes, N. D. Lawrence, D. D. Lee, M. Sugiyama, and R. Garnett, editors, _Advances in Neural Information Processing Systems 28_ , pages 2746–2754. Curran Associates, Inc., 2015.
  * Zhang et al. (2018) Amy Zhang, Harsh Satija, and Joelle Pineau.  Decoupling dynamics and reward for transfer learning.  _arXiv preprint arXiv:1804.10689_ , 2018.