Differentiable Model Predictive Excitation (DMPE):

This repository implements an algorithm for the excitation of systems with unknown (usually non-linear) systems. The inner workings and lines of thought are outlined within the corresponding publication. If you found this repository useful for your research, please cite the current preprint version as:

@Article{Vater2024,
  author      = {Vater, Hendrik and Wallscheid, Oliver},
  title       = {Differentiable Model Predictive Excitation: Generating Optimal Data Sets for Learning of Dynamical System Models},
  journal     = {TechRxiv preprint},
  year        = {2024},
  doi         = {10.36227/techrxiv.172840381.16440835/v1},
}

Installation:

Simplest way, using Python >= 3.11:

pip install dmpe

• Intended for a Linux system using an NVIDIA GPU where CUDA is set up
• Theoretically, it can be used without a GPU and also on Windows, but performance will likely be suboptimal and the results are not exactly reproducible with the GPU/Linux results.
• Depends on exciting_environments
• As of now, the requirements/dependencies are strict. It is likely that other versions work as well, but the given setup has been used extensively. (The requirements will likely be extended in the future.)
• As this repository is actively being worked on, it is possible that a more recent version is accessible in the DMPE GitHub repository.

Alternative installation:

Download the current state of the exciting_environments repository, e.g.:

git clone [email protected]:ExcitingSystems/exciting-environments.git

and install it in your python environment by moving to the downloaded folder and running pip install .. Then, download the DMPE source code, e.g.:

git clone [email protected]:ExcitingSystems/DMPE.git

Afterwards, install it from within the repository folder via pip install -e . for an editable version or with pip install . if you do not plan to make changes to the code.

Structure:

The repository is structured as follows:

• dmpe/ contains the source code for the DMPE algorithm and also for the GOATS algorithms from the related work.
• eval/ contains the code used in the experiments in the corresponding publication [Vater2024].
• dev/ contains jupyter notebooks that are intended for development on the repository.
• examples/ contains some examples to get started
• fig/ contains example images (e.g., for the README)

Basic Usage:

To apply the algorithms onto a system, the systems structure must comply to a specific API (Naturally, this can be adapted in the future. Please open an issue or write an e-mail to [email protected], if you are interested in discussing this). Example environments following this API can be found in the exciting_environments repository.

Using the algorithm for such an environment is as simple as:

import jax.numpy as jnp
import diffrax

import exciting_environments as excenvs
from dmpe.models.models import NeuralEulerODEPendulum
from dmpe.algorithms import excite_with_dmpe
from dmpe.algorithms.algorithm_utils import default_dmpe_parameterization


env = excenvs.make(
    env_id="Pendulum-v0",
    batch_size=1,
    action_constraints={"torque": 5},
    static_params={"g": 9.81, "l": 1, "m": 1},
    solver=diffrax.Tsit5(),
    tau=2e-2,
)

def featurize_theta(obs):
    """Transform angle information with sin() and cos()."""
    feat_obs = jnp.stack([jnp.sin(obs[..., 0] * jnp.pi), jnp.cos(obs[..., 0] * jnp.pi), obs[..., 1]], axis=-1)
    return feat_obs

# get default parameterization
exp_params, proposed_actions, loader_key, expl_key = default_dmpe_parameterization(
    env, seed=0, featurize=featurize_theta, model_class=NeuralEulerODEPendulum
)
exp_params["n_time_steps"] = 1000  # reduce to N=1000 steps

# run excitation
observations, actions, model, density_estimate, losses, proposed_actions = excite_with_dmpe(
    env,
    exp_params,
    proposed_actions,
    loader_key,
    expl_key,
)

# visualize
from dmpe.evaluation.plotting_utils import plot_sequence
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['text.usetex'] = True
mpl.rcParams.update({'font.size': 10})
mpl.rcParams['text.latex.preamble']=r"\usepackage{bm}\usepackage{amsmath}"

fig = plot_sequence(observations, actions, env.tau, env.obs_description, env.action_description)
plt.show()