Koopman-Assisted Reinforcement Learning

About

KoopmanRL is a reinforcement learning (RL) package designed around the two Koopman-Assisted RL (KARL) algorithms, Soft Koopman Value Iteration (SKVI), and Soft Actor Koopman-Critic (SAKC). It provides the utilities to build upon parts of its algorithms by either using only the Koopman tensor itself, or only components of the two KARL algorithms. In addition it provides utilities for automatic hyperparameter tuning of KARL algorithms, as well as 4 environments rooted in the control literature.

Getting Started

Having uv installed, one can easily get the project's environment up-and-running by cloning the repository and syncing the local project environment with the lock file:

git clone https://github.com/dynamicslab/KoopmanRL.git
cd KoopmanRL && uv sync

At which point we have the local project environment we need to run the computations. Alternatively, one can create a dedicated virtual environment with uv, i.e.:

uv venv --python 3.10 && source .venv/bin/activate

After which one can install the KoopmanRL package into the virtual environment

uv pip install .

At which point one can run the first experiment, taking the (fast) Linear Quadratic Regulator as an illustrative example here:

uv run -m koopmanrl.linear_quadratic_regulator   # after `uv sync`
python -m koopmanrl.linear_quadratic_regulator   # in the activated virtual environment

The commands in this README use uv run, which first syncs the project environment with the lock file. In a virtual environment set up with uv pip install ., run them with python in place of uv run instead (e.g. python -m koopmanrl.sakc_optuna_opt ...), so that the installed environment is used as it is.

Which will default to the linear system with the default hyperparameters. For all algorithms, and the hyperparameter optimization routines query the available hyperparameters with the --help identifier after the routine at which point all options will be listed in the terminal.

uv run -m koopmanrl.linear_quadratic_regulator --help

Running Individual Experiments

KoopmanRL provides a number of algorithm implementations, which can all be run as quasi-scripts with their own (typed) argument parsers. The main files here are:

Control Algorithms:

• koopmanrl.linear_quadratic_regulator
• koopmanrl.sac_continuous_action
• koopmanrl.value_based_sac_continuous_action
• koopmanrl.soft_koopman_value_iteration
• koopmanrl.soft_actor_koopman_critic

Hyperparameter Optimization:

• koopmanrl.skvi_optuna_opt
• koopmanrl.sakc_optuna_opt

Koopman Tensor:

• koopmanrl.koopman_tensor.generate_tensor, which identifies a Koopman tensor of one environment and stores it; --regressor {ols,ridge,sindy,rrr} selects the regression

Each of the algorithms, be it a control algorithm, or one of the two hyperparameter optimization routines can be applied to one of the four environments. The implemented environments are:

1. LinearSystem-v0
2. FluidFlow-v0
3. Lorenz-v0
4. DoubleWell-v0

The algorithms can be applied to any of the environments with the --env_id flag. Utilizing the example of the Linear Quadratic Regulator again:

uv run -m koopmanrl.linear_quadratic_regulator --env_id FluidFlow-v0

Note

Note that the hyperparameter optimizations need to be started with the python executable of the environments. E.g. with a uv environment:

uv run python -m koopmanrl.sakc_optuna_opt --env_id FluidFlow-v0 \
   --storage_dir . --output_file my_sakc_fluid_flow --cpu_cores_per_trial 4

or when working with an activated virtual environment

python -m koopmanrl.sakc_optuna_opt --env_id FluidFlow-v0 \
   --storage_dir . --output_file my_sakc_fluid_flow --cpu_cores_per_trial 4

The best configuration is written to <storage_dir>/<output_file>.json. --storage_dir is required and has to exist. Keep --cpu_cores_per_trial (default 28) at or below the number of cores of your machine, since Ray cannot schedule a trial that asks for more.

Using Optimized Hyperparameter Configurations

The configurations/ directory contains the best-found hyperparameter configurations for SAKC and SKVI across each supported environment, produced by the Optuna optimization routines. The naming scheme is <algo>_<env_slug>_hparams.json, e.g.:

configurations/sakc_fluid_flow_hparams.json
configurations/skvi_lorenz_hparams.json
configurations/sakc_double_well_hparams.json

Next to them, episodic_returns.json, ablations.json and tsne.json list the runs of the three reproduction pipelines described below, which both the launchers and the Snakemake workflows read.

Both koopmanrl.soft_actor_koopman_critic and koopmanrl.soft_koopman_value_iteration accept a --config_file flag that loads hyperparameters directly from one of these JSON files:

uv run python -m koopmanrl.soft_actor_koopman_critic \
    --config_file configurations/sakc_fluid_flow_hparams.json

uv run python -m koopmanrl.soft_koopman_value_iteration \
    --config_file configurations/skvi_lorenz_hparams.json

Any flag passed explicitly on the command line takes precedence over the value in the config file, so individual hyperparameters can be overridden without editing the JSON:

uv run python -m koopmanrl.soft_actor_koopman_critic \
    --config_file configurations/sakc_fluid_flow_hparams.json \
    --seed 42 \
    --total_timesteps 100000

Reproducing the Results of the Paper

The raw data of the figures and the tables the figures are drawn from can be made again with the launchers in koopmanrl_utils/. Each guide below describes the path from the runs to the pgfplots-ready tables:

Result	Launcher	Guide
Episodic returns of LQR, SAC (Q), SAC (V), SKVI and SAKC	koopmanrl_utils.run_optimized_experiments	koopmanrl_utils/EPISODIC_RETURNS.md
Ablations of SKVI and SAKC	koopmanrl_utils.run_ablations	koopmanrl_utils/ABLATIONS.md
t-SNE embedding of the Koopman tensors	koopmanrl_utils.tsne_koopman_tensor	koopmanrl_utils/TSNE.md

The same three pipelines are available as Snakemake workflows, which make every run in a job of its own and only the results that are missing:

uv sync
uvx --python 3.12 snakemake -n episodic_returns          # list the jobs without running them
uvx --python 3.12 snakemake --cores 8 episodic_returns   # or: ablations, tsne

The episodic returns take 495 runs and the ablations 1,440 runs, which is several days of computing; workflow/README.md describes how to run a part of them and how to run them on a cluster.

License

KoopmanRL is released under the MIT License (LICENSE.md). The observables in koopmanrl/koopman_observables.py and koopmanrl/koopman_tensor/observables/ are adapted from d3s and remain under the LGPL-3.0, and the SAC-based scripts are adapted from CleanRL (MIT); see LICENSES/README.md. The package as a whole is distributed under MIT AND LGPL-3.0-only.

Citation

If you use KoopmanRL, please cite the paper:

@article{rozwood2024koopman,
  title   = {Koopman-Assisted Reinforcement Learning},
  author  = {Rozwood, Preston and Mehrez, Edward and Paehler, Ludger and Sun, Wen and Brunton, Steven L.},
  journal = {arXiv preprint arXiv:2403.02290},
  year    = {2024}
}

The software itself can be cited with the metadata in CITATION.cff, which GitHub also offers under "Cite this repository".