├── .gitignore ├── .gitlab-ci.yml ├── .gitmodules ├── .pre-commit-config.yaml ├── LICENSE ├── README.md ├── examples ├── inverted_pendulum │ ├── inverted_pendulum_model_acados.py │ ├── inverted_pendulum_residual_learning_zoro.ipynb │ ├── inverted_pendulum_residual_learning_zoro.py │ └── utils.py └── l4casadi_vs_l4acados │ ├── README.md │ ├── acados.py │ ├── analyze_results.ipynb │ ├── analyze_results.py │ ├── build_acados_num_threads.sh │ ├── data │ └── .gitkeep │ ├── l4casadi_vs_l4acados.ipynb │ ├── l4casadi_vs_l4acados.py │ ├── run_all_experiments.py │ └── run_single_experiment.py ├── external └── gpytorch_utils │ ├── gp_hyperparam_training.py │ └── gp_utils.py ├── pyproject.toml ├── setup.py ├── src └── l4acados │ ├── __init__.py │ ├── controllers │ ├── __init__.py │ ├── residual_learning_mpc.py │ ├── zero_order_gpmpc.py │ └── zoro_acados_utils.py │ └── models │ ├── __init__.py │ ├── pytorch_models │ ├── __init__.py │ ├── gpytorch_models │ │ ├── __init__.py │ │ ├── gpytorch_data_processing_strategy.py │ │ ├── gpytorch_gp.py │ │ └── gpytorch_residual_model.py │ ├── pytorch_feature_selector.py │ ├── pytorch_residual_model.py │ └── pytorch_utils.py │ └── residual_model.py └── test ├── acados_python_examples └── getting_started │ ├── __init__.py │ ├── pendulum_model.py │ ├── test_minimal_example_closed_loop.py │ └── utils.py ├── gpytorch_model └── test_gpytorch_learning_model.py └── l4casadi_comparison ├── __init__.py ├── l4casadi_with_acados.py ├── nn_params.pt └── test_l4casadi_l4acados_comparison.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Notebooks in test folder 2 | test/*/*.ipynb 3 | test/*.ipynb 4 | 5 | # Environment 6 | .envrc 7 | 8 | # General 9 | .DS_Store 10 | 11 | # Test 12 | test/legacy_controllers/*.json 13 | 14 | # Examples 15 | examples/inverted_pendulum/*.json 16 | examples/l4casadi_vs_l4acados/data/ 17 | 18 | # PDF 19 | *.pdf 20 | 21 | # CasADi / acados 22 | c_generated_code/ 23 | jit_tmp*.c 24 | tmp_casadi_compiler_shell*.o 25 | *sim*.json 26 | *ocp*.json 27 | _l4c_generated/ 28 | 29 | # Byte-compiled / optimized / DLL files 30 | __pycache__/ 31 | *.py[cod] 32 | *$py.class 33 | 34 | # C extensions 35 | *.so 36 | 37 | # Distribution / packaging 38 | .Python 39 | build/ 40 | develop-eggs/ 41 | dist/ 42 | downloads/ 43 | eggs/ 44 | .eggs/ 45 | lib/ 46 | lib64/ 47 | parts/ 48 | sdist/ 49 | var/ 50 | wheels/ 51 | share/python-wheels/ 52 | *.egg-info/ 53 | .installed.cfg 54 | *.egg 55 | MANIFEST 56 | 57 | # PyInstaller 58 | # Usually these files are written by a python script from a template 59 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 60 | *.manifest 61 | *.spec 62 | 63 | # Installer logs 64 | pip-log.txt 65 | pip-delete-this-directory.txt 66 | 67 | # Unit test / coverage reports 68 | htmlcov/ 69 | .tox/ 70 | .nox/ 71 | .coverage 72 | .coverage.* 73 | .cache 74 | nosetests.xml 75 | coverage.xml 76 | *.cover 77 | *.py,cover 78 | .hypothesis/ 79 | .pytest_cache/ 80 | cover/ 81 | 82 | # Translations 83 | *.mo 84 | *.pot 85 | 86 | # Django stuff: 87 | *.log 88 | local_settings.py 89 | db.sqlite3 90 | db.sqlite3-journal 91 | 92 | # Flask stuff: 93 | instance/ 94 | .webassets-cache 95 | 96 | # Scrapy stuff: 97 | .scrapy 98 | 99 | # Sphinx documentation 100 | docs/_build/ 101 | 102 | # PyBuilder 103 | .pybuilder/ 104 | target/ 105 | 106 | # Jupyter Notebook 107 | .ipynb_checkpoints 108 | 109 | # IPython 110 | profile_default/ 111 | ipython_config.py 112 | 113 | # pyenv 114 | # For a library or package, you might want to ignore these files since the code is 115 | # intended to run in multiple environments; otherwise, check them in: 116 | .python-version 117 | 118 | # pipenv 119 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 120 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 121 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 122 | # install all needed dependencies. 123 | #Pipfile.lock 124 | 125 | # poetry 126 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 127 | # This is especially recommended for binary packages to ensure reproducibility, and is more 128 | # commonly ignored for libraries. 129 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 130 | #poetry.lock 131 | 132 | # pdm 133 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 134 | #pdm.lock 135 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 136 | # in version control. 137 | # https://pdm.fming.dev/#use-with-ide 138 | .pdm.toml 139 | 140 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 141 | __pypackages__/ 142 | 143 | # Celery stuff 144 | celerybeat-schedule 145 | celerybeat.pid 146 | 147 | # SageMath parsed files 148 | *.sage.py 149 | 150 | # Environments 151 | .env 152 | .venv 153 | env/ 154 | venv/ 155 | ENV/ 156 | env.bak/ 157 | venv.bak/ 158 | 159 | # Spyder project settings 160 | .spyderproject 161 | .spyproject 162 | 163 | # Rope project settings 164 | .ropeproject 165 | 166 | # mkdocs documentation 167 | /site 168 | 169 | # mypy 170 | .mypy_cache/ 171 | .dmypy.json 172 | dmypy.json 173 | 174 | # Pyre type checker 175 | .pyre/ 176 | 177 | # pytype static type analyzer 178 | .pytype/ 179 | 180 | # Cython debug symbols 181 | cython_debug/ 182 | 183 | # PyCharm 184 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 185 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 186 | # and can be added to the global gitignore or merged into this file. For a more nuclear 187 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 188 | #.idea/ 189 | -------------------------------------------------------------------------------- /.gitlab-ci.yml: -------------------------------------------------------------------------------- 1 | image: python:3.9 2 | stages: 3 | - format 4 | - build 5 | - python 6 | - test 7 | 8 | ### ============================================================= 9 | ### ============== JOB TEMPLATES ========================= 10 | ### ============================================================= 11 | 12 | .deep_thought_job: 13 | retry: 0 14 | tags: 15 | - deep-thought 16 | 17 | .python_parallel_matrix_job: 18 | extends: .deep_thought_job 19 | image: "python:$PYTHON_VERSION" 20 | parallel: 21 | matrix: 22 | - PYTHON_VERSION: ["3.8", "3.9", "3.10"] 23 | before_script: 24 | - python --version ; pip --version # For debugging 25 | - python -m venv $VENV_DIR 26 | - source "$VENV_DIR/bin/activate" || exit 1 27 | - pip list # For debugging 28 | # l4casadi 29 | - cd $CI_PROJECT_DIR && git submodule update --recursive --init external/l4casadi || exit 1 30 | - cd $CI_PROJECT_DIR/external/l4casadi 31 | - pip install torch>=2.0 --index-url https://download.pytorch.org/whl/cpu || exit 1 32 | - pip install -r requirements_build.txt || exit 1 33 | - pip install . --no-build-isolation || exit 1 34 | # l4acados 35 | - cd $CI_PROJECT_DIR 36 | - pip install -e ./external/acados/interfaces/acados_template || exit 1 37 | - pip install -e .[test] --no-build-isolation || exit 1 38 | - pip install -e .[gpytorch-exo] || exit 1 39 | variables: 40 | PIP_CACHE_DIR: "$CI_PROJECT_DIR/.cache/pip" 41 | VENV_DIR: "$CI_PROJECT_DIR/.venv_$PYTHON_VERSION" 42 | resource_group: deep-thought-l4acados-$PYTHON_VERSION 43 | 44 | .pytest_job: 45 | extends: .python_parallel_matrix_job 46 | stage: test 47 | script: 48 | - pytest --ignore external/ 49 | variables: 50 | ACADOS_SOURCE_DIR: "$CI_PROJECT_DIR/external/acados" 51 | LD_LIBRARY_PATH: "$ACADOS_SOURCE_DIR/lib" 52 | cache: 53 | key: pip-cache-$PYTHON_VERSION 54 | paths: 55 | - .cache/pip 56 | # - $VENV_DIR 57 | policy: pull 58 | 59 | 60 | ### ============================================================= 61 | ### ============== FORMATTING STAGE ====================== 62 | ### ============================================================= 63 | 64 | run formatters: # Check whether code adheres to the required formats 65 | extends: .deep_thought_job 66 | stage: format 67 | script: 68 | - pip install pre-commit && pre-commit install 69 | - pre-commit run -a 70 | 71 | ### ============================================================= 72 | ### ============== PREPARATION STAGE ===================== 73 | ### ============================================================= 74 | 75 | install acados: 76 | extends: .deep_thought_job 77 | stage: build 78 | image: python:3.8 79 | script: 80 | - apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y cmake wget 81 | - mkdir -p $CI_PROJECT_DIR/external/acados/build && cd $CI_PROJECT_DIR/external/acados/build && cmake -DACADOS_PYTHON=ON .. && make install -j4 || exit 1 82 | - cd $CI_PROJECT_DIR/external/acados/bin && wget https://github.com/acados/tera_renderer/releases/download/v0.0.34/t_renderer-v0.0.34-linux -O t_renderer && chmod +x t_renderer || exit 1 83 | retry: 0 84 | tags: 85 | - deep-thought 86 | artifacts: 87 | paths: 88 | - $CI_PROJECT_DIR/external/acados/ 89 | variables: 90 | GIT_SUBMODULE_STRATEGY: recursive 91 | GIT_SUBMODULE_PATHS: external/acados 92 | 93 | 94 | install python dependencies: 95 | extends: .python_parallel_matrix_job 96 | stage: python 97 | script: 98 | - python -c "import numpy, gpytorch, acados_template, l4casadi" || exit 1 99 | cache: 100 | key: pip-cache-$PYTHON_VERSION 101 | paths: 102 | - .cache/pip 103 | # - $VENV_DIR # TODO: fix caching of venv 104 | # artifacts: 105 | # paths: 106 | # - "$VENV_DIR/" 107 | 108 | 109 | ### ============================================================= 110 | ### ============== TESTING STAGE ========================= 111 | ### ============================================================= 112 | 113 | run tests: 114 | extends: .pytest_job 115 | script: 116 | - pytest --ignore external/ 117 | resource_group: l4acados_tests 118 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "external/acados"] 2 | path = external/acados 3 | url = https://github.com/acados/acados.git 4 | [submodule "external/l4casadi"] 5 | path = external/l4casadi 6 | url = git@github.com:lahramon/l4casadi.git 7 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | # See https://pre-commit.com for more information 2 | # See https://pre-commit.com/hooks.html for more hooks 3 | repos: 4 | - repo: https://github.com/mwouts/jupytext 5 | rev: v1.16.4 # CURRENT_TAG/COMMIT_HASH 6 | hooks: 7 | - id: jupytext 8 | # args: [--sync, --pipe, black] 9 | args: [--sync] 10 | additional_dependencies: 11 | - black==24.3.0 # Matches hook 12 | - repo: https://github.com/psf/black-pre-commit-mirror 13 | rev: 24.3.0 14 | hooks: 15 | - id: black 16 | - repo: https://github.com/pre-commit/pre-commit-hooks 17 | rev: v4.6.0 18 | hooks: 19 | - id: trailing-whitespace 20 | - id: end-of-file-fixer 21 | - id: check-yaml 22 | - id: mixed-line-ending 23 | args: [-f, lf] 24 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Amon Lahr, Joshua Näf, Andrea Carron, Melanie N. Zeilinger 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # L4acados: Learning-based models for `acados` 2 | > Paper: https://arxiv.org/abs/2411.19258 3 | 4 | Integrate learning-based Python models into acados for real-time model predictive control. 5 | 6 | ## Usage 7 | 8 | 1. Define your [`AcadosOcp`](https://docs.acados.org/python_interface/index.html#acados_template.acados_ocp.AcadosOcp), including the nominal dynamics model 9 | ```python 10 | from acados_template import AcadosOcp 11 | ocp = AcadosOcp() 12 | # ... 13 | ``` 14 | 15 | 2. Define the residual model using the `L4acados` `ResidualModel` (here as a [PyTorchResidualModel](https://github.com/IntelligentControlSystems/l4acados/blob/main/src/l4acados/models/pytorch_models/pytorch_residual_model.py) example): 16 | ```python 17 | import l4acados as l4a 18 | residual_model = l4a.PyTorchResidualModel(your_pytorch_model) 19 | ``` 20 | 21 | > [Other models](https://github.com/IntelligentControlSystems/l4acados/tree/main/src/l4acados/models) can be straightforwardly implemented using as a [`ResidualModel`](https://github.com/IntelligentControlSystems/l4acados/blob/main/src/l4acados/models/residual_model.py) instance; see [here](https://github.com/IntelligentControlSystems/l4acados?tab=readme-ov-file#install-l4acados-with-optional-dependencies-of-your-choice) for already available residual models. 22 | 23 | 3. Generate the `L4acados` solver object 24 | ```python 25 | l4acados_solver = l4a.ResidualLearningMPC( 26 | ocp=ocp, 27 | residual_model=residual_model, 28 | use_cython=True, # accelerate get/set operations by using the acados Cython interface 29 | ) 30 | ``` 31 | 32 | 4. Done! The `ResidualLearningMPC` object can be interfaced like the `AcadosOcpSolver`: 33 | - `l4acados_solver.set(...)` 34 | - `l4acados_solver.solve()` 35 | - `l4acados_solver.get(...)` 36 | > Not all solver interface functions are mapped by the `ResidualLearningMPC`. You can still access the underlying `AcadosOcpSolver` object through the `l4acados_solver.ocp_solver` property. Besides the dynamics model and parameter definition, the `l4acados_solver.ocp_solver` is equivalent to the `acados_solver.ocp_solver` generated from the original `ocp`. 37 | 38 | 39 | ## Installation 40 | 41 | ### Install `acados` 42 | 43 | 1. Clone 44 | ```bash 45 | git submodule update --recursive --init external/acados 46 | ``` 47 | 1. Build the submodule `acados` according to the [installation instructions](https://docs.acados.org/installation/index.html): 48 | ```bash 49 | mkdir -p external/acados/build 50 | cd external/acados/build 51 | cmake -DACADOS_PYTHON=ON .. # do not forget the ".." 52 | make install -j4 53 | ``` 54 | 55 | > You can also install `acados` separately; in this case, refer to the acados version of the submodule in case you encounter errors with a newer version of `acados`. 56 | 57 | 2. Install acados Python interface 58 | ```bash 59 | pip install -e external/acados/interfaces/acados_template 60 | ``` 61 | 62 | 3. Set environment variables (where `$PROJECT_DIR` is the directory you cloned the repository into in Step 1): 63 | ```bash 64 | export ACADOS_SOURCE_DIR=$PROJECT_DIR/external/acados 65 | export LD_LIBRARY_PATH=$ACADOS_SOURCE_DIR/lib 66 | ``` 67 | 68 | ### Install `L4acados` with optional dependencies of your choice 69 | 70 | Install `L4acados` with optional dependencies of your choice 71 | 72 | ```bash 73 | pip install -e .[] 74 | ``` 75 | 76 | Available options: 77 | - (without optional dependencies): Basic ResidualLearningMPC for custom implementations (e.g. Jacobian approximations, finite-difference approximations for e.g. models without sensitivity information) 78 | - `[pytorch]`: PyTorch models. 79 | - `[gpytorch]`: GPyTorch models. 80 | - `[gpytorch-exo]`: GpyTorch models with online-learning improvements. 81 | - not supported yet: JAX, TensorFlow 82 | 83 | 84 | ## Contributing 85 | 86 | If you would like to contribute features to `L4acados`, please follow the development installation instructions below to set up your working environment. 87 | 88 | ### Development installation 89 | 90 | 1. Install `L4acados` with development dependencies (in addition to the used learning framework, see above): 91 | 92 | ```bash 93 | pip install -e .[dev] 94 | ``` 95 | 96 | 2. Sync notebooks with jupytext: 97 | ```bash 98 | jupytext --set-formats ipynb,py examples/*/*.ipynb 99 | ``` 100 | 101 | 3. Add pre-commit hooks 102 | ```bash 103 | pre-commit install 104 | ``` 105 | 106 | ## Citing us 107 | 108 | If you use this software, please cite our corresponding articles as written below. 109 | 110 | ### General software 111 | 112 | ``` 113 | @online{lahr_l4acados_2024, 114 | title = {L4acados: {{Learning-based}} Models for Acados, Applied to {{Gaussian}} Process-Based Predictive Control}, 115 | shorttitle = {L4acados}, 116 | author = {Lahr, Amon and Näf, Joshua and Wabersich, Kim P. and Frey, Jonathan and Siehl, Pascal and Carron, Andrea and Diehl, Moritz and Zeilinger, Melanie N.}, 117 | date = {2024-11-28}, 118 | eprint = {2411.19258}, 119 | eprinttype = {arXiv}, 120 | doi = {10.48550/arXiv.2411.19258}, 121 | pubstate = {prepublished} 122 | } 123 | ``` 124 | 125 | ### Zero-Order GP-MPC algorithm 126 | 127 | ``` 128 | @article{lahr_zero-order_2023, 129 | title = {Zero-Order optimization for {{Gaussian}} process-based model predictive control}, 130 | author = {Lahr, Amon and Zanelli, Andrea and Carron, Andrea and Zeilinger, Melanie N.}, 131 | year = {2023}, 132 | journal = {European Journal of Control}, 133 | pages = {100862}, 134 | issn = {0947-3580}, 135 | doi = {10.1016/j.ejcon.2023.100862} 136 | } 137 | ``` 138 | -------------------------------------------------------------------------------- /examples/inverted_pendulum/inverted_pendulum_model_acados.py: -------------------------------------------------------------------------------- 1 | # %% 2 | import numpy as np 3 | from acados_template import AcadosModel, AcadosOcp 4 | from casadi import DM, SX, MX, vertcat, sin, cos, Function 5 | from scipy.linalg import block_diag 6 | 7 | 8 | # %% 9 | def export_simplependulum_ode_model( 10 | noise=False, 11 | only_lower_bounds=False, 12 | add_residual_dynamics=False, 13 | model_name="simplependulum_ode", 14 | ): 15 | # set up states & controls 16 | theta = SX.sym("theta") 17 | dtheta = SX.sym("dtheta") 18 | 19 | x = vertcat(theta, dtheta) 20 | u = SX.sym("u") 21 | 22 | # xdot 23 | theta_dot = SX.sym("theta_dot") 24 | dtheta_dot = SX.sym("dtheta_dot") 25 | xdot = vertcat(theta_dot, dtheta_dot) 26 | 27 | # parameters 28 | p = [] 29 | if noise: 30 | w = SX.sym("w", 2, 1) 31 | p += [w] 32 | p = vertcat(*p) 33 | 34 | # dynamics 35 | f_expl = vertcat(dtheta, sin(theta) + u) 36 | if noise: 37 | f_expl += w 38 | if add_residual_dynamics: 39 | f_expl += vertcat(DM(0), 0.5 * sin(theta**2)) 40 | 41 | f_impl = xdot - f_expl 42 | 43 | # constraints 44 | if only_lower_bounds: 45 | con_h_expr = vertcat(theta, -theta) # for lower bound # for upper bound 46 | else: 47 | con_h_expr = vertcat( 48 | theta, # theta 49 | ) 50 | 51 | # acados model 52 | model = AcadosModel() 53 | model.f_impl_expr = f_impl 54 | model.f_expl_expr = f_expl 55 | model.x = x 56 | model.xdot = xdot 57 | model.u = u 58 | # model.z = z 59 | model.p = p 60 | model.con_h_expr = con_h_expr 61 | model.con_h_expr_e = con_h_expr 62 | model.name = model_name 63 | 64 | return model 65 | 66 | 67 | def export_ocp_nominal(N, T, ocp_opts=None, only_lower_bounds=False, **model_kwargs): 68 | # constraints 69 | x0 = np.array([np.pi, 0]) 70 | lb_u = -2.0 71 | ub_u = 2.0 72 | lb_theta = (0.0 / 360.0) * 2 * np.pi 73 | ub_theta = (200.0 / 360.0) * 2 * np.pi 74 | 75 | # cost 76 | cost_theta = 5 77 | cost_omega = 1 78 | cost_fac_e = 4 79 | Q = np.diagflat(np.array([cost_theta, cost_omega])) 80 | R = np.array(1) 81 | 82 | model = export_simplependulum_ode_model( 83 | only_lower_bounds=only_lower_bounds, **model_kwargs 84 | ) 85 | 86 | # generate acados OCP for INITITIALIZATION 87 | ocp = AcadosOcp() 88 | 89 | ocp.model = model 90 | ocp.dims.N = N 91 | 92 | # dimensions 93 | nx = model.x.shape[0] 94 | nu = model.u.shape[0] 95 | ny = nx + nu 96 | ny_e = nx 97 | 98 | ocp.dims.nx = nx 99 | ocp.dims.nu = nu 100 | ocp.dims.nh = model.con_h_expr.shape[0] 101 | ocp.dims.np = model.p.shape[0] if isinstance(model.p, SX) else 0 102 | 103 | # cost 104 | ocp.cost.cost_type = "LINEAR_LS" 105 | ocp.cost.cost_type_e = "LINEAR_LS" 106 | ocp.cost.W = block_diag(Q, R) 107 | ocp.cost.W_0 = ocp.cost.W 108 | ocp.cost.W_e = cost_fac_e * Q 109 | 110 | ocp.cost.Vx = np.zeros((ny, nx)) 111 | ocp.cost.Vx[:nx, :nx] = np.eye(nx) 112 | 113 | ocp.cost.Vu = np.zeros((ny, nu)) 114 | ocp.cost.Vu[nx : nx + nu, :] = np.eye(nu) 115 | 116 | ocp.cost.Vx_e = np.eye(ny_e) 117 | 118 | ocp.cost.yref = np.zeros((ny,)) 119 | ocp.cost.yref_e = np.zeros((ny_e,)) 120 | 121 | # constraints 122 | ocp.constraints.constr_type = "BGH" 123 | ocp.constraints.lbu = np.array([lb_u]) 124 | ocp.constraints.ubu = np.array([ub_u]) 125 | ocp.constraints.idxbu = np.array(range(nu)) 126 | 127 | # TODO: automate this for arbitrary models (and tightened-constraint indices) 128 | if only_lower_bounds: 129 | inf_num = 1e6 130 | ocp.constraints.lh = np.array([lb_theta, -ub_theta]) 131 | ocp.constraints.uh = np.array([inf_num, inf_num]) 132 | else: 133 | ocp.constraints.lh = np.array([lb_theta]) 134 | ocp.constraints.uh = np.array([ub_theta]) 135 | 136 | # terminal constraints 137 | ocp.constraints.lh_e = ocp.constraints.lh 138 | ocp.constraints.uh_e = ocp.constraints.uh 139 | ocp.constraints.C_e = ocp.constraints.C 140 | ocp.constraints.lg_e = ocp.constraints.lg 141 | ocp.constraints.ug_e = ocp.constraints.ug 142 | 143 | ocp.constraints.x0 = x0 144 | 145 | # solver options 146 | 147 | ocp.solver_options.integrator_type = "ERK" 148 | # ocp.solver_options.integrator_type = 'DISCRETE' # 'IRK' 149 | 150 | ocp.solver_options.qp_solver = "PARTIAL_CONDENSING_HPIPM" # FULL_CONDENSING_QPOASES 151 | # ocp.solver_options.qp_solver = 'FULL_CONDENSING_QPOASES' 152 | 153 | ocp.solver_options.hessian_approx = "GAUSS_NEWTON" 154 | # ocp.solver_options.hessian_approx = 'EXACT' 155 | 156 | ocp.solver_options.nlp_solver_type = "SQP" 157 | ocp.solver_options.nlp_solver_max_iter = 50 158 | # ocp.solver_options.nlp_solver_type = 'SQP' # , SQP_RTI 159 | 160 | ocp.solver_options.tf = T 161 | ocp.solver_options.Tsim = T / N 162 | 163 | ocp.solver_options.tol = 1e-4 164 | 165 | return ocp 166 | -------------------------------------------------------------------------------- /examples/inverted_pendulum/inverted_pendulum_residual_learning_zoro.py: -------------------------------------------------------------------------------- 1 | # --- 2 | # jupyter: 3 | # jupytext: 4 | # formats: ipynb,py 5 | # text_representation: 6 | # extension: .py 7 | # format_name: light 8 | # format_version: '1.5' 9 | # jupytext_version: 1.16.4 10 | # kernelspec: 11 | # display_name: l4acados_dev 12 | # language: python 13 | # name: python3 14 | # --- 15 | 16 | # + endofcell="--" 17 | # # + metadata={} 18 | import sys, os 19 | 20 | sys.path += ["../../external/"] 21 | 22 | # # + metadata={} 23 | # %load_ext autoreload 24 | # %autoreload 1 25 | # %aimport l4acados 26 | 27 | # # + metadata={} 28 | import numpy as np 29 | from scipy.stats import norm 30 | import casadi as cas 31 | from acados_template import ( 32 | AcadosOcp, 33 | AcadosSim, 34 | AcadosSimSolver, 35 | AcadosOcpSolver, 36 | AcadosOcpOptions, 37 | ZoroDescription, 38 | ) 39 | import matplotlib.pyplot as plt 40 | import torch 41 | import gpytorch 42 | import copy 43 | 44 | # zoRO imports 45 | import l4acados 46 | from l4acados.controllers import ( 47 | ZeroOrderGPMPC, 48 | ) 49 | from l4acados.controllers.zoro_acados_utils import setup_sim_from_ocp 50 | from inverted_pendulum_model_acados import ( 51 | export_simplependulum_ode_model, 52 | export_ocp_nominal, 53 | ) 54 | from utils import * 55 | 56 | # gpytorch_utils 57 | from gpytorch_utils.gp_hyperparam_training import ( 58 | generate_train_inputs_acados, 59 | generate_train_outputs_at_inputs, 60 | train_gp_model, 61 | ) 62 | from gpytorch_utils.gp_utils import ( 63 | gp_data_from_model_and_path, 64 | gp_derivative_data_from_model_and_path, 65 | plot_gp_data, 66 | generate_grid_points, 67 | ) 68 | from l4acados.models.pytorch_models.gpytorch_models.gpytorch_gp import ( 69 | BatchIndependentMultitaskGPModel, 70 | ) 71 | 72 | # - 73 | 74 | 75 | # ## Define model parameters 76 | # 77 | # We model the inverted pendulum 78 | # 79 | # $$ 80 | # \dot{x} = f(x,u) = \begin{bmatrix} \dot{\theta} \\ \ddot{\theta} \end{bmatrix} = \begin{bmatrix} \dot{\theta} \\ -\sin(\theta) + u \end{bmatrix}, 81 | # $$ 82 | # 83 | # which is to be controlled from the hanging-down resting position, $(\theta_0, \dot{\theta}_0) = (\pi, 0)$, to the upright position $(\theta_r, \dot{\theta}_r) = (0,0)$, subject to the constraints that overshoot should be avoided, i.e., 84 | # 85 | # $$ 86 | # \theta_{lb} \leq \theta \leq \theta_{ub}. 87 | # $$ 88 | # 89 | # The model setup and controller definition can be found in the functions `export_simplependulum_ode_model()`, `export_ocp_nominal()` in the `inverted_pendulum_model_acados.py` file. 90 | 91 | # # + metadata={} 92 | # build C code again? 93 | build_c_code = True 94 | 95 | # # + metadata={} 96 | # discretization 97 | N = 30 98 | T = 5 99 | dT = T / N 100 | 101 | # constraints 102 | x0 = np.array([np.pi, 0]) 103 | nx = 2 104 | nu = 1 105 | 106 | 107 | # # + metadata={} 108 | prob_x = 0.95 109 | prob_tighten = norm.ppf(prob_x) 110 | 111 | # noise 112 | # uncertainty dynamics 113 | sigma_theta = (0.0001 / 360.0) * 2 * np.pi 114 | sigma_omega = (0.0001 / 360.0) * 2 * np.pi 115 | w_theta = 0.005 116 | w_omega = 0.005 117 | Sigma_x0 = np.array([[sigma_theta**2, 0], [0, sigma_omega**2]]) 118 | Sigma_W = np.array([[w_theta**2, 0], [0, w_omega**2]]) 119 | # - 120 | 121 | # ## Set up nominal solver 122 | 123 | # # + metadata={} 124 | ocp_init = export_ocp_nominal(N, T, model_name="simplependulum_ode_init") 125 | ocp_init.solver_options.nlp_solver_type = "SQP" 126 | 127 | # # + metadata={} 128 | ocp_init.solver_options.Tsim 129 | 130 | # # + metadata={} 131 | acados_ocp_init_solver = AcadosOcpSolver( 132 | ocp_init, json_file="acados_ocp_init_simplependulum_ode.json" 133 | ) 134 | 135 | # # + metadata={} 136 | ocp_init.solver_options.Tsim 137 | # - 138 | 139 | # ## Open-loop planning with nominal solver 140 | 141 | # # + metadata={} 142 | X_init, U_init = get_solution(acados_ocp_init_solver, x0) 143 | 144 | # # + metadata={} 145 | # integrator for nominal model 146 | sim = setup_sim_from_ocp(ocp_init) 147 | 148 | acados_integrator = AcadosSimSolver( 149 | sim, json_file="acados_sim_" + sim.model.name + ".json" 150 | ) 151 | # - 152 | 153 | # ## Simulator object 154 | # 155 | # To automatically discretize the model (and obtain sensitivities of the discrete-time model) within the zero-order implementation, we create the `AcadosSimSolver` object to pass to the solver. 156 | 157 | # # + metadata={} 158 | # generate training data for GP with "real model" 159 | model_actual = export_simplependulum_ode_model( 160 | model_name=sim.model.name + "_actual", add_residual_dynamics=True 161 | ) 162 | 163 | sim_actual = setup_sim_from_ocp(ocp_init) 164 | sim_actual.model = model_actual 165 | 166 | # acados_ocp_solver = AcadosOcpSolver(ocp, json_file = 'acados_ocp_' + model.name + '.json') 167 | acados_integrator_actual = AcadosSimSolver( 168 | sim_actual, json_file="acados_sim_" + model_actual.name + ".json" 169 | ) 170 | # - 171 | 172 | # ## Simulation results (nominal) 173 | 174 | # # + metadata={} 175 | X_init_sim = simulate_solution(acados_integrator_actual, x0, N, nx, nu, U_init) 176 | 177 | # # + metadata={} 178 | lb_theta = -ocp_init.constraints.lh[0] 179 | fig, ax = base_plot(lb_theta=lb_theta) 180 | 181 | plot_data_nom = EllipsoidTubeData2D(center_data=X_init, ellipsoid_data=None) 182 | plot_data_nom_sim = EllipsoidTubeData2D(center_data=X_init_sim, ellipsoid_data=None) 183 | add_plot_trajectory(ax, plot_data_nom, prob_tighten=None, color_fun=plt.cm.Blues) 184 | add_plot_trajectory(ax, plot_data_nom_sim, prob_tighten=None, color_fun=plt.cm.Blues) 185 | # - 186 | 187 | # # GP training 188 | # 189 | # We use a model with different parameters to emulate the real-world model and obtain some training data. Also create simulator object for real-world model to evaluate our results later (not used in solver). 190 | 191 | # ## Generate training data 192 | # 193 | # We generate training data (one-step ahead residuals `y_train` for starting point `x_train`) here by running robustified (cautious) solver without GP. 194 | 195 | # # + metadata={} 196 | random_seed = 123 197 | N_sim_per_x0 = 1 198 | N_x0 = 10 199 | x0_rand_scale = 0.1 200 | 201 | x_train, x0_arr = generate_train_inputs_acados( 202 | acados_ocp_init_solver, 203 | x0, 204 | N_sim_per_x0, 205 | N_x0, 206 | random_seed=random_seed, 207 | x0_rand_scale=x0_rand_scale, 208 | ) 209 | 210 | y_train = generate_train_outputs_at_inputs( 211 | x_train, acados_integrator, acados_integrator_actual, Sigma_W 212 | ) 213 | 214 | # # + metadata={} 215 | x_train 216 | # - 217 | 218 | # ## Hyper-parameter training for GP model 219 | # 220 | # Optimize hyper-parameters of GP model (kernel function parameters, ...) 221 | 222 | # # + metadata={} 223 | x_train_tensor = torch.Tensor(x_train) 224 | y_train_tensor = torch.Tensor(y_train) 225 | nout = y_train.shape[1] 226 | 227 | likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=nout) 228 | gp_model = BatchIndependentMultitaskGPModel(x_train_tensor, y_train_tensor, likelihood) 229 | 230 | # # + metadata={} 231 | load_gp_model_from_state_dict = False 232 | state_dict_path_gp_model = "gp_model_state_dict.pth" 233 | state_dict_path_likelihood = "gp_model_likelihood_state_dict.pth" 234 | train_data_path = "gp_model_train_data.pth" 235 | 236 | if load_gp_model_from_state_dict: 237 | # Load state dict 238 | gp_model.load_state_dict(torch.load(state_dict_path_gp_model)) 239 | likelihood.load_state_dict(torch.load(state_dict_path_likelihood)) 240 | else: 241 | training_iterations = 200 242 | rng_seed = 456 243 | 244 | gp_model, likelihood = train_gp_model( 245 | gp_model, torch_seed=rng_seed, training_iterations=training_iterations 246 | ) 247 | 248 | # EVAL MODE 249 | gp_model.eval() 250 | likelihood.eval() 251 | 252 | # # + metadata={} 253 | # save GP hyper-params 254 | torch.save(gp_model.state_dict(), state_dict_path_gp_model) 255 | torch.save(likelihood.state_dict(), state_dict_path_likelihood) 256 | torch.save({"x_train": x_train_tensor, "y_train": y_train_tensor}, train_data_path) 257 | 258 | 259 | # # + metadata={} 260 | data_dict = torch.load(train_data_path) 261 | data_dict 262 | # - 263 | 264 | # ## Plot GP predictions 265 | # 266 | # We plot GP predictions along the predicted trajectory of the robustified solver by projecting the multivariate plot down to a line. 267 | 268 | # # + metadata={} 269 | x_train.shape, y_train.shape 270 | 271 | # # + metadata={} 272 | num_samples = 5 273 | use_likelihood = False 274 | 275 | num_points_between_samples = 30 276 | t_lin = np.linspace(0, 1, num_points_between_samples, endpoint=False) 277 | 278 | x_plot_waypts = np.hstack((X_init[1:, :], U_init)) 279 | x_plot = [] 280 | for i in range(x_plot_waypts.shape[0] - 1): 281 | x_plot += [ 282 | x_plot_waypts[i, :] + (x_plot_waypts[i + 1, :] - x_plot_waypts[i, :]) * t 283 | for t in t_lin 284 | ] 285 | x_plot = np.vstack(x_plot) 286 | 287 | gp_data = gp_data_from_model_and_path( 288 | gp_model, likelihood, x_plot, num_samples=num_samples, use_likelihood=use_likelihood 289 | ) 290 | plot_gp_data([gp_data], marker_size_lim=[1, 15]) 291 | # - 292 | 293 | # We can also plot the derivative of the GP. Note that the projected Jacobian is not smooth since our path is not smooth either (jump projection direction = jump in Jacobian); however, the actual Jacobian should be smooth here (squared exponential kernel). 294 | 295 | # # + metadata={} 296 | gp_derivative_data = gp_derivative_data_from_model_and_path( 297 | gp_model, likelihood, x_plot, num_samples=0 298 | ) 299 | plot_gp_data([gp_derivative_data], marker_size_lim=[5, 20], plot_train_data=False) 300 | # - 301 | 302 | # Compare with plotting along a slice of the dimension. Since we generated training data along the path of the robustified controller, the GP looks pretty untrained along a slice of the coordinates. 303 | 304 | # # + metadata={} 305 | # plot along axis 306 | x_dim_lims = np.array([[0, np.pi], [-2, 1], [-2, 2]]) 307 | x_dim_slice = np.array([1 * np.pi, 0, 0]) 308 | x_dim_plot = 2 309 | x_grid = generate_grid_points(x_dim_lims, x_dim_slice, x_dim_plot, num_points=800) 310 | 311 | gp_grid_data = gp_data_from_model_and_path( 312 | gp_model, likelihood, x_grid, num_samples=num_samples, use_likelihood=use_likelihood 313 | ) 314 | fig, ax = plot_gp_data([gp_grid_data], marker_size_lim=[5, 50]) 315 | 316 | y_lim_0 = ax[0].get_ylim() 317 | y_lim_1 = ax[1].get_ylim() 318 | # - 319 | 320 | # Jacobian... not much going on away from the data points (this is good!) 321 | 322 | # # + metadata={} 323 | gp_derivative_grid_data = gp_derivative_data_from_model_and_path( 324 | gp_model, likelihood, x_grid, num_samples=0 325 | ) 326 | fig, ax = plot_gp_data( 327 | [gp_derivative_grid_data], marker_size_lim=[5, 50], plot_train_data=False 328 | ) 329 | 330 | ax[0].set_ylim(*y_lim_0) 331 | ax[1].set_ylim(*y_lim_1) 332 | plt.draw() 333 | # - 334 | 335 | # # Residual-Model MPC 336 | # -- 337 | 338 | # + endofcell="--" metadata={} 339 | from l4acados.models import GPyTorchResidualModel 340 | 341 | # # + metadata={} 342 | residual_model = GPyTorchResidualModel(gp_model) 343 | 344 | # # + metadata={} 345 | residual_model.evaluate(x_plot_waypts[0:3, :]) 346 | 347 | # # + metadata={} 348 | residual_model.jacobian(x_plot_waypts[0:3, :]) 349 | 350 | # # + metadata={} 351 | residual_model.value_and_jacobian(x_plot_waypts[0:3, :]) 352 | # - 353 | 354 | # create zoro_description 355 | zoro_description = ZoroDescription() 356 | zoro_description.backoff_scaling_gamma = norm.ppf(prob_x) 357 | zoro_description.P0_mat = Sigma_x0 358 | zoro_description.fdbk_K_mat = np.zeros((nu, nx)) 359 | # zoro_description.unc_jac_G_mat = B 360 | """G in (nx, nw) describes how noise affects dynamics. I.e. x+ = ... + G@w""" 361 | zoro_description.W_mat = Sigma_W 362 | """W in (nw, nw) describes the covariance of the noise on the system""" 363 | zoro_description.input_P0_diag = True 364 | zoro_description.input_P0 = False 365 | zoro_description.input_W_diag = True 366 | zoro_description.input_W_add_diag = True 367 | zoro_description.output_P_matrices = True 368 | zoro_description.idx_lh_t = [0] 369 | zoro_description.idx_lh_e_t = [0] 370 | ocp_init.zoro_description = zoro_description 371 | 372 | # # + metadata={} 373 | residual_mpc = ZeroOrderGPMPC( 374 | ocp_init, 375 | residual_model=residual_model, 376 | use_cython=False, 377 | path_json_ocp="residual_mpc_ocp_solver_config.json", 378 | path_json_sim="residual_mpc_sim_solver_config.json", 379 | build_c_code=True, 380 | ) 381 | 382 | # # + metadata={} 383 | for i in range(N): 384 | residual_mpc.ocp_solver.set(i, "x", X_init[i, :]) 385 | residual_mpc.ocp_solver.set(i, "u", U_init[i, :]) 386 | residual_mpc.ocp_solver.set(N, "x", X_init[N, :]) 387 | 388 | residual_mpc.solve(acados_sqp_mode=True) 389 | X_res, U_res = residual_mpc.get_solution() 390 | P_res_arr = residual_mpc.covariances_array 391 | 392 | P_res = [] 393 | for i in range(N + 1): 394 | P_res.append(P_res_arr[i * nx**2 : (i + 1) * nx**2].reshape((nx, nx))) 395 | P_res = np.array(P_res) 396 | 397 | # # + metadata={} 398 | X_res_sim = np.zeros_like(X_res) 399 | X_res_sim[0, :] = x0 400 | for i in range(N): 401 | acados_integrator_actual.set("x", X_res_sim[i, :]) 402 | acados_integrator_actual.set("u", U_res[i, :]) 403 | acados_integrator_actual.solve() 404 | X_res_sim[i + 1, :] = acados_integrator_actual.get("x") 405 | 406 | # # + metadata={} 407 | lb_theta = -ocp_init.constraints.lh[0] 408 | fig, ax = base_plot(lb_theta=lb_theta) 409 | 410 | plot_data_res = EllipsoidTubeData2D(center_data=X_res, ellipsoid_data=P_res) 411 | plot_data_res_sim = EllipsoidTubeData2D(center_data=X_res_sim, ellipsoid_data=None) 412 | add_plot_trajectory(ax, plot_data_nom, color_fun=plt.cm.Blues) 413 | add_plot_trajectory(ax, plot_data_nom_sim, color_fun=plt.cm.Blues) 414 | add_plot_trajectory( 415 | ax, plot_data_res, prob_tighten=prob_tighten, color_fun=plt.cm.Oranges 416 | ) 417 | add_plot_trajectory(ax, plot_data_res_sim, color_fun=plt.cm.Oranges) 418 | # -- 419 | 420 | residual_mpc.print_statistics() 421 | 422 | residual_mpc.get_stats("res_stat_all") 423 | 424 | residual_mpc.get_stats("residuals") 425 | 426 | residual_mpc.get_stats("time_preparation_all"), residual_mpc.get_stats( 427 | "time_preparation" 428 | ) 429 | -------------------------------------------------------------------------------- /examples/inverted_pendulum/utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy import linalg as npla 3 | from dataclasses import dataclass 4 | import matplotlib.pyplot as plt 5 | from inverted_pendulum_model_acados import export_simplependulum_ode_model 6 | 7 | # Plotting 8 | 9 | 10 | @dataclass 11 | class EllipsoidTubeData2D: 12 | center_data: np.ndarray = (None,) 13 | ellipsoid_data: np.ndarray = (None,) 14 | ellipsoid_colors: np.ndarray = None 15 | 16 | 17 | def base_plot(lb_theta=None): 18 | fig, ax = plt.subplots() 19 | 20 | if lb_theta is not None: 21 | ax.axvline(lb_theta) 22 | 23 | return fig, ax 24 | 25 | 26 | def add_plot_ellipse(ax, E, e0, n=50, **plot_args): 27 | # sample angle uniformly from [0,2pi] and length from [0,1] 28 | radius = 1.0 29 | theta_arr = np.linspace(0, 2 * np.pi, n) 30 | w_rad_arr = [[radius, theta] for theta in theta_arr] 31 | w_one_arr = np.array( 32 | [ 33 | [w_rad[0] * np.cos(w_rad[1]), w_rad[0] * np.sin(w_rad[1])] 34 | for w_rad in w_rad_arr 35 | ] 36 | ) 37 | w_ell = np.array([e0 + E @ w_one for w_one in w_one_arr]) 38 | h = ax.plot(w_ell[:, 0], w_ell[:, 1], **plot_args) 39 | return h 40 | 41 | 42 | def add_plot_trajectory( 43 | ax, 44 | tube_data: EllipsoidTubeData2D, 45 | color_fun=plt.cm.Blues, 46 | prob_tighten=1, 47 | **plot_args, 48 | ): 49 | n_data = tube_data.center_data.shape[0] 50 | evenly_spaced_interval = np.linspace(0.6, 1, n_data) 51 | colors = [color_fun(x) for x in evenly_spaced_interval] 52 | 53 | h_plot = ax.plot( 54 | tube_data.center_data[:, 0], tube_data.center_data[:, 1], **plot_args 55 | ) 56 | # set color 57 | h_plot[0].set_color(colors[-1]) 58 | 59 | for i, color in enumerate(colors): 60 | center_i = tube_data.center_data[i, :] 61 | if tube_data.ellipsoid_data is not None: 62 | ellipsoid_i = tube_data.ellipsoid_data[i, :, :] 63 | # get eigenvalues 64 | eig_val, eig_vec = npla.eig(ellipsoid_i) 65 | ellipsoid_i_sqrt = ( 66 | prob_tighten 67 | * eig_vec 68 | @ np.diag(np.sqrt(np.maximum(eig_val, 0))) 69 | @ np.transpose(eig_vec) 70 | ) 71 | # print(i, eig_val, ellipsoid_i) 72 | h_ell = add_plot_ellipse(ax, ellipsoid_i_sqrt, center_i, **plot_args) 73 | h_ell[0].set_color(color) 74 | 75 | 76 | # OCP stuff 77 | def get_solution(ocp_solver, x0): 78 | N = ocp_solver.acados_ocp.dims.N 79 | 80 | # get initial values 81 | X = np.zeros((N + 1, ocp_solver.acados_ocp.dims.nx)) 82 | U = np.zeros((N, ocp_solver.acados_ocp.dims.nu)) 83 | 84 | # xcurrent = x0 85 | X[0, :] = x0 86 | 87 | # solve 88 | status = ocp_solver.solve() 89 | 90 | if status != 0: 91 | raise Exception("acados ocp_solver returned status {}. Exiting.".format(status)) 92 | 93 | # get data 94 | for i in range(N): 95 | X[i, :] = ocp_solver.get(i, "x") 96 | U[i, :] = ocp_solver.get(i, "u") 97 | 98 | X[N, :] = ocp_solver.get(N, "x") 99 | return X, U 100 | 101 | 102 | def simulate_solution(sim_solver, x0, N, nx, nu, U): 103 | # get initial values 104 | X = np.zeros((N + 1, nx)) 105 | 106 | # xcurrent = x0 107 | X[0, :] = x0 108 | 109 | # simulate 110 | for i in range(N): 111 | sim_solver.set("x", X[i, :]) 112 | sim_solver.set("u", U[i, :]) 113 | status = sim_solver.solve() 114 | if status != 0: 115 | raise Exception( 116 | "acados sim_solver returned status {}. Exiting.".format(status) 117 | ) 118 | X[i + 1, :] = sim_solver.get("x") 119 | 120 | return X 121 | 122 | 123 | def init_ocp_solver(ocp_solver, X, U): 124 | # initialize with nominal solution 125 | N = ocp_solver.acados_ocp.dims.N 126 | print(f"N = {N}, size_X = {X.shape}") 127 | for i in range(N): 128 | ocp_solver.set(i, "x", X[i, :]) 129 | ocp_solver.set(i, "u", U[i, :]) 130 | ocp_solver.set(N, "x", X[N, :]) 131 | 132 | 133 | # GP model 134 | 135 | 136 | def get_gp_model(ocp_solver, sim_solver, sim_solver_actual, x0, Sigma_W): 137 | random_seed = 123 138 | N_sim_per_x0 = 1 139 | N_x0 = 10 140 | x0_rand_scale = 0.1 141 | 142 | x_train, x0_arr = generate_train_inputs_acados( 143 | ocp_solver, 144 | x0, 145 | N_sim_per_x0, 146 | N_x0, 147 | random_seed=random_seed, 148 | x0_rand_scale=x0_rand_scale, 149 | ) 150 | 151 | y_train = generate_train_outputs_at_inputs( 152 | x_train, sim_solver, sim_solver_actual, Sigma_W 153 | ) 154 | -------------------------------------------------------------------------------- /examples/l4casadi_vs_l4acados/README.md: -------------------------------------------------------------------------------- 1 | # Example: L4CasADi vs. L4acados 2 | 3 | ## Installation 4 | 5 | To execute, you need to additionally install L4CasADi (https://github.com/Tim-Salzmann/l4casadi). 6 | -------------------------------------------------------------------------------- /examples/l4casadi_vs_l4acados/acados.py: -------------------------------------------------------------------------------- 1 | import casadi as cs 2 | import numpy as np 3 | import torch 4 | import l4casadi as l4c 5 | from acados_template import AcadosOcpSolver, AcadosOcp, AcadosModel 6 | import time 7 | 8 | COST = "LINEAR_LS" # NONLINEAR_LS 9 | 10 | 11 | class MultiLayerPerceptron(torch.nn.Module): 12 | def __init__(self, n_inputs=2, hidden_layers=2, hidden_size=512, n_outputs=1): 13 | super().__init__() 14 | 15 | self.input_layer = torch.nn.Linear(n_inputs, hidden_size) 16 | 17 | all_hidden_layers = [] 18 | for i in range(hidden_layers): 19 | all_hidden_layers.append(torch.nn.Linear(hidden_size, hidden_size)) 20 | 21 | self.hidden_layer = torch.nn.ModuleList(all_hidden_layers) 22 | self.out_layer = torch.nn.Linear(hidden_size, n_outputs) 23 | 24 | # Model is not trained -- setting output to zero 25 | with torch.no_grad(): 26 | self.out_layer.bias.fill_(0.0) 27 | self.out_layer.weight.fill_(0.0) 28 | 29 | def forward(self, x): 30 | x = self.input_layer(x) 31 | for layer in self.hidden_layer: 32 | x = torch.tanh(layer(x)) 33 | x = self.out_layer(x) 34 | return x 35 | 36 | 37 | class DoubleIntegratorWithLearnedDynamics: 38 | def __init__(self, learned_dyn=None, name="wr", batched=False, batch_dim=1): 39 | self.learned_dyn = learned_dyn 40 | self.name = name 41 | self.batched = batched 42 | self.batch_dim = batch_dim 43 | 44 | def model(self): 45 | s = cs.MX.sym("s", 1) 46 | s_dot = cs.MX.sym("s_dot", 1) 47 | s_dot_dot = cs.MX.sym("s_dot_dot", 1) 48 | u = cs.MX.sym("u", 1) 49 | x = cs.vertcat(s, s_dot) 50 | x_dot = cs.vertcat(s_dot, s_dot_dot) 51 | 52 | if self.batched: 53 | x_l4casadi = cs.repmat(cs.reshape(x, 1, 2), self.batch_dim, 1) 54 | else: 55 | x_l4casadi = x 56 | 57 | if self.learned_dyn is None: 58 | res_model = cs.MX.zeros(2, 1) 59 | else: 60 | res_model = self.learned_dyn(x_l4casadi) 61 | 62 | B = cs.DM.ones(self.batch_dim, 1) 63 | 64 | # f_expl = cs.vertcat(s_dot, u - cs.sin(3.1 * s) ** 2) + B.T @ res_model 65 | f_expl = cs.vertcat(s_dot, u) + B.T @ res_model 66 | 67 | x_start = np.zeros((2,)) 68 | 69 | # store to struct 70 | model = cs.types.SimpleNamespace() 71 | model.x = x 72 | model.xdot = x_dot 73 | model.u = u 74 | model.z = cs.vertcat([]) 75 | model.p = cs.vertcat([]) 76 | model.f_expl = f_expl 77 | model.x_start = x_start 78 | model.constraints = cs.vertcat([]) 79 | model.name = self.name 80 | 81 | return model 82 | 83 | 84 | class MPC: 85 | def __init__( 86 | self, 87 | model, 88 | N, 89 | external_shared_lib_dir=None, 90 | external_shared_lib_name=None, 91 | num_threads_acados_openmp=1, 92 | ): 93 | self.N = N 94 | self.model = model 95 | self.external_shared_lib_dir = external_shared_lib_dir 96 | self.external_shared_lib_name = external_shared_lib_name 97 | self.num_threads_acados_openmp = num_threads_acados_openmp 98 | 99 | @property 100 | def solver(self): 101 | return AcadosOcpSolver(self.ocp()) 102 | 103 | def ocp(self): 104 | model = self.model 105 | 106 | t_horizon = 1.0 107 | N = self.N 108 | 109 | # Get model 110 | model_ac = self.acados_model(model=model) 111 | model_ac.p = model.p 112 | 113 | # Dimensions 114 | nx = 2 115 | nu = 1 116 | ny = 1 117 | 118 | # Create OCP object to formulate the optimization 119 | ocp = AcadosOcp() 120 | ocp.model = model_ac 121 | 122 | ocp.dims.N = N 123 | ocp.dims.nx = nx 124 | ocp.dims.nu = nu 125 | ocp.dims.ny = ny 126 | ocp.solver_options.tf = t_horizon 127 | 128 | if COST == "LINEAR_LS": 129 | # Initialize cost function 130 | ocp.cost.cost_type = "LINEAR_LS" 131 | ocp.cost.cost_type_e = "LINEAR_LS" 132 | 133 | ocp.cost.W = np.array([[1.0]]) 134 | 135 | ocp.cost.Vx = np.zeros((ny, nx)) 136 | ocp.cost.Vx[0, 0] = 1.0 137 | ocp.cost.Vu = np.zeros((ny, nu)) 138 | ocp.cost.Vz = np.array([[]]) 139 | ocp.cost.Vx_e = np.zeros((ny, nx)) 140 | 141 | l4c_y_expr = None 142 | else: 143 | ocp.cost.cost_type = "NONLINEAR_LS" 144 | ocp.cost.cost_type_e = "NONLINEAR_LS" 145 | 146 | x = ocp.model.x 147 | 148 | ocp.cost.W = np.array([[1.0]]) 149 | 150 | # Trivial PyTorch index 0 151 | l4c_y_expr = l4c.L4CasADi( 152 | lambda x: x[0], name="y_expr", model_expects_batch_dim=False 153 | ) 154 | 155 | ocp.model.cost_y_expr = l4c_y_expr(x) 156 | ocp.model.cost_y_expr_e = x[0] 157 | 158 | ocp.cost.W_e = np.array([[0.0]]) 159 | ocp.cost.yref_e = np.array([0.0]) 160 | 161 | # Initial reference trajectory (will be overwritten) 162 | ocp.cost.yref = np.zeros(1) 163 | 164 | # Initial state (will be overwritten) 165 | ocp.constraints.x0 = model.x_start 166 | 167 | # Set constraints 168 | a_max = 10 169 | ocp.constraints.lbu = np.array([-a_max]) 170 | ocp.constraints.ubu = np.array([a_max]) 171 | ocp.constraints.idxbu = np.array([0]) 172 | 173 | # Solver options 174 | ocp.solver_options.qp_solver = "FULL_CONDENSING_HPIPM" 175 | ocp.solver_options.hessian_approx = "GAUSS_NEWTON" 176 | # ocp.solver_options.hessian_approx = "EXACT" 177 | ocp.solver_options.integrator_type = "ERK" 178 | ocp.solver_options.nlp_solver_type = "SQP_RTI" 179 | ocp.solver_options.sim_method_num_stages = 1 # Euler 180 | ocp.solver_options.sim_method_num_steps = 1 181 | ocp.solver_options.nlp_solver_max_iter = 1 182 | # ocp.solver_options.num_threads_in_batch_solve = self.num_threads_acados_openmp 183 | 184 | if self.external_shared_lib_dir is not None: 185 | ocp.solver_options.model_external_shared_lib_dir = ( 186 | self.external_shared_lib_dir 187 | ) 188 | 189 | if COST == "LINEAR_LS": 190 | ocp.solver_options.model_external_shared_lib_name = ( 191 | self.external_shared_lib_name 192 | ) 193 | else: 194 | ocp.solver_options.model_external_shared_lib_name = ( 195 | self.external_shared_lib_name + " -l" + l4c_y_expr.name 196 | ) 197 | 198 | return ocp 199 | 200 | def acados_model(self, model): 201 | model_ac = AcadosModel() 202 | model_ac.f_impl_expr = model.xdot - model.f_expl 203 | model_ac.f_expl_expr = model.f_expl 204 | model_ac.x = model.x 205 | model_ac.xdot = model.xdot 206 | model_ac.u = model.u 207 | model_ac.name = model.name 208 | return model_ac 209 | 210 | 211 | def run(): 212 | N = 10 213 | learned_dyn_model = l4c.L4CasADi( 214 | MultiLayerPerceptron(), model_expects_batch_dim=True, name="learned_dyn" 215 | ) 216 | 217 | model = DoubleIntegratorWithLearnedDynamics(learned_dyn_model) 218 | 219 | print(learned_dyn_model.shared_lib_dir) 220 | print(learned_dyn_model.name) 221 | 222 | solver = MPC( 223 | model=model.model(), 224 | N=N, 225 | external_shared_lib_dir=learned_dyn_model.shared_lib_dir, 226 | external_shared_lib_name=learned_dyn_model.name, 227 | ).solver 228 | 229 | x = [] 230 | x_ref = [] 231 | ts = 1.0 / N 232 | xt = np.array([1.0, 0.0]) 233 | opt_times = [] 234 | 235 | for i in range(50): 236 | now = time.time() 237 | t = np.linspace(i * ts, i * ts + 1.0, 10) 238 | yref = np.sin(0.5 * t + np.pi / 2) 239 | x_ref.append(yref[0]) 240 | for t, ref in enumerate(yref): 241 | solver.set(t, "yref", ref) 242 | solver.set(0, "lbx", xt) 243 | solver.set(0, "ubx", xt) 244 | solver.solve() 245 | xt = solver.get(1, "x") 246 | x.append(xt) 247 | 248 | x_l = [] 249 | for i in range(N): 250 | x_l.append(solver.get(i, "x")) 251 | 252 | elapsed = time.time() - now 253 | opt_times.append(elapsed) 254 | 255 | print( 256 | f"Mean iteration time: {1000*np.mean(opt_times):.1f}ms -- {1/np.mean(opt_times):.0f}Hz)" 257 | ) 258 | 259 | 260 | if __name__ == "__main__": 261 | run() 262 | -------------------------------------------------------------------------------- /examples/l4casadi_vs_l4acados/analyze_results.py: -------------------------------------------------------------------------------- 1 | # --- 2 | # jupyter: 3 | # jupytext: 4 | # formats: ipynb,py 5 | # text_representation: 6 | # extension: .py 7 | # format_name: light 8 | # format_version: '1.5' 9 | # jupytext_version: 1.16.4 10 | # kernelspec: 11 | # display_name: zero-order-gpmpc-package-3.10 12 | # language: python 13 | # name: python3 14 | # --- 15 | 16 | import numpy as np 17 | import matplotlib.pyplot as plt 18 | import os 19 | 20 | plt.rcParams.update( 21 | { 22 | "text.usetex": True, 23 | "font.family": "serif", 24 | "font.sans-serif": "Times", 25 | "font.size": 10, 26 | } 27 | ) 28 | 29 | dir_data = "data" 30 | 31 | # only use 4 plots -> leave out cpu-t0-a0 32 | N_arr = [int(i) for i in np.ceil(np.logspace(0, 3, 10))] 33 | batch_dim = 1 34 | hidden_layers_arr = [1, 10, 20] 35 | solve_steps = 1000 36 | device_arr = ["cpu", "cpu", "cuda", "cuda"] 37 | num_threads_arr = [1, 10, 1, 1] 38 | num_threads_acados_openmp_arr = [0, 0, 0, 10] 39 | warmup_iter = 1 40 | 41 | assert len(num_threads_arr) == len(device_arr) == len(num_threads_acados_openmp_arr) 42 | 43 | device_threads_arr = list( 44 | zip(device_arr, num_threads_arr, num_threads_acados_openmp_arr) 45 | ) 46 | device_threads_arr 47 | 48 | N_grid, hidden_layers_grid = np.meshgrid(N_arr, hidden_layers_arr) 49 | 50 | N_grid_flat = N_grid.flatten() 51 | hidden_layers_flat = hidden_layers_grid.flatten() 52 | 53 | N_experiments = len(N_arr) * len(hidden_layers_arr) * len(device_threads_arr) 54 | N_experiments 55 | 56 | # + 57 | avg_times = { 58 | "l4casadi": np.zeros((len(hidden_layers_arr), len(N_arr), len(device_threads_arr))), 59 | "l4acados": np.zeros((len(hidden_layers_arr), len(N_arr), len(device_threads_arr))), 60 | } 61 | 62 | for i, hidden_layers in enumerate(hidden_layers_arr): 63 | for j, N in enumerate(N_arr): 64 | for k, device_threads in enumerate(device_threads_arr): 65 | npzfile = np.load( 66 | os.path.join( 67 | dir_data, 68 | f"l4casadi_vs_l4acados_N{N}_layers{hidden_layers}_steps{solve_steps}_{device_threads[0]}_threads{device_threads[1]}_acados_{device_threads[2]}.npz", 69 | # f"l4casadi_vs_l4acados_N{N}_layers{hidden_layers}_steps{solve_steps}_{device_threads[0]}_threads{device_threads[1]}.npz", 70 | ) 71 | ) 72 | x_l4casadi = npzfile["x_l4casadi"] 73 | opt_times_l4casadi = npzfile["opt_times_l4casadi"] 74 | x_l4acados = npzfile["x_l4acados"] 75 | opt_times_l4acados = npzfile["opt_times_l4acados"] 76 | 77 | avg_times["l4casadi"][i, j, k] = np.mean(opt_times_l4casadi[warmup_iter:]) 78 | avg_times["l4acados"][i, j, k] = np.mean(opt_times_l4acados[warmup_iter:]) 79 | 80 | # + 81 | from matplotlib import colormaps 82 | 83 | cmap_blue = colormaps["Blues"] 84 | cmap_orange = colormaps["Oranges"] 85 | 86 | cmap_l4casadi = colormaps["winter"] 87 | # cmap_l4acados = colormaps["summer"] 88 | cmap_l4acados = colormaps["winter"] 89 | 90 | 91 | # cmap_l4acados = colormaps["YlOrBr"] 92 | # cmap_l4casadi = colormaps["YlOrBr"] 93 | # - 94 | 95 | import matplotlib.lines as mlines 96 | 97 | # + 98 | colormap_exp = 1.0 99 | fig, ax = plt.subplots( 100 | 1, len(device_threads_arr), figsize=(1.5 * 7, 1.5 * 2.5), sharey=True 101 | ) 102 | ax_ylim_arr = np.zeros((2, len(device_threads_arr))) 103 | 104 | for l, device_threads in enumerate(device_threads_arr): 105 | for i, hidden_layers in enumerate(hidden_layers_arr): 106 | h_l4casadi = ax[l].plot( 107 | N_arr, 108 | avg_times["l4casadi"][i, :, l], 109 | color=cmap_l4casadi(((i + 1) / len(hidden_layers_arr)) ** colormap_exp), 110 | linestyle="--", 111 | ) 112 | h_l4acados = ax[l].plot( 113 | N_arr, 114 | avg_times["l4acados"][i, :, l], 115 | color=cmap_l4acados(((i + 1) / len(hidden_layers_arr)) ** colormap_exp), 116 | ) 117 | 118 | ax_ylim_arr[:, l] += np.array(ax[l].get_ylim()) 119 | ax[l].set_xscale("log") 120 | ax[l].set_yscale("log") 121 | ax[l].set_xlabel("N") 122 | ax[l].set_title( 123 | f"{device_threads[0]}-torch{device_threads[1]}-acados{device_threads[2]}" 124 | ) 125 | ax[l].grid() 126 | 127 | ax_ylim_min = min(ax_ylim_arr[0, :]) 128 | ax_ylim_max = max(ax_ylim_arr[1, :]) 129 | for l in range(len(device_threads_arr)): 130 | ax[l].set_ylim((ax_ylim_min, ax_ylim_max)) 131 | ax[l].set_xlim( 132 | ( 133 | ax[l].get_xlim()[0], 134 | 10 135 | ** ( 136 | np.log10(ax[l].get_xlim()[0]) 137 | + np.log10(ax[l].get_ylim()[1]) 138 | - np.log10(ax[l].get_ylim()[0]) 139 | ), 140 | ) 141 | ) 142 | ax[l].set_box_aspect(1.0) 143 | 144 | legend_handles = [mlines.Line2D([], [], color="k", linestyle="-", label="l4acados")] 145 | legend_handles += [mlines.Line2D([], [], color="k", linestyle="--", label="l4casadi")] 146 | legend_handles += [ 147 | mlines.Line2D([], [], color="tab:gray", linestyle="none", label="", marker="+") 148 | ] 149 | legend_handles += [ 150 | mlines.Line2D( 151 | [], 152 | [], 153 | color=cmap_l4acados(((i + 1) / len(hidden_layers_arr)) ** colormap_exp), 154 | marker="o", 155 | linestyle="none", 156 | label=f"$\#_{{\\mathrm{{layers}}}} = {hidden_layers_arr[i]:d}$", 157 | ) 158 | for i in range(len(hidden_layers_arr) - 1, -1, -1) 159 | ] 160 | 161 | ax[0].set_ylabel("Time per SQP iteration [s]") 162 | ax[-1].legend(loc="center left", bbox_to_anchor=(1.05, 0.5), handles=legend_handles) 163 | plt.tight_layout() 164 | 165 | fig.savefig("l4casadi_vs_l4acados.pdf", dpi=600) 166 | # - 167 | 168 | avg_times 169 | 170 | # + 171 | from matplotlib.colors import LogNorm 172 | from matplotlib import colormaps 173 | 174 | # scale_exp = 1.5 175 | scale_max = 50 176 | scale_exp = np.log10(scale_max) 177 | width_ratios = np.ones( 178 | ( 179 | len( 180 | device_threads_arr, 181 | ) 182 | ) 183 | ) 184 | # width_ratios[-1] = 1.25 185 | width_ratios_normalized = width_ratios / sum(width_ratios) 186 | fig, ax = plt.subplots( 187 | 1, 188 | len(device_threads_arr), 189 | figsize=(7, 2.5), 190 | width_ratios=width_ratios_normalized, 191 | sharey=True, 192 | ) 193 | for l, device_threads in enumerate(device_threads_arr): 194 | c = ax[l].pcolor( 195 | N_grid, 196 | hidden_layers_grid, 197 | avg_times["l4casadi"][:, :, l] / avg_times["l4acados"][:, :, l], 198 | shading="nearest", 199 | norm=LogNorm(vmin=10 ** (-scale_exp), vmax=10**scale_exp), 200 | cmap=colormaps["BrBG"], 201 | ) 202 | ax[l].set_xlabel("N") 203 | ax[l].set_xticks(N_arr) 204 | ax[l].set_yticks(hidden_layers_arr) 205 | ax[l].set_xscale("log") 206 | ax[l].set_xscale("log") 207 | ax[l].set_xlim((min(N_arr), max(N_arr))) 208 | 209 | ax[l].set_box_aspect(1.0) 210 | 211 | ax[l].set_title(f"{device_threads[0]}-{device_threads[1]}") 212 | ax[0].set_ylabel("Hidden layers") 213 | 214 | scale_exp_int_range_min = int(np.ceil(-scale_exp)) 215 | scale_exp_int_range_max = int(np.floor(scale_exp)) 216 | scale_exp_int_range = ( 217 | [-scale_exp] 218 | + list(range(scale_exp_int_range_min, scale_exp_int_range_max + 1)) 219 | + [scale_exp] 220 | ) 221 | cbar_yticks_list = ( 222 | [10 ** (-scale_exp)] 223 | + [10**i for i in list(range(scale_exp_int_range_min, scale_exp_int_range_max + 1))] 224 | + [10 ** (scale_exp)] 225 | ) 226 | cbar_yticks_labels = [ 227 | f"$\\frac{{1}}{{{10**abs(s):.0f}}}$" if s < 0 else f"${10**s:.0f}$" 228 | for s in scale_exp_int_range 229 | ] 230 | 231 | fig.subplots_adjust(right=0.8) 232 | # cbar_ax = fig.add_axes([0.85, 0.15, 0.01, 0.7]) 233 | # cbar = fig.colorbar(c, ticks=cbar_yticks_list, ax=cbar_ax) 234 | cbar = fig.colorbar(c, ticks=cbar_yticks_list, ax=ax.ravel().tolist()) 235 | cbar.ax.set_yticklabels(cbar_yticks_labels) 236 | cbar.ax.set_title("speedup") 237 | 238 | # plt.tight_layout() 239 | plt.savefig("l4casadi_vs_l4acados_speedup.pdf") 240 | # - 241 | 242 | device_threads_arr 243 | 244 | # + 245 | fig, ax = plt.subplots( 246 | 2, len(device_threads_arr), figsize=(7.1413, 3.5), sharex=True, sharey="row" 247 | ) 248 | # fig, ax = plt.subplots(2, len(device_threads_arr), figsize=(3*7.1413,3*3.5), sharex=True, sharey="row") 249 | 250 | # Timing plots 251 | 252 | colormap_exp = 1.0 253 | # fig, ax = plt.subplots(1, len(device_threads_arr), figsize=(1.5*7,1.5*2.5), sharey=True) 254 | ax_ylim_arr = np.zeros((2, len(device_threads_arr))) 255 | 256 | for l, device_threads in enumerate(device_threads_arr): 257 | for i, hidden_layers in enumerate(hidden_layers_arr): 258 | h_l4casadi = ax[0, l].plot( 259 | N_arr, 260 | avg_times["l4casadi"][i, :, l], 261 | color=cmap_l4casadi(((i + 1) / len(hidden_layers_arr)) ** colormap_exp), 262 | linestyle="--", 263 | ) 264 | h_l4acados = ax[0, l].plot( 265 | N_arr, 266 | avg_times["l4acados"][i, :, l], 267 | color=cmap_l4acados(((i + 1) / len(hidden_layers_arr)) ** colormap_exp), 268 | ) 269 | 270 | # ax[0,l].plot(N_arr, np.array(N_arr) * 5e-2, linestyle=":", color="gray") 271 | 272 | ax_ylim_arr[:, l] += np.array(ax[0, l].get_ylim()) 273 | ax[0, l].set_xscale("log") 274 | ax[0, l].set_yscale("log") 275 | # ax[0,l].set_xlabel("N") 276 | # ax[0,l].set_title(f"{device_threads[0]}-torch{device_threads[1]}-acados{device_threads[2]}") 277 | ax[0, l].set_title(f"{device_threads[0]}-t{device_threads[1]}-a{device_threads[2]}") 278 | ax[0, l].grid() 279 | 280 | ax_ylim_min = min(ax_ylim_arr[0, :]) 281 | ax_ylim_max = max(ax_ylim_arr[1, :]) 282 | for l in range(len(device_threads_arr)): 283 | # ax[0,l].set_ylim((ax_ylim_min, ax_ylim_max)) 284 | # ax[0,l].set_xlim((ax[0,l].get_xlim()[0], 10**(np.log10(ax[0,l].get_xlim()[0]) + np.log10(ax[0,l].get_ylim()[1]) - np.log10(ax[0,l].get_ylim()[0])))) 285 | ax[0, l].set_box_aspect(1.0) 286 | 287 | legend_handles = [mlines.Line2D([], [], color="k", linestyle="-", label="l4acados")] 288 | legend_handles += [mlines.Line2D([], [], color="k", linestyle="--", label="l4casadi")] 289 | legend_handles += [ 290 | mlines.Line2D([], [], color="tab:gray", linestyle="none", label="", marker="+") 291 | ] 292 | legend_handles += [ 293 | mlines.Line2D( 294 | [], 295 | [], 296 | color=cmap_l4acados(((i + 1) / len(hidden_layers_arr)) ** colormap_exp), 297 | marker="o", 298 | linestyle="none", 299 | label=f"$\#_{{\\mathrm{{layers}}}} = {hidden_layers_arr[i]:d}$", 300 | ) 301 | for i in range(len(hidden_layers_arr) - 1, -1, -1) 302 | ] 303 | 304 | ax[0, 0].set_ylabel("Time per SQP iteration [s]") 305 | 306 | # Color plots 307 | 308 | # scale_exp = 1.5 309 | scale_max = 50 310 | scale_exp = np.log10(scale_max) 311 | # width_ratios = np.ones((len(device_threads_arr,))) 312 | # width_ratios[-1] = 1.25 313 | # width_ratios_normalized = width_ratios / sum(width_ratios) 314 | # fig, ax = plt.subplots(1, len(device_threads_arr), figsize=(7,2.5), width_ratios=width_ratios_normalized, sharey=True) 315 | for l, device_threads in enumerate(device_threads_arr): 316 | c = ax[1, l].pcolor( 317 | N_grid, 318 | hidden_layers_grid, 319 | avg_times["l4casadi"][:, :, l] / avg_times["l4acados"][:, :, l], 320 | shading="nearest", 321 | norm=LogNorm(vmin=10 ** (-scale_exp), vmax=10**scale_exp), 322 | cmap=colormaps["BrBG"], 323 | ) 324 | ax[1, l].set_xlabel("N") 325 | ax[1, l].set_xticks(N_arr) 326 | ax[1, l].set_xscale("log") 327 | ax[1, l].set_xlim((min(N_arr), max(N_arr))) 328 | # ax[1,l].set_yscale("log") 329 | # ax[1,l].set_ylim((min(hidden_layers_arr),max(hidden_layers_arr))) 330 | ax[1, l].set_yticks(hidden_layers_arr) 331 | 332 | ax[1, l].set_box_aspect(1.0) 333 | 334 | ax[1, 0].set_ylabel("Hidden layers") 335 | 336 | scale_exp_int_range_min = int(np.ceil(-scale_exp)) 337 | scale_exp_int_range_max = int(np.floor(scale_exp)) 338 | scale_exp_int_range = ( 339 | [-scale_exp] 340 | + list(range(scale_exp_int_range_min, scale_exp_int_range_max + 1)) 341 | + [scale_exp] 342 | ) 343 | cbar_yticks_list = ( 344 | [10 ** (-scale_exp)] 345 | + [10**i for i in list(range(scale_exp_int_range_min, scale_exp_int_range_max + 1))] 346 | + [10 ** (scale_exp)] 347 | ) 348 | cbar_yticks_labels = [ 349 | f"$\\frac{{1}}{{{10**abs(s):.0f}}}$" if s < 0 else f"${10**s:.0f}$" 350 | for s in scale_exp_int_range 351 | ] 352 | 353 | # Legend 354 | ax[0, -1].legend(loc="center left", bbox_to_anchor=(1.05, 0.5), handles=legend_handles) 355 | 356 | # Colorbar 357 | ax_shift_y = 5 358 | ax_last_lim = ax[1, -1].get_ylim() 359 | cbar_ax = ax[1, -1].inset_axes( 360 | [N_arr[-1] * 3.0, ax_last_lim[0], N_arr[-1] * 3.0, ax_last_lim[1] - ax_last_lim[0]], 361 | transform=ax[1, -1].transData, 362 | ) 363 | cbar = fig.colorbar(c, ticks=cbar_yticks_list, ax=ax, cax=cbar_ax) 364 | cbar.ax.set_yticklabels(cbar_yticks_labels) 365 | cbar.set_label("Speedup") 366 | 367 | plt.tight_layout() 368 | 369 | fig.savefig("l4casadi_vs_l4acados_all.pdf", dpi=600) 370 | -------------------------------------------------------------------------------- /examples/l4casadi_vs_l4acados/build_acados_num_threads.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | ACADOS_NUM_THREADS=$1 3 | ACADOS_PYTHON=ON 4 | ACADOS_WITH_OPENMP=${2:-ON} 5 | 6 | mkdir -p $ACADOS_SOURCE_DIR/build 7 | cd $ACADOS_SOURCE_DIR/build 8 | 9 | cmake -DACADOS_PYTHON=$ACADOS_PYTHON -DACADOS_WITH_OPENMP=$ACADOS_WITH_OPENMP -DACADOS_NUM_THREADS=$ACADOS_NUM_THREADS .. 10 | make install -j4 11 | -------------------------------------------------------------------------------- /examples/l4casadi_vs_l4acados/data/.gitkeep: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/IntelligentControlSystems/l4acados/a56ec0184c52d832f6801ffe7e39c463d6fa1d37/examples/l4casadi_vs_l4acados/data/.gitkeep -------------------------------------------------------------------------------- /examples/l4casadi_vs_l4acados/l4casadi_vs_l4acados.py: -------------------------------------------------------------------------------- 1 | # --- 2 | # jupyter: 3 | # jupytext: 4 | # text_representation: 5 | # extension: .py 6 | # format_name: light 7 | # format_version: '1.5' 8 | # jupytext_version: 1.16.4 9 | # kernelspec: 10 | # display_name: zero-order-gpmpc-package-3.10 11 | # language: python 12 | # name: python3 13 | # --- 14 | 15 | # + 16 | import acados 17 | import importlib 18 | 19 | importlib.reload(acados) 20 | # - 21 | 22 | from acados import run, MultiLayerPerceptron, DoubleIntegratorWithLearnedDynamics, MPC 23 | import l4casadi as l4c 24 | import numpy as np 25 | import time 26 | import l4acados as l4a 27 | from typing import Optional, Union 28 | import torch 29 | import casadi as cs 30 | 31 | from l4acados.controllers.residual_learning_mpc import ResidualLearningMPC 32 | from l4acados.models import ResidualModel, PyTorchFeatureSelector 33 | from l4acados.controllers.zoro_acados_utils import setup_sim_from_ocp 34 | 35 | import copy 36 | 37 | from run_single_experiment import * 38 | 39 | N = 20 40 | ts = 1.0 / N 41 | batch_dim = 1 42 | hidden_layers = 5 43 | warmup_iter = 100 44 | solve_steps = 1000 45 | num_threads = 1 46 | # device = "cpu" 47 | device = "cuda" 48 | num_threads_acados_openmp = 4 49 | 50 | x_l4casadi, opt_times_l4casadi, x_l4acados, opt_times_l4acados = run( 51 | N, 52 | hidden_layers, 53 | solve_steps, 54 | device=device, 55 | num_threads_acados_openmp=num_threads_acados_openmp, 56 | ) 57 | 58 | import matplotlib.pyplot as plt 59 | 60 | opt_times_l4casadi_avg = np.cumsum(opt_times_l4casadi[warmup_iter:]) / np.arange( 61 | 1, len(opt_times_l4casadi[warmup_iter:]) + 1 62 | ) 63 | opt_times_l4acados_avg = np.cumsum(opt_times_l4acados[warmup_iter:]) / np.arange( 64 | 1, len(opt_times_l4acados[warmup_iter:]) + 1 65 | ) 66 | 67 | h_l4casadi = plt.plot( 68 | np.arange(warmup_iter, len(opt_times_l4casadi)), 69 | opt_times_l4casadi_avg, 70 | label="l4casadi", 71 | ) 72 | h_l4acados = plt.plot( 73 | np.arange(warmup_iter, len(opt_times_l4casadi)), 74 | opt_times_l4acados_avg, 75 | label="l4acados", 76 | ) 77 | plt.plot( 78 | opt_times_l4casadi, 79 | label="l4casadi", 80 | color=h_l4casadi[0].get_color(), 81 | alpha=0.3, 82 | ) 83 | plt.plot( 84 | opt_times_l4acados, 85 | label="l4acados", 86 | color=h_l4acados[0].get_color(), 87 | alpha=0.3, 88 | ) 89 | plt.axvline( 90 | x=warmup_iter, color="k", linestyle="--", linewidth=1, label="warmup", alpha=0.5 91 | ) 92 | # axes scaling log 93 | # plt.xscale("log") 94 | plt.yscale("log") 95 | # axes title y 96 | plt.ylabel("Time [s]") 97 | plt.xlabel("Iteration") 98 | plt.legend() 99 | # plt.ylim([1e-3, 1e-1]) 100 | plt.grid() 101 | 102 | opt_times_l4casadi_avg[-1], opt_times_l4acados_avg[-1], opt_times_l4casadi_avg[ 103 | -1 104 | ] / opt_times_l4acados_avg[-1] 105 | 106 | plt.plot(x_l4casadi, linewidth=3) 107 | plt.plot(x_l4acados, linewidth=1) 108 | 109 | np.linalg.norm(np.array(x_l4casadi) - np.array(x_l4acados)) 110 | 111 | plt.plot(np.array(x_l4acados) - np.array(x_l4casadi)) 112 | -------------------------------------------------------------------------------- /examples/l4casadi_vs_l4acados/run_all_experiments.py: -------------------------------------------------------------------------------- 1 | from run_single_experiment import * 2 | import subprocess 3 | 4 | N_arr = [int(i) for i in np.ceil(np.logspace(0, 3, 10))] 5 | batch_dim = 1 6 | hidden_layers_arr = [1, 10, 20] 7 | solve_steps = 1000 8 | device_arr = ["cpu", "cpu", "cuda", "cpu", "cuda"] 9 | num_threads_arr = [1, 10, 1, 1, 1] 10 | num_threads_acados_openmp_arr = [0, 0, 0, 10, 10] 11 | save_data = True 12 | 13 | assert len(num_threads_arr) == len(device_arr) == len(num_threads_acados_openmp_arr) 14 | device_threads_arr = list( 15 | zip(device_arr, num_threads_arr, num_threads_acados_openmp_arr) 16 | ) 17 | 18 | print( 19 | f"Running experiments with\nN={N_arr}\nhidden_layers={hidden_layers_arr}\ndevices={device_arr}\nnum_threads_torch={num_threads_arr}\nnum_threads_acados={num_threads_acados_openmp_arr}" 20 | ) 21 | print( 22 | f"Total number of experiments: {len(N_arr)*len(hidden_layers_arr)*len(device_threads_arr)}" 23 | ) 24 | 25 | num_threads_acados_openmp_previous = -1 26 | for device, num_threads, num_threads_acados_openmp in device_threads_arr: 27 | 28 | build_acados = False 29 | if not num_threads_acados_openmp == num_threads_acados_openmp_previous: 30 | build_acados = True 31 | 32 | num_threads_acados_openmp_previous = num_threads_acados_openmp 33 | 34 | for i, N in enumerate(N_arr): 35 | for hidden_layers in hidden_layers_arr: 36 | print( 37 | f"Calling subprocess with N={N}, hidden_layers={hidden_layers}, device={device}, num_threads={num_threads}, num_threads_acados_openmp={num_threads_acados_openmp}" 38 | ) 39 | 40 | if build_acados: 41 | build_acados_arg = "--build_acados" 42 | else: 43 | build_acados_arg = "--no-build_acados" 44 | 45 | build_acados = False 46 | 47 | subprocess_call_list = [ 48 | "python", 49 | "run_single_experiment.py", 50 | "--N", 51 | str(N), 52 | "--hidden_layers", 53 | str(hidden_layers), 54 | "--device", 55 | str(device), 56 | "--num_threads", 57 | str(num_threads), 58 | "--num_threads_acados_openmp", 59 | str(num_threads_acados_openmp), 60 | "--solve_steps", 61 | str(solve_steps), 62 | build_acados_arg, 63 | ] 64 | subprocess_call_str = " ".join(subprocess_call_list) 65 | 66 | print(f"Start subprocess: {subprocess_call_str}") 67 | subprocess.check_call(subprocess_call_list) 68 | print(f"Finished subprocess: {subprocess_call_str}") 69 | -------------------------------------------------------------------------------- /examples/l4casadi_vs_l4acados/run_single_experiment.py: -------------------------------------------------------------------------------- 1 | from acados import run, MultiLayerPerceptron, DoubleIntegratorWithLearnedDynamics, MPC 2 | import l4casadi as l4c 3 | import numpy as np 4 | import time 5 | import l4acados as l4a 6 | from typing import Optional, Union 7 | import torch 8 | import casadi as cs 9 | from l4acados.controllers import ResidualLearningMPC 10 | from l4acados.models import ResidualModel, PyTorchResidualModel, PyTorchFeatureSelector 11 | from l4acados.controllers.zoro_acados_utils import setup_sim_from_ocp 12 | import argparse 13 | import os, shutil, re 14 | import subprocess 15 | 16 | 17 | class PyTorchResidualModel(ResidualModel): 18 | """Basic PyTorch residual model class. 19 | 20 | Args: 21 | - model: Any Pytorch torch.nn.Module. 22 | - feature_selector: Optional feature selector mapping (state, input) -> (NN input dimension). If set to None, then no selection 23 | is performed. 24 | """ 25 | 26 | def __init__( 27 | self, 28 | model: torch.nn.Module, 29 | feature_selector: Optional[PyTorchFeatureSelector] = None, 30 | device="cpu", 31 | ): 32 | self.model = model 33 | 34 | self._gp_feature_selector = ( 35 | feature_selector 36 | if feature_selector is not None 37 | else PyTorchFeatureSelector() 38 | ) 39 | 40 | if device == "cuda": 41 | self.to_tensor = lambda X: torch.Tensor(X).cuda() 42 | self.to_numpy = lambda T: T.cpu().detach().numpy() 43 | self.model.cuda() 44 | elif device == "cpu": 45 | self.to_tensor = lambda X: torch.Tensor(X) 46 | self.to_numpy = lambda T: T.detach().numpy() 47 | self.model.cpu() 48 | else: 49 | raise ValueError(f"Unknown device {device}, should be 'cpu' or 'gpu'") 50 | 51 | def _mean_fun_sum(self, y): 52 | """Helper function for jacobian computation 53 | 54 | sums up the mean predictions along the first dimension 55 | (i.e. along the horizon). 56 | """ 57 | self.evaluate(y, require_grad=True) 58 | return self.predictions.sum(dim=0) 59 | 60 | def evaluate(self, y, require_grad=False): 61 | y_tensor = self.to_tensor(y) 62 | if require_grad: 63 | self.predictions = self.model(self._gp_feature_selector(y_tensor)) 64 | else: 65 | with torch.no_grad(): 66 | self.predictions = self.model(self._gp_feature_selector(y_tensor)) 67 | 68 | self.current_mean = self.to_numpy(self.predictions) 69 | return self.current_mean 70 | 71 | def jacobian(self, y): 72 | y_tensor = self.to_tensor(y) 73 | mean_dy = torch.autograd.functional.jacobian(self._mean_fun_sum, y_tensor) 74 | return self.to_numpy(mean_dy) 75 | 76 | def value_and_jacobian(self, y): 77 | """Computes the necessary values for GPMPC 78 | 79 | Args: 80 | - x_input: (N, state_dim) tensor 81 | 82 | Returns: 83 | - mean: (N, residual_dim) tensor 84 | - mean_dy: (residual_dim, N, state_dim) tensor 85 | - covariance: (N, residual_dim) tensor 86 | """ 87 | self.current_mean_dy = self.jacobian(y) 88 | 89 | return self.current_mean, self.current_mean_dy 90 | 91 | 92 | def init_l4casadi( 93 | N: int, 94 | hidden_layers: int, 95 | batch_dim: int = 1, 96 | batched: bool = True, 97 | device="cpu", 98 | num_threads_acados_openmp=1, 99 | ): 100 | learned_dyn_model = l4c.L4CasADi( 101 | MultiLayerPerceptron(hidden_layers=hidden_layers), 102 | batched=batched, 103 | name="learned_dyn", 104 | device=device, 105 | ) 106 | 107 | model = DoubleIntegratorWithLearnedDynamics( 108 | learned_dyn_model, batched=batched, batch_dim=batch_dim 109 | ) 110 | 111 | mpc_obj = MPC( 112 | model=model.model(), 113 | N=N, 114 | external_shared_lib_dir=learned_dyn_model.shared_lib_dir, 115 | external_shared_lib_name=learned_dyn_model.name, 116 | num_threads_acados_openmp=num_threads_acados_openmp, 117 | ) 118 | solver = mpc_obj.solver 119 | 120 | if num_threads_acados_openmp > 1: 121 | assert ( 122 | solver.acados_lib_uses_omp == True 123 | ), "Acados not compiled with OpenMP, cannot use multiple threads." 124 | 125 | ocp = mpc_obj.ocp() 126 | return solver 127 | 128 | 129 | def init_l4acados( 130 | N: int, 131 | hidden_layers: int, 132 | batch_dim: int = 1, 133 | batched: bool = True, 134 | device="cpu", 135 | use_cython=False, 136 | num_threads_acados_openmp=1, 137 | ): 138 | feature_selector = PyTorchFeatureSelector([1, 1, 0], device=device) 139 | residual_model = PyTorchResidualModel( 140 | MultiLayerPerceptron(hidden_layers=hidden_layers), 141 | feature_selector, 142 | device=device, 143 | ) 144 | B_proj = np.ones((1, batch_dim)) 145 | model_new = DoubleIntegratorWithLearnedDynamics(None, name="wr_new") 146 | mpc_obj_nolib = MPC( 147 | model=model_new.model(), 148 | N=N, 149 | num_threads_acados_openmp=num_threads_acados_openmp, 150 | ) 151 | solver_nolib = mpc_obj_nolib.solver 152 | ocp_nolib = mpc_obj_nolib.ocp() 153 | sim_nolib = setup_sim_from_ocp(ocp_nolib) 154 | 155 | solver_l4acados = ResidualLearningMPC( 156 | ocp=ocp_nolib, 157 | B=B_proj, 158 | residual_model=residual_model, 159 | use_cython=use_cython, 160 | ) 161 | 162 | return solver_l4acados 163 | 164 | 165 | def run_timing_experiment(N, solver, solve_call, solve_steps=1e3): 166 | x = [] 167 | u = [] 168 | xt = np.array([1.0, 0.0]) 169 | T = 1.0 170 | ts = T / N 171 | opt_times = [] 172 | 173 | for i in range(solve_steps): 174 | # t = np.linspace(i * ts, (i + 1) * ts, N) # TODO: this ts should be T IMO, maybe with lower frequency? 175 | t = np.linspace( 176 | i * T, (i + 1) * T, N 177 | ) # TODO: this ts should be T IMO, maybe with lower frequency? 178 | yref = np.sin(0.1 * t + np.pi / 2) 179 | for t, ref in enumerate(yref): 180 | solver.set(t, "yref", np.array([ref])) 181 | solver.set(0, "lbx", xt) 182 | solver.set(0, "ubx", xt) 183 | 184 | # now = time.time() 185 | elapsed = solve_call() 186 | 187 | xt = solver.get(1, "x") 188 | 189 | opt_times.append(elapsed) 190 | x.append(xt) 191 | 192 | print(f"Running timing experiment: {i}/{solve_steps}") 193 | 194 | return x, opt_times 195 | 196 | 197 | def time_fun_call(fun): 198 | now = time.perf_counter() 199 | fun() 200 | return time.perf_counter() - now 201 | 202 | 203 | def delete_file_by_pattern(dir_path, pattern): 204 | for f in os.listdir(dir_path): 205 | if re.search(pattern, f): 206 | os.remove(os.path.join(dir_path, f)) 207 | 208 | 209 | def run( 210 | N, 211 | hidden_layers, 212 | solve_steps, 213 | device="cpu", 214 | save_data=False, 215 | save_dir="data", 216 | num_threads: int = -1, 217 | num_threads_acados_openmp: int = 1, 218 | build_acados: bool = True, 219 | ): 220 | 221 | if build_acados: 222 | if num_threads_acados_openmp >= 1: 223 | build_acados_with_openmp = "ON" 224 | else: 225 | build_acados_with_openmp = "OFF" 226 | 227 | subprocess.check_call( 228 | [ 229 | os.path.join( 230 | os.path.dirname(os.path.abspath(__file__)), 231 | "build_acados_num_threads.sh", 232 | ), 233 | str(num_threads_acados_openmp), 234 | build_acados_with_openmp, 235 | ] 236 | ) 237 | 238 | if device == "cuda": 239 | num_threads = 1 240 | elif num_threads == -1: 241 | num_threads = os.cpu_count() // 2 242 | torch.set_num_threads(num_threads) 243 | 244 | solver_l4casadi = init_l4casadi( 245 | N, 246 | hidden_layers, 247 | device=device, 248 | num_threads_acados_openmp=num_threads_acados_openmp, 249 | ) 250 | x_l4casadi, opt_times_l4casadi = run_timing_experiment( 251 | N, 252 | solver_l4casadi, 253 | lambda: time_fun_call(solver_l4casadi.solve), 254 | solve_steps=solve_steps, 255 | ) 256 | 257 | shutil.rmtree("c_generated_code") 258 | shutil.rmtree("_l4c_generated") 259 | delete_file_by_pattern("./", r".*[ocp|sim].*\.json") 260 | 261 | solver_l4acados = init_l4acados( 262 | N, 263 | hidden_layers, 264 | device=device, 265 | use_cython=True, 266 | num_threads_acados_openmp=num_threads_acados_openmp, 267 | ) 268 | x_l4acados, opt_times_l4acados = run_timing_experiment( 269 | N, 270 | solver_l4acados.ocp_solver, 271 | lambda: time_fun_call(lambda: solver_l4acados.solve()), 272 | solve_steps=solve_steps, 273 | ) 274 | 275 | shutil.rmtree("c_generated_code") 276 | delete_file_by_pattern("./", r".*[ocp|sim].*\.json") 277 | 278 | if save_data: 279 | # save data 280 | print("Saving data") 281 | np.savez( 282 | os.path.join( 283 | save_dir, 284 | f"l4casadi_vs_l4acados_N{N}_layers{hidden_layers}_steps{solve_steps}_{device}_threads{num_threads}_acados_{num_threads_acados_openmp}.npz", 285 | ), 286 | x_l4casadi=x_l4casadi, 287 | opt_times_l4casadi=opt_times_l4casadi, 288 | x_l4acados=x_l4acados, 289 | opt_times_l4acados=opt_times_l4acados, 290 | ) 291 | 292 | del solver_l4casadi, solver_l4acados 293 | 294 | return x_l4casadi, opt_times_l4casadi, x_l4acados, opt_times_l4acados 295 | 296 | 297 | if __name__ == "__main__": 298 | # parse arguments with argparse 299 | parser = argparse.ArgumentParser() 300 | parser.add_argument("--N", type=int, default=10) 301 | parser.add_argument("--hidden_layers", type=int, default=1) 302 | parser.add_argument("--solve_steps", type=int, default=1000) 303 | parser.add_argument("--num_threads", type=int, default=-1) 304 | parser.add_argument("--num_threads_acados_openmp", type=int, default=1) 305 | parser.add_argument("--device", type=str, default="cpu") 306 | parser.add_argument( 307 | "--build_acados", action=argparse.BooleanOptionalAction, default=True 308 | ) 309 | args = parser.parse_args() 310 | args_dict = vars(args) 311 | 312 | print_str = "Run_single_experiment: " 313 | for arg_name, arg_value in args_dict.items(): 314 | print_str += f"{arg_name}={arg_value}, " 315 | print(f"{print_str}\n") 316 | 317 | run( 318 | args.N, 319 | args.hidden_layers, 320 | args.solve_steps, 321 | device=args.device, 322 | num_threads=args.num_threads, 323 | num_threads_acados_openmp=args.num_threads_acados_openmp, 324 | build_acados=args.build_acados, 325 | save_data=True, 326 | ) 327 | 328 | # python run_single_experiment.py --N 10 --hidden_layers 20 --solve_steps 100 --num_threads 1 --num_threads_acados_openmp 14 --device cpu --(no-)build_acados 329 | -------------------------------------------------------------------------------- /external/gpytorch_utils/gp_hyperparam_training.py: -------------------------------------------------------------------------------- 1 | import sys, os 2 | import numpy as np 3 | import gpytorch 4 | import torch 5 | import random 6 | 7 | 8 | def generate_train_inputs_zoro( 9 | zoro_solver, x0_nom, N_sim_per_x0, N_x0, random_seed=None, x0_rand_scale=0.1 10 | ): 11 | if random_seed is not None: 12 | np.random.seed(random_seed) 13 | 14 | nx = zoro_solver.nx 15 | nu = zoro_solver.nu 16 | N = zoro_solver.N 17 | 18 | x0_arr = np.zeros((N_x0, nx)) 19 | X_inp = np.zeros((N_x0 * N_sim_per_x0 * N, nx + nu)) 20 | 21 | i_fac = N_sim_per_x0 * N 22 | j_fac = N 23 | for i in range(N_x0): 24 | x0_arr[i, :] = x0_nom + x0_rand_scale * (2 * np.random.rand(nx) - 1) 25 | 26 | zoro_solver.ocp_solver.set(0, "lbx", x0_arr[i, :]) 27 | zoro_solver.ocp_solver.set(0, "ubx", x0_arr[i, :]) 28 | zoro_solver.solve() 29 | 30 | X, U, P = zoro_solver.get_solution() 31 | 32 | # store training points 33 | for j in range(N_sim_per_x0): 34 | for k in range(N): 35 | ijk = i * i_fac + j * j_fac + k 36 | X_inp[ijk, :] = np.hstack((X[k, :], U[k, :])) 37 | 38 | return X_inp, x0_arr 39 | 40 | 41 | def generate_train_inputs_acados( 42 | ocp_solver, x0_nom, N_sim_per_x0, N_x0, random_seed=None, x0_rand_scale=0.1 43 | ): 44 | if random_seed is not None: 45 | np.random.seed(random_seed) 46 | 47 | nx = ocp_solver.acados_ocp.dims.nx 48 | nu = ocp_solver.acados_ocp.dims.nu 49 | N = ocp_solver.acados_ocp.dims.N 50 | 51 | x0_arr = np.zeros((N_x0, nx)) 52 | X_inp = np.zeros((N_x0 * N_sim_per_x0 * N, nx + nu)) 53 | 54 | i_fac = N_sim_per_x0 * N 55 | j_fac = N 56 | for i in range(N_x0): 57 | x0_arr[i, :] = x0_nom + x0_rand_scale * (2 * np.random.rand(nx) - 1) 58 | 59 | ocp_solver.set(0, "lbx", x0_arr[i, :]) 60 | ocp_solver.set(0, "ubx", x0_arr[i, :]) 61 | ocp_solver.solve() 62 | 63 | # store training points 64 | for j in range(N_sim_per_x0): 65 | for k in range(N): 66 | ijk = i * i_fac + j * j_fac + k 67 | X_inp[ijk, :] = np.hstack( 68 | (ocp_solver.get(k, "x"), ocp_solver.get(k, "u")) 69 | ) 70 | 71 | return X_inp, x0_arr 72 | 73 | 74 | def generate_train_data_acados( 75 | acados_ocp_solver, 76 | integrator_nom, 77 | integrator_sim, 78 | x0_nom, 79 | Sigma_W, 80 | N_sim_per_x0, 81 | N_sim, 82 | B=None, 83 | N_x0=1, 84 | random_seed=None, 85 | x0_rand_scale=0.1, 86 | ): 87 | if random_seed is not None: 88 | np.random.seed(random_seed) 89 | 90 | if B is None: 91 | B = np.eye(nw) 92 | B_inv = np.linalg.pinv(B) 93 | 94 | nx = acados_ocp_solver.acados_ocp.model.x.size()[0] 95 | nu = acados_ocp_solver.acados_ocp.model.u.size()[0] 96 | nw = Sigma_W.shape[0] 97 | 98 | B_inv = np.linalg.pinv(B) 99 | x0_arr = np.zeros((N_x0, nx)) 100 | X_inp = np.zeros((N_x0 * N_sim_per_x0 * N_sim, nx + nu)) 101 | Y_out = np.zeros((N_x0 * N_sim_per_x0 * N_sim, nw)) 102 | 103 | i_fac = N_sim_per_x0 * N_sim 104 | j_fac = N_sim 105 | for i in range(N_x0): 106 | ijk = i * i_fac 107 | x0_arr[i, :] = x0_nom + x0_rand_scale * (2 * np.random.rand(nx) - 1) 108 | xcurrent = x0_arr[i, :] 109 | for j in range(N_sim_per_x0): 110 | for k in range(N_sim): 111 | acados_ocp_solver.set(0, "lbx", xcurrent) 112 | acados_ocp_solver.set(0, "ubx", xcurrent) 113 | acados_ocp_solver.solve() 114 | 115 | u = acados_ocp_solver.get(0, "u") 116 | 117 | # integrate nominal model 118 | integrator_nom.set("x", xcurrent) 119 | integrator_nom.set("u", u) 120 | integrator_nom.solve() 121 | xnom = integrator_nom.get("x") 122 | 123 | # integrate real model 124 | integrator_sim.set("x", xcurrent) 125 | integrator_sim.set("u", u) 126 | integrator_sim.solve() 127 | xcurrent = integrator_sim.get("x") 128 | 129 | # difference 130 | w = np.random.multivariate_normal(np.zeros((nw,)), Sigma_W) 131 | 132 | # store training points 133 | ijk = i * i_fac + j * j_fac + k 134 | X_inp[ijk, :] = np.hstack((xcurrent, u)) 135 | 136 | Y_out[ijk, :] = B_inv @ (xcurrent - xnom) + w 137 | 138 | return X_inp, Y_out 139 | 140 | 141 | def generate_train_outputs_at_inputs( 142 | X_inp, integrator_nom, integrator_sim, Sigma_W, B=None 143 | ): 144 | 145 | nx = integrator_nom.acados_sim.model.x.size()[0] 146 | nu = integrator_nom.acados_sim.model.u.size()[0] 147 | nw = Sigma_W.shape[0] 148 | 149 | if B is None: 150 | B = np.eye(nw) 151 | B_inv = np.linalg.pinv(B) 152 | 153 | n_train = X_inp.shape[0] 154 | Y_out = np.zeros((n_train, nw)) 155 | for i in range(n_train): 156 | integrator_sim.set("x", X_inp[i, 0:nx]) 157 | integrator_sim.set("u", X_inp[i, nx : nx + nu]) 158 | integrator_sim.solve() 159 | 160 | integrator_nom.set("x", X_inp[i, 0:nx]) 161 | integrator_nom.set("u", X_inp[i, nx : nx + nu]) 162 | integrator_nom.solve() 163 | 164 | w = np.random.multivariate_normal(np.zeros((nw,)), Sigma_W) 165 | 166 | Y_out[i, :] = B_inv @ (integrator_sim.get("x") - integrator_nom.get("x")) + w 167 | return Y_out 168 | 169 | 170 | def train_gp_model( 171 | gp_model, torch_seed=None, training_iterations=200, learning_rate=0.1 172 | ): 173 | if torch_seed is not None: 174 | torch.manual_seed(torch_seed) 175 | 176 | likelihood = gp_model.likelihood 177 | train_x = gp_model.train_inputs[0] 178 | train_y = gp_model.train_targets 179 | 180 | # Find optimal model hyperparameters 181 | gp_model.train() 182 | likelihood.train() 183 | 184 | # Use the adam optimizer 185 | optimizer = torch.optim.Adam( 186 | filter(lambda p: p.requires_grad, gp_model.parameters()), lr=learning_rate 187 | ) # Includes GaussianLikelihood parameters 188 | 189 | # "Loss" for GPs - the marginal log likelihood 190 | mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp_model) 191 | 192 | prev_loss = np.inf 193 | num_loss_below_threshold = 0 194 | 195 | for i in range(training_iterations): 196 | optimizer.zero_grad() 197 | output = likelihood(gp_model(train_x)) 198 | loss = -mll(output, train_y.reshape((train_y.numel(),))) 199 | 200 | num_loss_below_threshold = ( 201 | num_loss_below_threshold + 1 if abs(loss - prev_loss) < 5e-4 else 0 202 | ) 203 | 204 | if num_loss_below_threshold > 5: 205 | print(f"stopping GP optimization early after {i} iterations.") 206 | break 207 | 208 | prev_loss = loss 209 | 210 | loss.backward() 211 | if (i + 1) % 20 == 0: 212 | print("Iter %d/%d - Loss: %.3f" % (i + 1, training_iterations, loss.item())) 213 | 214 | optimizer.step() 215 | 216 | gp_model.eval() 217 | likelihood.eval() 218 | return gp_model, likelihood 219 | 220 | 221 | def set_gp_param_value(gp_model, name: str, value: torch.Tensor): 222 | constraint = gp_model.constraint_for_parameter_name(name) 223 | parameter = gp_model.get_parameter(name) 224 | state_dict = gp_model.state_dict() 225 | 226 | # transform value 227 | if constraint is not None: 228 | value_transform = constraint.inverse_transform(value) 229 | else: 230 | value_transform = value 231 | 232 | state_dict[name] = value_transform 233 | gp_model.load_state_dict(state_dict) 234 | 235 | return parameter 236 | 237 | 238 | def get_gp_param_value(gp_model, name: str): 239 | constraint = gp_model.constraint_for_parameter_name(name) 240 | parameter = gp_model.get_parameter(name) 241 | state_dict = gp_model.state_dict() 242 | value_transform = state_dict[name] 243 | 244 | # transform value 245 | if constraint is not None: 246 | value = constraint.transform(value_transform) 247 | else: 248 | value = value_transform 249 | return value 250 | 251 | 252 | def get_gp_param_names(gp_model): 253 | names = [] 254 | for name, parameter in gp_model.named_parameters(): 255 | names += [name] 256 | return names 257 | 258 | 259 | def get_gp_param_names_values(gp_model): 260 | names = get_gp_param_names(gp_model) 261 | values = [] 262 | for name in names: 263 | values += [get_gp_param_value(gp_model, name)] 264 | return zip(names, values) 265 | 266 | 267 | def get_prior_covariance(gp_model): 268 | 269 | dim = gp_model.train_inputs[0].shape[1] 270 | 271 | # cuda check 272 | if gp_model.train_inputs[0].device.type == "cuda": 273 | to_numpy = lambda T: T.cpu().numpy() 274 | y_test = torch.Tensor(np.ones((1, dim))).cuda() 275 | else: 276 | to_numpy = lambda T: T.numpy() 277 | y_test = torch.Tensor(np.ones((1, dim))) 278 | 279 | gp_model.eval() 280 | # This only works for batched shape now 281 | # TODO: Make this work again for general multitask kernel 282 | prior_covar = np.diag(to_numpy(gp_model.covar_module(y_test)).squeeze()) 283 | 284 | return prior_covar 285 | -------------------------------------------------------------------------------- /external/gpytorch_utils/gp_utils.py: -------------------------------------------------------------------------------- 1 | import sys, os 2 | import numpy as np 3 | import gpytorch 4 | import torch 5 | from matplotlib import pyplot as plt 6 | from dataclasses import dataclass 7 | from typing import List 8 | 9 | 10 | @dataclass 11 | class GPlotData: 12 | x_path: np.ndarray = (None,) 13 | mean_on_path: np.ndarray = (None,) 14 | stddev_on_path: np.ndarray = (None,) 15 | conf_lower_on_path: np.ndarray = (None,) 16 | conf_upper_on_path: np.ndarray = (None,) 17 | sample_on_path: np.ndarray = (None,) 18 | train_x: np.ndarray = (None,) 19 | train_y: np.ndarray = None 20 | 21 | 22 | def gp_data_from_model_and_path( 23 | gp_model, likelihood, x_path, num_samples=0, use_likelihood=False 24 | ): 25 | # GP data 26 | if gp_model.train_inputs[0].device.type == "cuda": 27 | # train_x = gp_model.train_inputs[0].cpu().numpy() 28 | # train_y = gp_model.train_targets.cpu().numpy() 29 | x_path_tensor = torch.Tensor(x_path).cuda() 30 | to_numpy = lambda T: T.cpu().numpy() 31 | else: 32 | # train_x = gp_model.train_inputs[0].numpy() 33 | # train_y = gp_model.train_targets.numpy() 34 | x_path_tensor = torch.Tensor(x_path) 35 | to_numpy = lambda T: T.numpy() 36 | 37 | train_x = to_numpy(gp_model.train_inputs[0]) 38 | train_y = to_numpy(gp_model.train_targets) 39 | 40 | # dimensions 41 | num_points = x_path.shape[0] 42 | nout = train_y.shape[1] 43 | 44 | # Make predictions 45 | with torch.no_grad(), gpytorch.settings.fast_pred_var(): 46 | if use_likelihood: 47 | predictions = likelihood(gp_model(x_path_tensor)) # predictions with noise 48 | else: 49 | predictions = gp_model( 50 | x_path_tensor 51 | ) # only model (we want to find true function) 52 | 53 | mean_on_path = predictions.mean 54 | stddev_on_path = predictions.stddev 55 | conf_lower_on_path, conf_upper_on_path = predictions.confidence_region() 56 | 57 | # reshape 58 | mean_on_path = to_numpy(mean_on_path).reshape((num_points, nout)) 59 | stddev_on_path = to_numpy(stddev_on_path).reshape((num_points, nout)) 60 | conf_lower_on_path = to_numpy(conf_lower_on_path).reshape((num_points, nout)) 61 | conf_upper_on_path = to_numpy(conf_upper_on_path).reshape((num_points, nout)) 62 | 63 | # draw function samples 64 | sample_on_path = np.zeros((num_points, nout, num_samples)) 65 | for j in range(0, num_samples): 66 | sample_on_path[:, :, j] = to_numpy(predictions.sample()).reshape( 67 | (num_points, nout) 68 | ) 69 | 70 | return GPlotData( 71 | x_path, 72 | mean_on_path, 73 | stddev_on_path, 74 | conf_lower_on_path, 75 | conf_upper_on_path, 76 | sample_on_path, 77 | train_x, 78 | train_y, 79 | ) 80 | 81 | 82 | def gp_derivative_data_from_model_and_path(gp_model, likelihood, x_path, num_samples=0): 83 | # GP data 84 | train_x = gp_model.train_inputs[0].numpy() 85 | train_y = gp_model.train_targets.numpy() 86 | 87 | # dimensions 88 | num_points = x_path.shape[0] 89 | nout = train_y.shape[1] 90 | 91 | x_path_tensor = torch.Tensor(x_path) 92 | # Make predictions 93 | with gpytorch.settings.fast_pred_var(): 94 | # predictions = likelihood(gp_model(test_x)) 95 | # predictions = gp_model(x_path_tensor) # only model (we want to find true function) 96 | # mean_on_path = predictions.mean.detach().numpy() 97 | # stddev_on_path = predictions.stddev.detach().numpy() 98 | # conf_lower_on_path, conf_upper_on_path = predictions.confidence_region().detach().numpy() 99 | 100 | # DERIVATIVE 101 | mean_dx = torch.autograd.functional.jacobian( 102 | lambda x: gp_model(x).mean.sum(dim=0), x_path_tensor 103 | ) 104 | 105 | # project derivative along path 106 | x_path_diff = x_path[1:, :] - x_path[:-1, :] 107 | x_path_norm = x_path_diff / np.tile( 108 | np.sqrt(np.sum(x_path_diff**2, axis=1)).reshape((num_points - 1, 1)), 109 | (1, x_path.shape[1]), 110 | ) 111 | mean_dx_on_path = np.array( 112 | [ 113 | np.inner(mean_dx[:, i, :], x_path_norm[i, :]) 114 | for i in range(num_points - 1) 115 | ] 116 | ) 117 | 118 | # kernel derivative 119 | # k = gp_model.covar_module 120 | # kernel_dx_left_at_train = torch.autograd.functional.jacobian( 121 | # lambda x: k(x,train_x).sum(dim=0).unsqueeze(0), 122 | # x_path_tensor 123 | # ) 124 | # kernel_dx_right_at_train = torch.autograd.functional.jacobian( 125 | # lambda x: k(train_x,x).sum(dim=1).unsqueeze(1), 126 | # x_path_tensor 127 | # ) 128 | # kernel_dx_left_at_eval = torch.autograd.functional.jacobian( 129 | # lambda x: k(x,train_x).sum(dim=0).unsqueeze(0), 130 | # train_x 131 | # ) 132 | # kernel_ddx_at_eval = torch.autograd.functional.jacobian( 133 | # lambda x: kernel_dx_left_at_eval(x,train_x).sum(dim=0).unsqueeze(0), 134 | # train_x 135 | # ) 136 | 137 | # draw function samples 138 | # sample_on_path = np.zeros((num_points,nout,num_samples)) 139 | # for j in range(0,num_samples): 140 | # sample_on_path[:,:,j] = predictions.sample() 141 | 142 | return GPlotData( 143 | x_path[1:, :], mean_dx_on_path, None, None, None, None, train_x, train_y 144 | ) 145 | 146 | 147 | def plot_gp_data( 148 | gp_data_list: List[GPlotData], 149 | marker_size_lim=[5, 100], 150 | marker_size_reldiff_zero=1e-3, 151 | marker_style="x", 152 | plot_train_data=True, 153 | x_path_mode="shared", 154 | ): 155 | # color_list 156 | cmap = plt.cm.tab10 # define the colormap 157 | # extract all colors from the .jet map 158 | color_list = [cmap(i) for i in range(cmap.N)] 159 | 160 | # Initialize plots 161 | nout = gp_data_list[0].train_y.shape[1] 162 | fig, ax = plt.subplots(nout, 1, figsize=(8, 3 * nout)) 163 | if nout == 1: 164 | ax = [ax] 165 | 166 | if x_path_mode == "sequential": 167 | x_plot_segments_step = 1.0 / len(gp_data_list) 168 | x_plot_segments = np.linspace(0, 1, len(gp_data_list) + 1) 169 | # x_plot_all = np.linspace(0, 1, num_points * len(gp_data_list)) 170 | 171 | for j, gp_data in enumerate(gp_data_list): 172 | num_points = gp_data.x_path.shape[0] 173 | if x_path_mode == "sequential": 174 | # x_plot_step = x_plot_segments_step / num_points 175 | # x_plot = np.arange(x_plot_segments[j], x_plot_segments[j+1], x_plot_step) 176 | x_plot = np.linspace(x_plot_segments[j], x_plot_segments[j + 1], num_points) 177 | else: 178 | x_plot = np.linspace(0, 1, num_points) 179 | 180 | for i in range(0, nout): 181 | add_gp_plot( 182 | gp_data, 183 | ax[i], 184 | i, 185 | x_plot=x_plot, 186 | color=color_list[j], 187 | marker_style=marker_style, 188 | marker_size_lim=marker_size_lim, 189 | marker_size_reldiff_zero=marker_size_reldiff_zero, 190 | plot_train_data=plot_train_data, 191 | ) 192 | 193 | # general settings 194 | fig.set_facecolor("white") 195 | 196 | return fig, ax 197 | 198 | 199 | def add_gp_plot( 200 | gp_data: GPlotData, 201 | ax, 202 | idx_out, 203 | x_plot=None, 204 | color="b", 205 | marker_size_lim=[5, 100], 206 | marker_size_reldiff_zero=1e-3, 207 | marker_style="x", 208 | plot_train_data=True, 209 | ): 210 | x_path = gp_data.x_path 211 | 212 | if plot_train_data: 213 | train_x = gp_data.train_x 214 | train_y_plot = gp_data.train_y[:, idx_out] 215 | 216 | # project on path 217 | train_x_on_path, train_x_dist_to_path, train_x_on_path_index = ( 218 | project_data_on_path(train_x, x_path) 219 | ) 220 | 221 | # square distance again for marker scaling (also quadratic) 222 | train_x_dist_to_path = train_x_dist_to_path**2 223 | 224 | # rescale to marker size limits 225 | train_x_dist_scale_to_zero = ( 226 | train_x_dist_to_path / np.sum(train_x_on_path**2, axis=1) 227 | >= marker_size_reldiff_zero 228 | ) * 1.0 229 | marker_size_reldiff = (train_x_dist_to_path - min(train_x_dist_to_path)) / ( 230 | max(train_x_dist_to_path) - min(train_x_dist_to_path) 231 | ) 232 | train_x_dist_scale = marker_size_lim[1] + ( 233 | marker_size_lim[0] - marker_size_lim[1] 234 | ) * (marker_size_reldiff * train_x_dist_scale_to_zero) 235 | 236 | num_points = x_path.shape[0] 237 | if x_plot is None: 238 | x_plot = np.linspace(0, 1, num_points) 239 | 240 | train_x_on_plot = x_plot[train_x_on_path_index] 241 | ax.scatter( 242 | train_x_on_plot, 243 | train_y_plot, 244 | s=train_x_dist_scale, 245 | marker=marker_style, 246 | color=color, 247 | alpha=0.5, 248 | ) 249 | 250 | # Predictive mean as blue line 251 | mean = gp_data.mean_on_path[:, idx_out] 252 | ax.plot(x_plot, mean, color=color) 253 | 254 | # Shade in confidence 255 | # stddev = gp_data.stddev_on_path[:, idx_out] 256 | if (gp_data.conf_upper_on_path is not None) and ( 257 | gp_data.conf_lower_on_path is not None 258 | ): 259 | upper = gp_data.conf_upper_on_path[:, idx_out] 260 | lower = gp_data.conf_lower_on_path[:, idx_out] 261 | y_plot_fill = ax.fill_between(x_plot, lower, upper, alpha=0.3, color=color) 262 | 263 | # plot samples 264 | if gp_data.sample_on_path is not None: 265 | num_samples = gp_data.sample_on_path.shape[2] 266 | for j in range(0, num_samples): 267 | sample = gp_data.sample_on_path[:, idx_out, j] 268 | ax.plot(x_plot, sample, color=color, linestyle=":", lw=1) 269 | 270 | # if gp_out_lims is not None: 271 | # ax.set_ylim(gp_out_lims[i,:]) 272 | 273 | ax.legend(["Observed Data", "Mean", "Confidence"]) 274 | ax.set_title(f"Observed Values (Likelihood), output {idx_out}") 275 | return ax 276 | 277 | 278 | def plot_gp_model( 279 | gp_model, 280 | likelihood, 281 | x_path, 282 | num_samples=0, 283 | marker_style="x", 284 | marker_size_lim=[1, 5], 285 | ): 286 | # gp_model, likelihood and x_path can be lists 287 | # Initialize plots 288 | nout = gp_model.train_targets[0].shape[0] 289 | fig, ax = plt.subplots(1, nout, figsize=(8, 3)) 290 | if nout == 1: 291 | ax = [ax] 292 | 293 | # along which dim to plot (in gp_dim_lims[dim], other values fixed to gp_dim_slice[other_dims]) 294 | # nout = gp_model.train_targets[0].shape[0] 295 | num_points = x_path.shape[0] 296 | 297 | # Set into eval mode 298 | gp_model.eval() 299 | likelihood.eval() 300 | 301 | # Make predictions 302 | with torch.no_grad(), gpytorch.settings.fast_pred_var(): 303 | # predictions = likelihood(gp_model(test_x)) 304 | predictions = gp_model( 305 | torch.Tensor(x_path) 306 | ) # only model (we want to find true function) 307 | mean = predictions.mean 308 | lower, upper = predictions.confidence_region() 309 | 310 | # draw function samples 311 | predictions_samples = np.zeros((num_points, nout, num_samples)) 312 | for j in range(0, num_samples): 313 | predictions_samples[:, :, j] = predictions.sample() 314 | 315 | # loop through all outputs and plot 316 | for i in range(0, nout): 317 | # TODO: Why do I need train_inputs[0]? Some batching possibilities? 318 | # plot data projected on path 319 | train_x = gp_model.train_inputs[0].numpy() 320 | train_y_plot = gp_model.train_targets[:, i].numpy() 321 | train_x_on_path, train_x_dist_to_path, train_x_on_path_index = ( 322 | project_data_on_path(train_x, x_path) 323 | ) 324 | 325 | # square distance again for marker scaling (also quadratic) 326 | train_x_dist_to_path = train_x_dist_to_path**2 327 | 328 | # rescale to marker size limits 329 | train_x_dist_scale = marker_size_lim[0] + ( 330 | marker_size_lim[1] - marker_size_lim[0] 331 | ) * (train_x_dist_to_path - min(train_x_dist_to_path)) / ( 332 | max(train_x_dist_to_path) - min(train_x_dist_to_path) 333 | ) 334 | test_x_plot = np.linspace(0, 1, num_points) 335 | train_x_on_plot = test_x_plot[train_x_on_path_index] 336 | ax[i].scatter( 337 | train_x_on_plot, 338 | train_y_plot, 339 | s=train_x_dist_scale, 340 | marker=marker_style, 341 | color="k", 342 | alpha=0.5, 343 | ) 344 | 345 | # Predictive mean as blue line 346 | ax[i].plot(test_x_plot, mean[:, i].numpy(), "b") 347 | 348 | # Shade in confidence 349 | y_plot_fill = ax[i].fill_between( 350 | test_x_plot, lower[:, i].numpy(), upper[:, i].numpy(), alpha=0.5 351 | ) 352 | 353 | # plot samples 354 | for j in range(0, num_samples): 355 | ax[i].plot(test_x_plot, predictions_samples[:, i, j], "b:", lw=1) 356 | 357 | # if gp_out_lims is not None: 358 | # ax[i].set_ylim(gp_out_lims[i,:]) 359 | 360 | ax[i].legend(["Observed Data", "Mean", "Confidence"]) 361 | ax[i].set_title(f"Observed Values (Likelihood), output {i}") 362 | 363 | # gp_data = gp_data_from_model_and_path(gp_model, likelihood, x_path, 364 | # num_samples = num_samples 365 | # ) 366 | 367 | # general settings 368 | fig.set_facecolor("white") 369 | 370 | 371 | def project_data_on_path(x_data: np.array, x_path: np.array): 372 | # x_path: n_plot x dim 373 | # x_data: n_data x dim 374 | n_data = x_data.shape[0] 375 | n_path = x_path.shape[0] 376 | i_path, i_data = np.meshgrid(np.arange(0, n_path), np.arange(0, n_data)) 377 | dist = np.sqrt(np.sum((x_data[i_data] - x_path[i_path]) ** 2, axis=2)) 378 | dist_min_data_index = np.argmin(dist, axis=1) 379 | dist_min_data = dist[np.arange(0, n_data), dist_min_data_index] 380 | x_data_on_path = x_path[dist_min_data_index, :] 381 | return x_data_on_path, dist_min_data, dist_min_data_index 382 | 383 | 384 | def generate_grid_points(x_dim_lims, x_dim_slice, x_dim_plot, num_points=200): 385 | x_dim_fix = np.arange(len(x_dim_slice)) 386 | for i in x_dim_fix: 387 | if i == x_dim_plot: 388 | x_add = np.linspace(x_dim_lims[i, 0], x_dim_lims[i, 1], num_points) 389 | else: 390 | x_add = x_dim_slice[i] * np.ones((num_points,)) 391 | 392 | # vstack 393 | if i == 0: 394 | x_grid = x_add 395 | else: 396 | x_grid = np.vstack((x_grid, x_add)) 397 | return x_grid.transpose() 398 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "l4acados" 3 | version = "0.0.1" 4 | authors = [ 5 | { name="Amon Lahr", email="amlahr@ethz.ch" }, 6 | { name="Joshua Näf", email="joshua.naef@ethz.ch" }, 7 | ] 8 | description = "Learning-based models for acados" 9 | readme = "README.md" 10 | requires-python = ">=3.8" 11 | classifiers = [ 12 | "Programming Language :: Python :: 3", 13 | "License :: OSI Approved :: MIT License", 14 | "Operating System :: OS Independent", 15 | ] 16 | dependencies = [ 17 | "casadi", 18 | "scipy", 19 | "numpy", 20 | ] 21 | 22 | [project.optional-dependencies] 23 | pytorch = [ 24 | "torch", 25 | ] 26 | gpytorch = [ 27 | "gpytorch", 28 | ] 29 | gpytorch-exo = [ 30 | "linear_operator @ git+https://github.com/naefjo/linear_operator@feature/exo-gp", 31 | "gpytorch @ git+https://github.com/naefjo/gpytorch@feature/exo-gp", 32 | ] 33 | test = [ 34 | "pytest", 35 | "l4casadi", 36 | ] 37 | dev = [ 38 | "l4acados[test]", 39 | "pre-commit", 40 | "jupytext", 41 | ] 42 | 43 | [build-system] 44 | requires = ["setuptools"] 45 | build-backend = "setuptools.build_meta" 46 | 47 | [tool.pytest.ini_options] 48 | addopts = [ 49 | "--import-mode=importlib", 50 | ] 51 | pythonpath = "test" 52 | 53 | [tool.setuptools.packages.find] 54 | where = ["src"] 55 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | from setuptools import setup 3 | 4 | if __name__ == "__main__": 5 | setup() 6 | -------------------------------------------------------------------------------- /src/l4acados/__init__.py: -------------------------------------------------------------------------------- 1 | from . import controllers, models 2 | -------------------------------------------------------------------------------- /src/l4acados/controllers/__init__.py: -------------------------------------------------------------------------------- 1 | from .residual_learning_mpc import ResidualLearningMPC 2 | from .zero_order_gpmpc import ZeroOrderGPMPC 3 | -------------------------------------------------------------------------------- /src/l4acados/controllers/residual_learning_mpc.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from acados_template import AcadosOcp, AcadosSim, AcadosSimSolver, AcadosOcpSolver 4 | from .zoro_acados_utils import ( 5 | transform_ocp, 6 | setup_sim_from_ocp, 7 | get_solve_opts_from_ocp, 8 | ) 9 | from l4acados.models import ResidualModel 10 | from typing import Tuple 11 | from time import perf_counter 12 | 13 | 14 | class ResidualLearningMPC: 15 | def __init__( 16 | self, 17 | ocp: AcadosOcp, 18 | B: np.ndarray = None, 19 | residual_model: ResidualModel = None, 20 | build_c_code: bool = True, 21 | use_cython: bool = True, 22 | use_ocp_model: bool = True, 23 | ignore_support_errors: bool = False, 24 | path_json_ocp: str = "residual_lbmpc_ocp_solver_config.json", 25 | path_json_sim: str = "residual_lbmpc_sim_solver_config.json", 26 | ) -> None: 27 | """ 28 | ocp: AcadosOcp for nominal problem 29 | B: Residual Jacobian mapping from resiudal dimension to state dimension. If None, 30 | it is assumed that redsidual_dim == state_dim 31 | residual_model: ResidualModel class 32 | build_c_code: Whether the solver should be built. Note, if you do not build it here, 33 | you will have to call the build function before calling the solver. 34 | The following args only apply if build_c_code == True: 35 | use_cython: Whether Acados' cython solver interface should be used. You probably 36 | want this enabled. 37 | use_ocp_model: Whether the OCP model should be used as the nominal model. If set to false, no nominal model will be used. 38 | path_json_ocp: Name of the json file where the resulting ocp will be dumped 39 | path_json_sim: Name of the json file where the resulting sim will be dumped 40 | """ 41 | 42 | ocp.make_consistent() 43 | 44 | # optional argument 45 | if B is None: 46 | B = np.eye(ocp.dims.nx) 47 | self.B = B 48 | 49 | # transform OCP to linear-params-model 50 | self.ocp, self.ocp_opts = transform_ocp( 51 | ocp, use_cython, ignore_errors=ignore_support_errors 52 | ) 53 | self.ocp_opts_tol_arr = np.array( 54 | [ 55 | self.ocp_opts["nlp_solver_tol_stat"], 56 | self.ocp_opts["nlp_solver_tol_eq"], 57 | self.ocp_opts["nlp_solver_tol_ineq"], 58 | self.ocp_opts["nlp_solver_tol_comp"], 59 | ] 60 | ) 61 | 62 | # get dimensions 63 | self.nx = self.ocp.dims.nx 64 | self.nu = self.ocp.dims.nu 65 | self.np_nonlin = ocp.dims.np 66 | self.np_linmdl = self.ocp.dims.np 67 | self.N = self.ocp.dims.N 68 | self.nw = B.shape[1] 69 | 70 | # allocation 71 | self.x_hat_all = np.zeros((self.N + 1, self.nx)) 72 | self.u_hat_all = np.zeros((self.N, self.nu)) 73 | self.y_hat_all = np.zeros((self.N, self.nx + self.nu)) 74 | self.p_hat_nonlin = np.array([ocp.parameter_values for _ in range(self.N + 1)]) 75 | self.p_hat_linmdl = np.array( 76 | [self.ocp.parameter_values for _ in range(self.N + 1)] 77 | ) 78 | self.nominal_fun = np.zeros((self.N, self.nx)) 79 | self.nominal_jac = np.zeros((self.N, self.nx, self.nx + self.nu)) 80 | self.residual_fun = np.zeros((self.N, self.nw)) 81 | self.residual_jac = np.zeros((self.nw, self.N, self.nx + self.nu)) 82 | self.nlp_residuals = np.zeros((self.ocp_opts["nlp_solver_max_iter"], 4)) 83 | self.time_preparation = np.zeros((self.ocp_opts["nlp_solver_max_iter"],)) 84 | self.time_feedback = np.zeros((self.ocp_opts["nlp_solver_max_iter"],)) 85 | self.time_residual = np.zeros((self.ocp_opts["nlp_solver_max_iter"],)) 86 | self.time_nominal = np.zeros((self.ocp_opts["nlp_solver_max_iter"],)) 87 | self.time_iter = np.zeros((self.ocp_opts["nlp_solver_max_iter"],)) 88 | self.num_iter = 0 89 | 90 | self.has_residual_model = False 91 | if residual_model is not None: 92 | self.has_residual_model = True 93 | self.residual_model = residual_model 94 | 95 | self.sim = None 96 | self.has_nominal_model = use_ocp_model 97 | if self.has_nominal_model: 98 | self.sim = setup_sim_from_ocp(ocp) 99 | 100 | self.build( 101 | use_cython=use_cython, 102 | build_c_code=build_c_code, 103 | path_json_ocp=path_json_ocp, 104 | path_json_sim=path_json_sim, 105 | ) 106 | 107 | self.init_last_iterate() 108 | 109 | def build( 110 | self, 111 | use_cython=False, 112 | build_c_code=True, 113 | path_json_ocp="residual_lbmpc_ocp_solver_config.json", 114 | path_json_sim="residual_lbmpc_sim_solver_config.json", 115 | ) -> None: 116 | """ 117 | build_c_code: Whether the solver should be built. If set to false, the solver 118 | will simply be read from the json files. 119 | use_cython: Whether Acados' cython solver interface should be used. You probably 120 | want this enabled. 121 | path_json_ocp: Name of the json file where the resulting ocp will be dumped 122 | (or was dumped if build_c_code == False) 123 | path_json_sim: Name of the json file where the resulting sim will be dumped 124 | (or was dumped if build_c_code == False) 125 | """ 126 | self.use_cython = use_cython 127 | if use_cython: 128 | if build_c_code: 129 | AcadosOcpSolver.generate(self.ocp, json_file=path_json_ocp) 130 | AcadosOcpSolver.build(self.ocp.code_export_directory, with_cython=True) 131 | self.ocp_solver = AcadosOcpSolver.create_cython_solver(path_json_ocp) 132 | 133 | if self.has_nominal_model: 134 | if build_c_code: 135 | AcadosSimSolver.generate(self.sim, json_file=path_json_sim) 136 | AcadosSimSolver.build( 137 | self.sim.code_export_directory, with_cython=True 138 | ) 139 | self.sim_solver = AcadosSimSolver.create_cython_solver(path_json_sim) 140 | else: 141 | self.ocp_solver = AcadosOcpSolver(self.ocp, json_file=path_json_ocp) 142 | if self.has_nominal_model: 143 | self.sim_solver = AcadosSimSolver(self.sim, json_file=path_json_sim) 144 | 145 | def solve(self, acados_sqp_mode=False): 146 | status_feed = 0 147 | for i in range(self.ocp_opts["nlp_solver_max_iter"]): 148 | time_iter_start = perf_counter() 149 | self.num_iter = i 150 | self.preparation() 151 | 152 | if acados_sqp_mode: 153 | self.store_last_iterate() 154 | 155 | status_feed = self.feedback() 156 | if status_feed != 0: 157 | raise Exception( 158 | "acados self.ocp_solver returned status {} in time step {}. Exiting.".format( 159 | status_feed, i 160 | ) 161 | ) 162 | self.time_iter[i] = perf_counter() - time_iter_start 163 | 164 | # ------------------- Check termination -------------------- 165 | if self.ocp.solver_options.rti_log_residuals: 166 | self.nlp_residuals[i, :] = self.get_initial_residuals() 167 | 168 | if np.all(self.nlp_residuals[i, :] < self.ocp_opts_tol_arr): 169 | if acados_sqp_mode: 170 | # restore previous iterate for which residuals are valid 171 | self.load_last_iterate() 172 | break 173 | 174 | return status_feed 175 | 176 | def preparation(self): 177 | time_preparation_start = perf_counter() 178 | # ------------------- Query nodes -------------------- 179 | # preparation rti_phase (solve() AFTER setting params to get right Jacobians) 180 | self.ocp_solver.options_set("rti_phase", 1) 181 | 182 | # get linearization points for all stages 183 | for stage in range(self.N): 184 | # current stage values 185 | self.x_hat_all[stage, :] = self.ocp_solver.get(stage, "x") 186 | self.u_hat_all[stage, :] = self.ocp_solver.get(stage, "u") 187 | self.y_hat_all[stage, :] = np.hstack( 188 | (self.x_hat_all[stage, :], self.u_hat_all[stage, :]) 189 | ).reshape((1, self.nx + self.nu)) 190 | 191 | self.x_hat_all[self.N, :] = self.ocp_solver.get(self.N, "x") 192 | 193 | # ------------------- Sensitivities -------------------- 194 | if self.has_residual_model: 195 | time_residual_start = perf_counter() 196 | self.residual_fun, self.residual_jac = ( 197 | self.residual_model.value_and_jacobian(self.y_hat_all) 198 | ) 199 | self.time_residual[self.num_iter] = perf_counter() - time_residual_start 200 | 201 | if self.has_nominal_model: 202 | time_nominal_start = perf_counter() 203 | for stage in range(self.N): 204 | # set parameters (linear matrices and offset) 205 | # ------------------- Integrate -------------------- 206 | self.sim_solver.set("x", self.x_hat_all[stage, :]) 207 | self.sim_solver.set("u", self.u_hat_all[stage, :]) 208 | self.sim_solver.set("p", self.p_hat_nonlin[stage, :]) 209 | status_integrator = self.sim_solver.solve() 210 | 211 | self.nominal_fun[stage, :] = self.sim_solver.get("x") 212 | self.nominal_jac[stage, :, 0 : self.nx] = self.sim_solver.get("Sx") 213 | self.nominal_jac[stage, :, self.nx : self.nx + self.nu] = ( 214 | self.sim_solver.get("Su") 215 | ) 216 | self.time_nominal[self.num_iter] = perf_counter() - time_nominal_start 217 | 218 | # ------------------- Update stages -------------------- 219 | for stage in range(self.N): 220 | # ------------------- Build linear model -------------------- 221 | A_total = ( 222 | self.nominal_jac[stage, :, 0 : self.nx] 223 | + self.B @ self.residual_jac[:, stage, 0 : self.nx] 224 | ) 225 | B_total = ( 226 | self.nominal_jac[stage, :, self.nx : self.nx + self.nu] 227 | + self.B @ self.residual_jac[:, stage, self.nx : self.nx + self.nu] 228 | ) 229 | 230 | f_hat = ( 231 | self.nominal_fun[stage, :] 232 | + self.B @ self.residual_fun[stage, :] 233 | - A_total @ self.x_hat_all[stage, :] 234 | - B_total @ self.u_hat_all[stage, :] 235 | ) 236 | 237 | # ------------------- Set sensitivities -------------------- 238 | A_reshape = np.reshape(A_total, (self.nx**2), order="F") 239 | B_reshape = np.reshape(B_total, (self.nx * self.nu), order="F") 240 | 241 | self.p_hat_linmdl[stage, :] = np.hstack( 242 | (A_reshape, B_reshape, f_hat, self.p_hat_nonlin[stage, :]) 243 | ) 244 | self.ocp_solver.set(stage, "p", self.p_hat_linmdl[stage, :]) 245 | 246 | # use stage N for linear part of last stage (unused anyway) 247 | self.p_hat_linmdl[self.N, :] = np.hstack( 248 | (A_reshape, B_reshape, f_hat, self.p_hat_nonlin[self.N, :]) 249 | ) 250 | self.ocp_solver.set(self.N, "p", self.p_hat_linmdl[self.N, :]) 251 | 252 | # feedback rti_phase 253 | # ------------------- Phase 1 -------------------- 254 | status = self.ocp_solver.solve() 255 | self.time_preparation[self.num_iter] = perf_counter() - time_preparation_start 256 | 257 | def feedback(self): 258 | time_feedback_start = perf_counter() 259 | # ------------------- Solve QP -------------------- 260 | self.ocp_solver.options_set("rti_phase", 2) 261 | status = self.ocp_solver.solve() 262 | self.time_feedback[self.num_iter] = perf_counter() - time_feedback_start 263 | return status 264 | 265 | def get_solution( 266 | self, 267 | ) -> Tuple[np.ndarray, np.ndarray]: 268 | X = np.zeros((self.N + 1, self.nx)) 269 | U = np.zeros((self.N, self.nu)) 270 | 271 | # get data 272 | for i in range(self.N): 273 | X[i, :] = self.ocp_solver.get(i, "x") 274 | U[i, :] = self.ocp_solver.get(i, "u") 275 | 276 | X[self.N, :] = self.ocp_solver.get(self.N, "x") 277 | 278 | return X, U 279 | 280 | def init_last_iterate(self): 281 | npi = len(self.ocp_solver.get(0, "pi")) 282 | nlam_0 = len(self.ocp_solver.get(0, "lam")) 283 | if self.N > 1: 284 | nlam = len(self.ocp_solver.get(1, "lam")) 285 | else: 286 | nlam = 0 287 | nlam_e = len(self.ocp_solver.get(self.N, "lam")) 288 | 289 | self.pi_hat_all = np.zeros((self.N, self.nx)) 290 | self.lam_hat_all_0 = np.zeros((nlam_0,)) 291 | self.lam_hat_all = np.zeros((self.N, nlam)) 292 | self.lam_hat_all_e = np.zeros((nlam_e,)) 293 | 294 | self.x_hat_all_lastiter = np.zeros((self.N + 1, self.nx)) 295 | self.u_hat_all_lastiter = np.zeros((self.N, self.nu)) 296 | self.pi_hat_all_lastiter = np.zeros((self.N, npi)) 297 | self.lam_hat_all_lastiter_0 = np.zeros((nlam_0,)) 298 | self.lam_hat_all_lastiter = np.zeros((self.N, nlam)) 299 | self.lam_hat_all_lastiter_e = np.zeros((nlam_e,)) 300 | 301 | def load_last_iterate(self): 302 | # TODO: remove this once cython support is enabled 303 | if self.use_cython: 304 | self.ocp_solver.set(0, "lam", self.lam_hat_all_lastiter_0[:]) 305 | for stage in range(self.N): 306 | self.ocp_solver.set(stage, "x", self.x_hat_all_lastiter[stage, :]) 307 | self.ocp_solver.set(stage, "u", self.u_hat_all_lastiter[stage, :]) 308 | self.ocp_solver.set(stage, "pi", self.pi_hat_all_lastiter[stage, :]) 309 | if stage > 0: 310 | self.ocp_solver.set( 311 | stage, "lam", self.lam_hat_all_lastiter[stage, :] 312 | ) 313 | self.ocp_solver.set(self.N, "x", self.x_hat_all_lastiter[self.N, :]) 314 | self.ocp_solver.set(self.N, "lam", self.lam_hat_all_lastiter_e[:]) 315 | else: 316 | self.ocp_solver.load_iterate_from_flat_obj(self.last_iterate) 317 | 318 | def store_last_iterate(self): 319 | # TODO: remove this once cython support is enabled 320 | if self.use_cython: 321 | self.lam_hat_all_lastiter_0[:] = self.ocp_solver.get(0, "lam") 322 | for stage in range(self.N): 323 | self.x_hat_all_lastiter[stage, :] = self.ocp_solver.get(stage, "x") 324 | self.u_hat_all_lastiter[stage, :] = self.ocp_solver.get(stage, "u") 325 | self.pi_hat_all_lastiter[stage, :] = self.ocp_solver.get(stage, "pi") 326 | if stage > 0: 327 | self.lam_hat_all_lastiter[stage, :] = self.ocp_solver.get( 328 | stage, "lam" 329 | ) 330 | self.x_hat_all_lastiter[self.N, :] = self.ocp_solver.get(self.N, "x") 331 | self.lam_hat_all_lastiter_e[:] = self.ocp_solver.get(self.N, "lam") 332 | else: 333 | self.last_iterate = self.ocp_solver.store_iterate_to_flat_obj() 334 | 335 | # ------------------- Forward OCP solver functions ------------------- 336 | def dump_last_qp_to_json(self, *args, **kwargs) -> None: 337 | self.ocp_solver.dump_last_qp_to_json(*args, **kwargs) 338 | 339 | def get(self, stage: int, var: str) -> np.ndarray: 340 | if var == "p": 341 | assert 0 <= stage <= self.N 342 | return self.p_hat_nonlin[stage, :] # params are set in preparation phase 343 | else: 344 | return self.ocp_solver.get(stage, var) 345 | 346 | def get_initial_residuals(self) -> np.ndarray: 347 | return self.ocp_solver.get_initial_residuals() 348 | 349 | def get_residuals(self, recompute=False, ignore_warning=False) -> np.ndarray: 350 | if ignore_warning: 351 | return self.ocp_solver.get_residuals(recompute, ignore_warning) 352 | raise ValueError( 353 | "Only logging of available residuals is supported. Use get_initial_residuals() instead. See https://github.com/acados/acados/pull/1346." 354 | ) 355 | 356 | def get_stats(self, stat: str): 357 | if stat == "res_stat_all": 358 | return self.nlp_residuals[: self.num_iter + 1, 0] 359 | if stat == "res_eq_all": 360 | return self.nlp_residuals[: self.num_iter + 1, 1] 361 | if stat == "res_ineq_all": 362 | return self.nlp_residuals[: self.num_iter + 1, 2] 363 | if stat == "res_comp_all": 364 | return self.nlp_residuals[: self.num_iter + 1, 3] 365 | if stat == "sqp_iter" or stat == "nlp_iter": 366 | return self.num_iter + 1 367 | if stat == "time_preparation": 368 | return self.time_preparation[self.num_iter] 369 | if stat == "time_feedback": 370 | return self.time_feedback[self.num_iter] 371 | if stat == "time_preparation_all": 372 | return self.time_preparation[: self.num_iter + 1] 373 | if stat == "time_feedback_all": 374 | return self.time_feedback[: self.num_iter + 1] 375 | if stat == "time_sim": 376 | return self.time_nominal[self.num_iter] 377 | if stat == "time_sim_all": 378 | return self.time_nominal[: self.num_iter + 1] 379 | if stat == "time_lin": 380 | return self.time_nominal[self.num_iter] + self.time_residual[self.num_iter] 381 | if stat == "time_lin_all": 382 | return ( 383 | self.time_nominal[: self.num_iter + 1] 384 | + self.time_residual[: self.num_iter + 1] 385 | ) 386 | if stat == "time_tot": 387 | return np.sum(self.time_iter) 388 | print( 389 | f"Warning: Returning acados solver status '{stat}', which might be incorrect when used with L4acados." 390 | ) 391 | return self.ocp_solver.get_stats(stat) 392 | 393 | def set(self, stage: int, var: str, value: np.ndarray) -> None: 394 | if var == "p": 395 | assert 0 <= stage <= self.N 396 | self.p_hat_nonlin[stage, :] = value 397 | else: 398 | self.ocp_solver.set(stage, var, value) 399 | 400 | def cost_set(self, stage: int, cost: str, value: np.ndarray) -> None: 401 | self.ocp_solver.cost_set(stage, cost, value) 402 | 403 | def constraints_set(self, stage: int, constraint: str, value: np.ndarray) -> None: 404 | self.ocp_solver.constraints_set(stage, constraint, value) 405 | 406 | def options_set(self, option: str, value) -> None: 407 | self.ocp_solver.options_set(option, value) 408 | 409 | def print_statistics(self) -> None: 410 | if self.ocp_opts["nlp_solver_type"] == "SQP": 411 | print("iter res_stat res_eq res_ineq res_comp") 412 | for i in range(self.num_iter + 1): 413 | print( 414 | f"{i:<7d} {self.nlp_residuals[i, 0]:12e} {self.nlp_residuals[i, 1]:12e} {self.nlp_residuals[i, 2]:12e} {self.nlp_residuals[i, 3]:11e}" 415 | ) 416 | return 417 | self.ocp_solver.print_statistics() 418 | 419 | def solve_for_x0( 420 | self, x0_bar, fail_on_nonzero_status=True, print_stats_on_failure=True 421 | ) -> np.ndarray: 422 | self.ocp_solver.set(0, "lbx", x0_bar) 423 | self.ocp_solver.set(0, "ubx", x0_bar) 424 | 425 | status = self.solve() 426 | 427 | if status != 0: 428 | if print_stats_on_failure: 429 | self.ocp_solver.print_statistics() 430 | if fail_on_nonzero_status: 431 | raise Exception(f"acados acados_ocp_solver returned status {status}") 432 | elif print_stats_on_failure: 433 | print(f"Warning: acados acados_ocp_solver returned status {status}") 434 | 435 | u0 = self.ocp_solver.get(0, "u") 436 | return u0 437 | -------------------------------------------------------------------------------- /src/l4acados/controllers/zero_order_gpmpc.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from .zoro_acados_utils import * 3 | from .residual_learning_mpc import ResidualLearningMPC 4 | from l4acados.models import ResidualModel 5 | from acados_template import AcadosOcp 6 | 7 | 8 | class ZeroOrderGPMPC(ResidualLearningMPC): 9 | def __init__( 10 | self, 11 | ocp: AcadosOcp, 12 | B: np.ndarray = None, 13 | residual_model: ResidualModel = None, 14 | use_cython: bool = True, 15 | path_json_ocp: str = "zogpmpc_ocp_solver_config.json", 16 | path_json_sim: str = "zogpmpc_sim_solver_config.json", 17 | build_c_code: bool = True, 18 | ) -> None: 19 | 20 | ocp = self.setup_custom_update(ocp) 21 | 22 | super().__init__( 23 | ocp, 24 | B=B, 25 | residual_model=residual_model, 26 | use_cython=use_cython, 27 | path_json_ocp=path_json_ocp, 28 | path_json_sim=path_json_sim, 29 | build_c_code=build_c_code, 30 | ) 31 | 32 | def preparation(self): 33 | super().preparation() 34 | self.do_custom_update() 35 | 36 | def setup_custom_update(self, ocp): 37 | template_c_file = "custom_update_function_zoro_template.in.c" 38 | template_h_file = "custom_update_function_zoro_template.in.h" 39 | custom_c_file = "custom_update_function.c" 40 | custom_h_file = "custom_update_function.h" 41 | 42 | # custom update: disturbance propagation 43 | ocp.solver_options.custom_update_filename = custom_c_file 44 | ocp.solver_options.custom_update_header_filename = custom_h_file 45 | 46 | ocp.solver_options.custom_templates = [ 47 | ( 48 | template_c_file, 49 | custom_c_file, 50 | ), 51 | ( 52 | template_h_file, 53 | custom_h_file, 54 | ), 55 | ] 56 | 57 | ocp.solver_options.custom_update_copy = False 58 | """NOTE(@naefjo): As far as I understand you need to set this variable to True if you just 59 | want to copy an existing custom_update.c/h into the export directory and to False if you want 60 | to render the custom_udpate files from the template""" 61 | 62 | if ocp.zoro_description.input_P0_diag: 63 | self.zoro_input_P0 = np.diag(ocp.zoro_description.P0_mat) 64 | else: 65 | self.zoro_input_P0 = ocp.zoro_description.P0_mat 66 | 67 | self.zoro_input_W_diag = np.diag(ocp.zoro_description.W_mat) 68 | self.covariances_array = np.zeros(((ocp.dims.N + 1) * ocp.dims.nx**2,)) 69 | 70 | return ocp 71 | 72 | def do_custom_update(self) -> None: 73 | """performs the acados custom update and propagates the covariances for the constraint tightening 74 | 75 | The array which is passed to the custom update function consists of an input array and 76 | an output array [cov_in, cov_out], where 77 | cov_in = [Sigma_x0, Sigma_w, [Sigma_GP_i forall i in (0, N-1)]] and 78 | cov_out = [-1*ones(3 * (N + 1)))] is a placeholder for the positional covariances used for 79 | visualization. 80 | 81 | Note that the function currently only supports setting the diagonal elements of the covariance matrices 82 | in the solver. 83 | """ 84 | if not self.has_residual_model: 85 | return 86 | 87 | covariances_in = np.concatenate( 88 | ( 89 | self.zoro_input_P0, 90 | self.zoro_input_W_diag, 91 | # self.residual_model.current_variance is updated with value_and_jacobian() call in preparation phase 92 | self.residual_model.current_variance.flatten(), 93 | ) 94 | ) 95 | covariances_in_len = covariances_in.size 96 | out_arr = np.ascontiguousarray( 97 | np.concatenate((covariances_in, -1.0 * np.ones(self.nx**2 * (self.N + 1)))) 98 | ) 99 | self.ocp_solver.custom_update(out_arr) 100 | self.covariances_array[:] = out_arr[covariances_in_len:] 101 | 102 | return 0 103 | -------------------------------------------------------------------------------- /src/l4acados/controllers/zoro_acados_utils.py: -------------------------------------------------------------------------------- 1 | from casadi import SX, MX, vertcat 2 | from acados_template import AcadosModel, AcadosSim 3 | import numpy as np 4 | from copy import deepcopy 5 | import re 6 | 7 | 8 | def export_linear_model(x, u, p): 9 | nx = x.shape[0] 10 | nu = u.shape[0] 11 | 12 | if isinstance(x, MX): 13 | VARX = MX 14 | elif isinstance(x, SX): 15 | VARX = SX 16 | else: 17 | raise ValueError("x must be of type SX or MX") 18 | 19 | nparam = p.shape[0] if isinstance(p, VARX) else 0 20 | 21 | # linear dynamics for every stage 22 | A = VARX.sym("A", nx, nx) 23 | B = VARX.sym("B", nx, nu) 24 | w = VARX.sym("w", nx, 1) 25 | xdot = VARX.sym("xdot", nx, 1) 26 | 27 | f_expl = A @ x + B @ u + w 28 | f_impl = xdot - f_expl 29 | 30 | # parameters 31 | p_lin = vertcat( 32 | A.reshape((nx**2, 1)), B.reshape((nx * nu, 1)), w, p # (P_vec, p_nom) 33 | ) 34 | 35 | # acados model 36 | model = AcadosModel() 37 | model.disc_dyn_expr = f_expl 38 | model.f_impl_expr = f_impl 39 | model.f_expl_expr = f_expl 40 | model.x = x 41 | model.xdot = xdot 42 | model.u = u 43 | model.p = p_lin 44 | model.name = f"linear_model_with_params_nx{nx}_nu{nu}_np{nparam}" 45 | 46 | return model 47 | 48 | 49 | def transform_ocp(ocp_input, use_cython, ignore_errors=False): 50 | ocp = deepcopy(ocp_input) 51 | ocp_opts = get_solve_opts_from_ocp(ocp) 52 | 53 | original_nparam = ocp.dims.np 54 | 55 | model_lin = export_linear_model(ocp.model.x, ocp.model.u, ocp.model.p) 56 | ocp.model.disc_dyn_expr = model_lin.disc_dyn_expr 57 | ocp.model.f_impl_expr = model_lin.f_impl_expr 58 | ocp.model.f_expl_expr = model_lin.f_expl_expr 59 | ocp.model.x = model_lin.x 60 | ocp.model.xdot = model_lin.xdot 61 | ocp.model.u = model_lin.u 62 | ocp.model.p = model_lin.p 63 | ocp.model.name = model_lin.name 64 | ocp.dims.np = model_lin.p.shape[0] 65 | 66 | ocp_parameter_values = ocp.parameter_values 67 | ocp.parameter_values = np.zeros((model_lin.p.shape[0],)) 68 | if original_nparam > 0: 69 | ocp.parameter_values[-original_nparam:] = ocp_parameter_values 70 | 71 | # adjust OCP options in case of SQP to get residuals 72 | if ocp_opts["nlp_solver_type"] == "SQP": 73 | if use_cython: 74 | ocp.solver_options.rti_log_only_available_residuals = 0 75 | ocp.solver_options.rti_log_residuals = 0 76 | print( 77 | "WARNING: Logging of (initial) residuals is not supported for the acados Cython interface. SQP iteration will be terminated based number of iterations only." 78 | ) 79 | else: 80 | ocp.solver_options.rti_log_only_available_residuals = 1 81 | ocp.solver_options.rti_log_residuals = 1 82 | print( 83 | "WARNING: Only logging of initial residuals is supported for SQP. SQP iteration will be terminated based on previous residual." 84 | ) 85 | 86 | # throw errors for unsupported options 87 | if not ignore_errors: 88 | if ocp_opts["nlp_solver_type"] == "SQP": 89 | if ocp_opts["globalization"] == "FIXED_STEP": 90 | if ocp_opts["globalization_fixed_step_length"] != 1.0: 91 | # TODO: implement this 92 | raise ValueError( 93 | "Only globalization_fixed_step_length = 1.0 is supported." 94 | ) 95 | else: 96 | raise ValueError("Only globalization = 'FIXED_STEP' is supported.") 97 | 98 | elif ocp_opts["nlp_solver_type"] == "SQP_RTI": 99 | if ( 100 | ocp.solver_options.rti_log_residuals == 1 101 | and ocp.solver_options.rti_log_only_available_residuals == 0 102 | ): 103 | raise ValueError( 104 | "Only logging of available residuals is supported. Please set rti_log_only_available_residuals to 1." 105 | ) 106 | else: 107 | raise ValueError("Only nlp_solver_type = {'SQP','SQP_RTI'} is supported.") 108 | 109 | ocp.solver_options.integrator_type = "DISCRETE" 110 | ocp.solver_options.nlp_solver_type = "SQP_RTI" 111 | ocp.solver_options.nlp_solver_max_iter = 1 112 | 113 | return ocp, ocp_opts 114 | 115 | 116 | def setup_sim_from_ocp(ocp): 117 | # integrator for nominal model 118 | sim = AcadosSim() 119 | 120 | sim.model = ocp.model 121 | sim.parameter_values = ocp.parameter_values 122 | 123 | for opt_name in dir(ocp.solver_options): 124 | if ( 125 | opt_name in dir(sim.solver_options) 126 | and re.search(r"__.*?__", opt_name) is None 127 | ): 128 | set_value = getattr(ocp.solver_options, opt_name) 129 | 130 | if opt_name == "sim_method_jac_reuse": 131 | set_value = array_to_int(set_value) 132 | 133 | print(f"Setting {opt_name} to {set_value}") 134 | setattr(sim.solver_options, opt_name, set_value) 135 | 136 | if ocp.solver_options.Tsim is not None and ocp.solver_options.tf is not None: 137 | assert np.isclose( 138 | ocp.solver_options.tf / ocp.dims.N, ocp.solver_options.Tsim 139 | ), "Both tf and Tsim are set in ocp, it should hold that (tf / N) == Tsim." 140 | 141 | if ocp.solver_options.Tsim is not None: 142 | sim.solver_options.T = ocp.solver_options.Tsim 143 | else: 144 | sim.solver_options.T = ocp.solver_options.tf / ocp.dims.N 145 | 146 | sim.solver_options.newton_iter = ocp.solver_options.sim_method_newton_iter 147 | sim.solver_options.newton_tol = ocp.solver_options.sim_method_newton_tol 148 | sim.solver_options.num_stages = array_to_int( 149 | ocp.solver_options.sim_method_num_stages 150 | ) 151 | sim.solver_options.num_steps = array_to_int(ocp.solver_options.sim_method_num_steps) 152 | sim.code_export_directory = ocp.code_export_directory 153 | 154 | return sim 155 | 156 | 157 | def get_solve_opts_from_ocp(ocp): 158 | solve_opts = {} 159 | for opt_name in dir(ocp.solver_options): 160 | if re.search(r"__.*?__", opt_name) is not None: 161 | continue 162 | if re.search(r"_AcadosOcpOptions__.*?", opt_name) is not None: 163 | continue 164 | 165 | try: 166 | set_value = getattr(ocp.solver_options, opt_name) 167 | except Exception as e: 168 | print(f"Error getting attribute: {e}") 169 | set_value = None 170 | 171 | print(f"Getting: {opt_name} = {set_value}") 172 | solve_opts[opt_name] = set_value 173 | 174 | return solve_opts 175 | 176 | 177 | def array_to_int(arr): 178 | value = deepcopy(arr) 179 | if type(value) is list or type(value) is np.ndarray: 180 | assert all(value == value[0]) 181 | value = value[0] 182 | 183 | return int(value) 184 | -------------------------------------------------------------------------------- /src/l4acados/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .residual_model import ResidualModel 2 | from .pytorch_models import PyTorchFeatureSelector, PyTorchResidualModel 3 | from .pytorch_models.gpytorch_models import GPyTorchResidualModel 4 | -------------------------------------------------------------------------------- /src/l4acados/models/pytorch_models/__init__.py: -------------------------------------------------------------------------------- 1 | try: 2 | import torch 3 | 4 | _torch_available = True 5 | except ImportError: 6 | _torch_available = False 7 | 8 | if _torch_available: 9 | 10 | from .pytorch_feature_selector import PyTorchFeatureSelector 11 | from .pytorch_residual_model import PyTorchResidualModel 12 | 13 | else: 14 | 15 | from l4acados.models import ResidualModel 16 | 17 | err_msg = "Torch is not available. 'l4acados.pytorch_models' requires the [pytorch] optional dependencies (see pyproject.toml)." 18 | 19 | class PyTorchFeatureSelector: 20 | def __init__(self, *args, **kwargs): 21 | raise RuntimeError(err_msg) 22 | 23 | class PyTorchResidualModel(ResidualModel): 24 | def __init__(self, *args, **kwargs): 25 | raise RuntimeError(err_msg) 26 | -------------------------------------------------------------------------------- /src/l4acados/models/pytorch_models/gpytorch_models/__init__.py: -------------------------------------------------------------------------------- 1 | try: 2 | import gpytorch 3 | 4 | _gpytorch_available = True 5 | except ImportError: 6 | _gpytorch_available = False 7 | 8 | if _gpytorch_available: 9 | 10 | from .gpytorch_residual_model import GPyTorchResidualModel 11 | 12 | else: 13 | 14 | err_msg = "GPyTorch is not available. 'l4acados.gpytorch_models' requires the [gpytorch] optional dependencies (see pyproject.toml)." 15 | 16 | class GPyTorchResidualModel: 17 | def __init__(self, *args, **kwargs): 18 | raise RuntimeError(err_msg) 19 | -------------------------------------------------------------------------------- /src/l4acados/models/pytorch_models/gpytorch_models/gpytorch_data_processing_strategy.py: -------------------------------------------------------------------------------- 1 | from abc import ABC 2 | import threading 3 | from typing import Optional 4 | 5 | import numpy as np 6 | import torch 7 | import gpytorch 8 | import linear_operator 9 | 10 | from ..pytorch_feature_selector import PyTorchFeatureSelector 11 | from ..pytorch_utils import to_numpy, to_tensor 12 | 13 | 14 | class DataProcessingStrategy(ABC): 15 | """Abstract base class for the data processing strategies. 16 | 17 | Depending on the experiment/operating mode of the controller, the data processing 18 | strategy might be different. To enable as much flexibility as possible, the data 19 | processing strategy is impelented using a strategy pattern, allowing the user to 20 | easily define their own strategies and swap them out when necessary 21 | """ 22 | 23 | def process( 24 | self, 25 | gp_model: gpytorch.models.ExactGP, 26 | x_input: np.array, 27 | y_target: np.array, 28 | gp_feature_selector: PyTorchFeatureSelector, 29 | timestamp: Optional[float], 30 | ) -> Optional[gpytorch.models.ExactGP]: 31 | """Function which is processed in the `record_datapoint` method of `ResidualGaussianProcess` 32 | 33 | Args: 34 | - residual_gp_instance: Instance of the `ResidualGaussianProcess` class so we can access 35 | relevant attributes 36 | - x_input: data which should be saved. Should have dimension (state_dimension,) or equivalent 37 | - y_target: the residual which was measured at x_input. Should have dimension 38 | (residual_dimension,) or equivalent 39 | """ 40 | raise NotImplementedError 41 | 42 | 43 | class VoidDataStrategy(DataProcessingStrategy): 44 | def process( 45 | self, 46 | gp_model: gpytorch.models.ExactGP, 47 | x_input: np.array, 48 | y_target: np.array, 49 | gp_feature_selector: PyTorchFeatureSelector, 50 | timestamp: Optional[float], 51 | ) -> Optional[gpytorch.models.ExactGP]: 52 | pass 53 | 54 | 55 | class RecordDataStrategy(DataProcessingStrategy): 56 | """Implements a processing strategy which saves the data continuously to a file. 57 | 58 | The strategy keeps a buffer of recent datapoints and asynchronously saves the buffer to 59 | a file. 60 | """ 61 | 62 | def __init__(self, x_data_path: str, y_data_path: str, buffer_size: int = 50): 63 | """Construct the data recorder. 64 | 65 | Args: 66 | - x_data_path: file path where x data should be saved. 67 | - y_data_path: file path where residual data should be saved 68 | """ 69 | self.x_data_path = x_data_path 70 | self.y_data_path = y_data_path 71 | self.buffer_size = buffer_size 72 | self._gp_training_data = {"x_training_data": [], "y_training_data": []} 73 | 74 | def process( 75 | self, 76 | gp_model: gpytorch.models.ExactGP, 77 | x_input: np.array, 78 | y_target: np.array, 79 | gp_feature_selector: PyTorchFeatureSelector, 80 | timestamp: Optional[float], 81 | ) -> Optional[gpytorch.models.ExactGP]: 82 | self._gp_training_data["x_training_data"].append(x_input) 83 | self._gp_training_data["y_training_data"].append(y_target) 84 | 85 | if len(self._gp_training_data["x_training_data"]) == self.buffer_size: 86 | # Do we need a local copy? 87 | save_data_x = np.array(self._gp_training_data["x_training_data"]) 88 | save_data_y = np.array(self._gp_training_data["y_training_data"]) 89 | 90 | self._gp_training_data["x_training_data"].clear() 91 | self._gp_training_data["y_training_data"].clear() 92 | 93 | threading.Thread( 94 | target=lambda: ( 95 | RecordDataStrategy._save_to_file(save_data_x, self.x_data_path), 96 | RecordDataStrategy._save_to_file(save_data_y, self.y_data_path), 97 | print("saved gp training data"), 98 | ) 99 | ).start() 100 | 101 | @staticmethod 102 | def _save_to_file(data: np.ndarray, filename: str) -> None: 103 | """Appends data to a file""" 104 | with open(filename, "ab") as f: 105 | # f.write(b"\n") 106 | np.savetxt( 107 | f, 108 | data, 109 | delimiter=",", 110 | ) 111 | 112 | 113 | class OnlineLearningStrategy(DataProcessingStrategy): 114 | """Implements an online learning strategy. 115 | 116 | The received data is incorporated in the GP and used for further predictions. 117 | This data processing strategy depends on the [online_gp] optional dependencies (see pyproject.toml). 118 | """ 119 | 120 | def __init__(self, max_num_points: int = 200, device: str = "cpu") -> None: 121 | self.max_num_points = max_num_points 122 | self.device = device 123 | 124 | def process( 125 | self, 126 | gp_model: gpytorch.models.ExactGP, 127 | x_input: np.array, 128 | y_target: np.array, 129 | gp_feature_selector: PyTorchFeatureSelector, 130 | timestamp: Optional[float], 131 | ) -> Optional[gpytorch.models.ExactGP]: 132 | 133 | # Convert to tensor 134 | if not torch.is_tensor(x_input): 135 | x_input = to_tensor(arr=x_input, device=self.device) 136 | 137 | if not torch.is_tensor(y_target): 138 | y_target = to_tensor(arr=y_target, device=self.device) 139 | 140 | # Extend to 2D for further computation 141 | x_input = torch.atleast_2d(x_input) 142 | y_target = torch.atleast_2d(y_target) 143 | 144 | if ( 145 | gp_model.prediction_strategy is None 146 | or gp_model.train_inputs is None 147 | or gp_model.train_targets is None 148 | ): 149 | if gp_model.train_inputs is not None: 150 | raise RuntimeError( 151 | "train_inputs in GP is not None. Something went wrong." 152 | ) 153 | 154 | # Set the training data and return (in-place modification) 155 | gp_model.set_train_data( 156 | gp_feature_selector(x_input), 157 | y_target, 158 | strict=False, 159 | ) 160 | return 161 | 162 | # Check if GP is already full 163 | if gp_model.train_inputs[0].shape[-2] >= self.max_num_points: 164 | with torch.no_grad(): 165 | selector = torch.ones(self.max_num_points, requires_grad=False) 166 | # TODO(@naefjo): Add super cool logic to determine which points to kick out. Ideally with O(-n^3) 167 | drop_idx = torch.randint( 168 | 0, self.max_num_points, torch.Size(), requires_grad=False 169 | ).item() 170 | selector[drop_idx] = 0 171 | # Calculate fantasy model with data selector 172 | try: 173 | fantasy_model = gp_model.get_fantasy_model( 174 | gp_feature_selector(x_input), 175 | y_target, 176 | data_selector=selector, 177 | ) 178 | except TypeError as err: 179 | # check if error message contains data_selector 180 | if "data_selector" in str(err): 181 | raise ImportError( 182 | "OnlineLearningStrategy requires the [online_gp] optional dependencies (see pyproject.toml)." 183 | ) 184 | raise err 185 | 186 | return fantasy_model 187 | 188 | with torch.no_grad(): 189 | # Add observation and return updated model 190 | fantasy_model = gp_model.get_fantasy_model( 191 | gp_feature_selector(x_input), y_target 192 | ) 193 | 194 | return fantasy_model 195 | -------------------------------------------------------------------------------- /src/l4acados/models/pytorch_models/gpytorch_models/gpytorch_gp.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, Union 2 | 3 | import gpytorch 4 | import torch 5 | 6 | 7 | class BatchIndependentMultitaskGPModel(gpytorch.models.ExactGP): 8 | def __init__( 9 | self, 10 | train_x: Optional[torch.tensor], 11 | train_y: Optional[torch.tensor], 12 | likelihood: gpytorch.likelihoods.Likelihood, 13 | use_ard: bool = False, 14 | residual_dimension: Optional[int] = None, 15 | input_dimension: Optional[int] = None, 16 | outputscale_prior: Optional[gpytorch.priors.Prior] = None, 17 | lengthscale_prior: Optional[gpytorch.priors.Prior] = None, 18 | ): 19 | super().__init__(train_x, train_y, likelihood) 20 | 21 | if train_y is None and residual_dimension is None: 22 | raise RuntimeError( 23 | "train_y and residual_dimension are both None. Please specify one." 24 | ) 25 | 26 | if ( 27 | train_y is not None 28 | and residual_dimension is not None 29 | and train_y.size(-1) != residual_dimension 30 | ): 31 | raise RuntimeError( 32 | f"train_y shape {train_y.shape()} and residual_dimension {residual_dimension}" 33 | " do not correspond." 34 | ) 35 | 36 | if train_x is None and input_dimension is None: 37 | raise RuntimeError( 38 | "train_x and input_dimension are both None. Please specify one." 39 | ) 40 | 41 | if use_ard: 42 | ard_input_shape = ( 43 | train_x.size(-1) if train_x is not None else input_dimension 44 | ) 45 | else: 46 | ard_input_shape = None 47 | 48 | residual_dimension = ( 49 | train_y.size(-1) if train_y is not None else residual_dimension 50 | ) 51 | 52 | self.mean_module = gpytorch.means.ZeroMean( 53 | batch_shape=torch.Size([residual_dimension]) 54 | ) 55 | self.covar_module = gpytorch.kernels.ScaleKernel( 56 | gpytorch.kernels.RBFKernel( 57 | ard_num_dims=ard_input_shape, 58 | batch_shape=torch.Size([residual_dimension]), 59 | lengthscale_prior=lengthscale_prior, 60 | ), 61 | batch_shape=torch.Size([residual_dimension]), 62 | outputscale_prior=outputscale_prior, 63 | ) 64 | 65 | def forward(self, x): 66 | mean_x = self.mean_module(x) 67 | covar_x = self.covar_module(x) 68 | return gpytorch.distributions.MultitaskMultivariateNormal.from_batch_mvn( 69 | gpytorch.distributions.MultivariateNormal(mean_x, covar_x) 70 | ) 71 | 72 | 73 | class BatchIndependentInducingPointGpModel(gpytorch.models.ExactGP): 74 | def __init__( 75 | self, 76 | train_x: Optional[torch.tensor], 77 | train_y: Optional[torch.tensor], 78 | likelihood: gpytorch.likelihoods.Likelihood, 79 | inducing_points: Union[int, torch.tensor] = 10, 80 | use_ard: bool = False, 81 | residual_dimension: Optional[int] = None, 82 | input_dimension: Optional[int] = None, 83 | outputscale_prior: Optional[gpytorch.priors.Prior] = None, 84 | lengthscale_prior: Optional[gpytorch.priors.Prior] = None, 85 | ): 86 | """ 87 | TODO(@naefjo): write doc 88 | """ 89 | super().__init__(train_x, train_y, likelihood) 90 | 91 | if train_y is None and residual_dimension is None: 92 | raise RuntimeError( 93 | "train_y and residual_dimension are both None. Please specify one." 94 | ) 95 | 96 | if ( 97 | train_y is not None 98 | and residual_dimension is not None 99 | and train_y.size(-1) != residual_dimension 100 | ): 101 | raise RuntimeError( 102 | f"train_y shape {train_y.shape()} and residual_dimension {residual_dimension}" 103 | " do not correspond." 104 | ) 105 | 106 | if train_x is None and input_dimension is None: 107 | raise RuntimeError( 108 | "train_x and input_dimension are both None. Please specify one." 109 | ) 110 | 111 | if use_ard: 112 | ard_input_shape = ( 113 | train_x.size(-1) if train_x is not None else input_dimension 114 | ) 115 | else: 116 | ard_input_shape = None 117 | 118 | residual_dimension = ( 119 | train_y.size(-1) if train_y is not None else residual_dimension 120 | ) 121 | 122 | self.mean_module = gpytorch.means.ZeroMean( 123 | batch_shape=torch.Size([residual_dimension]) 124 | ) 125 | base_covar_module = gpytorch.kernels.ScaleKernel( 126 | gpytorch.kernels.RBFKernel( 127 | ard_num_dims=ard_input_shape, 128 | batch_shape=torch.Size([residual_dimension]), 129 | lengthscale_prior=lengthscale_prior, 130 | ), 131 | batch_shape=torch.Size([residual_dimension]), 132 | outputscale_prior=outputscale_prior, 133 | ) 134 | 135 | if isinstance(inducing_points, int): 136 | permuted_indices = torch.randperm(train_x.size(0)) 137 | inducing_point_indices = permuted_indices[:inducing_points] 138 | inducing_points = train_x[inducing_point_indices, :] 139 | 140 | self.num_inducing_points = inducing_point_indices.numel() 141 | self.covar_module = gpytorch.kernels.InducingPointKernel( 142 | base_covar_module, inducing_points=inducing_points, likelihood=likelihood 143 | ) 144 | 145 | def forward(self, x): 146 | mean_x = self.mean_module(x) 147 | covar_x = self.covar_module(x) 148 | return gpytorch.distributions.MultitaskMultivariateNormal.from_batch_mvn( 149 | gpytorch.distributions.MultivariateNormal(mean_x, covar_x) 150 | ) 151 | -------------------------------------------------------------------------------- /src/l4acados/models/pytorch_models/gpytorch_models/gpytorch_residual_model.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | import numpy as np 3 | import torch 4 | import gpytorch 5 | 6 | from ..pytorch_feature_selector import PyTorchFeatureSelector 7 | from ..pytorch_residual_model import PyTorchResidualModel 8 | from .gpytorch_data_processing_strategy import DataProcessingStrategy, VoidDataStrategy 9 | 10 | 11 | class GPyTorchResidualModel(PyTorchResidualModel): 12 | """Gpytorch based GP residual model which offers strategies for data collection 13 | and online learning. 14 | 15 | Args: 16 | - gp_model: A conditioned and trained instance of a gpytorch.models.ExactGP. 17 | - data_processing_strategy: Specifies how the incoming data should be handled. 18 | E.g. update the GP or saved to a file. 19 | - feature_selector: Optional feature selector if certain state dimensions are 20 | are known to be irrelevant for the GP inference. If set to None, then no selection 21 | is performed. 22 | """ 23 | 24 | def __init__( 25 | self, 26 | gp_model: gpytorch.models.ExactGP, 27 | feature_selector: Optional[PyTorchFeatureSelector] = None, 28 | data_processing_strategy: Optional[DataProcessingStrategy] = VoidDataStrategy, 29 | ): 30 | super().__init__(gp_model, feature_selector) 31 | self.gp_model = gp_model 32 | self._data_processing_strategy = data_processing_strategy 33 | 34 | def _predictions_fun_sum(self, y): 35 | """Helper function for jacobian computation 36 | 37 | sums up the mean predictions along the first dimension 38 | (i.e. along the horizon). 39 | """ 40 | self.evaluate(y, require_grad=True) 41 | return self.predictions.mean.sum(dim=0) 42 | 43 | def evaluate(self, y, require_grad=False): 44 | # NOTE(@naefjo): covar_root_decomposition=False forces linear_operator to use cholesky. 45 | # This is needed as otherwise our approach falls apart. 46 | with gpytorch.settings.fast_pred_var(), gpytorch.settings.fast_computations( 47 | covar_root_decomposition=False 48 | ): 49 | y_tensor = self.to_tensor(y) 50 | if require_grad: 51 | self.predictions = self.gp_model(self._feature_selector(y_tensor)) 52 | else: 53 | with torch.no_grad(): 54 | self.predictions = self.gp_model(self._feature_selector(y_tensor)) 55 | 56 | self.current_prediction = self.to_numpy(self.predictions.mean) 57 | 58 | # NOTE: legacy compatibility, to be removed. 59 | self.current_mean = self.current_prediction 60 | 61 | # NOTE(@naefjo): If we skipped posterior covariances, we keep the old covars in cache for hewing method. 62 | if gpytorch.settings.skip_posterior_variances.off(): 63 | self.current_variance = self.to_numpy(self.predictions.variance) 64 | 65 | return self.current_prediction 66 | 67 | def record_datapoint( 68 | self, x_input: np.array, y_target: np.array, timestamp: Optional[float] = None 69 | ) -> None: 70 | """Record one datapoint to the training dataset. 71 | 72 | Args: 73 | - x_input: (N, state_dim) input features 74 | - y_target: (N, residual_dim) array of size nw with the targets we want to predict. 75 | """ 76 | 77 | # Process datapoint and update model if needed 78 | if ( 79 | updated_gp_model := self._data_processing_strategy.process( 80 | gp_model=self.gp_model, 81 | x_input=x_input, 82 | y_target=y_target, 83 | gp_feature_selector=self._feature_selector, 84 | timestamp=timestamp, 85 | ) 86 | ) is not None: 87 | self.gp_model = updated_gp_model 88 | -------------------------------------------------------------------------------- /src/l4acados/models/pytorch_models/pytorch_feature_selector.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from typing import Optional, Union 4 | 5 | 6 | class PyTorchFeatureSelector: 7 | """Creates the Feature selector which is used for GP input dimensionality reduction. 8 | 9 | I.e. generates :math:`F \in \mathbb{R}^{gp_input_dim \times state_dim}` for :math:`y \sim gp(Fx)`, 10 | where :math:`gp_input_dim \leq state_dim`. The `input_selection_matrix` is either the full matrix 11 | or an array of size :math:`n_x` with either zeros or ones where the nonzero elements indicate 12 | the input features and the sum of nonzero elements corresponds to :math:`gp_input_dim`. 13 | 14 | Args: 15 | - input_selection: Either a vector or an (gp_input_dim, state_dim) matrix providing the mapping from 16 | full feature space to GP input feature space. If None, then this represents an identity map. 17 | - external_inputs: Optional constant feature vector which are concatenated to the extracted features. 18 | """ 19 | 20 | def __init__( 21 | self, 22 | input_selection: Optional[Union[list, np.ndarray, torch.Tensor]] = None, 23 | external_inputs: Optional[Union[list, np.ndarray, torch.Tensor]] = None, 24 | device="cpu", 25 | ) -> torch.Tensor: 26 | if input_selection is not None: 27 | if isinstance(input_selection, (list, np.ndarray)): 28 | input_selection = torch.Tensor(input_selection) 29 | 30 | if input_selection.ndim == 1: 31 | input_selection_matrix = get_input_selection_matrix(input_selection) 32 | elif input_selection.ndim == 2: 33 | input_selection_matrix = input_selection 34 | else: 35 | raise ValueError( 36 | f"input_selection must be 1D or 2D but has shape {input_selection.size()}" 37 | ) 38 | 39 | self._input_selection_matrix = input_selection_matrix.to(device=device) 40 | 41 | else: 42 | self._input_selection_matrix = None 43 | 44 | if external_inputs is not None: 45 | if isinstance(external_inputs, (list, np.ndarray)): 46 | self.external_inputs = torch.Tensor(external_inputs, device=device) 47 | elif isinstance(external_inputs, torch.Tensor): 48 | self.external_inputs = external_inputs.to(device=device) 49 | else: 50 | raise RuntimeError( 51 | f"Unexpected type {type(external_inputs)} for constant_features." 52 | f"Supported types are list, np.array, torch.tensor." 53 | ) 54 | 55 | if self.external_inputs.ndim == 1: 56 | self.external_inputs = self.external_inputs.unsqueeze(-1) 57 | elif self.external_inputs.ndim != 2: 58 | raise ValueError("constant_features has invalid number of dimensions") 59 | 60 | else: 61 | self.external_inputs = None 62 | 63 | def __str__(self) -> str: 64 | return str(self._input_selection_matrix) 65 | 66 | def __call__(self, x_input: torch.Tensor) -> torch.Tensor: 67 | """Extracts GP features from the full state vector and adds constant features if available. 68 | 69 | Args: 70 | - x_input in (N, state_dim) 71 | 72 | Returns: 73 | - extracted features in (N, gp_input_dim) 74 | """ 75 | 76 | if self._input_selection_matrix is not None: 77 | features = torch.matmul(x_input, self._input_selection_matrix.T) 78 | else: 79 | features = x_input 80 | 81 | if self.external_inputs is not None: 82 | features = torch.cat( 83 | (features, self.external_inputs[: x_input.size(0), :]), dim=-1 84 | ) 85 | 86 | return features 87 | 88 | 89 | def get_input_selection_matrix( 90 | input_selection: Union[list, np.ndarray, torch.Tensor] 91 | ) -> torch.Tensor: 92 | """Generate an input selection matrix from a 1D vector. 93 | 94 | Args: 95 | - input_selection: 1D selection vector of type list, np.ndarray or torch.Tensor. 96 | The nonzero elements indicate the selected input features. 97 | 98 | Returns: 99 | - input_selection_matrix: (gp_input_dim, state_dim) 2D matrix of type torch.Tensor. 100 | Provides a mapping from the full feature space to the selected GP input feature space. 101 | """ 102 | if isinstance(input_selection, (list, np.ndarray)): 103 | input_selection = torch.Tensor(input_selection) 104 | 105 | if input_selection.ndim != 1: 106 | raise ValueError("input_selection must be a 1D vector") 107 | 108 | input_selection_matrix = torch.diag(input_selection)[~(input_selection == 0), :] 109 | 110 | return input_selection_matrix 111 | -------------------------------------------------------------------------------- /src/l4acados/models/pytorch_models/pytorch_residual_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from typing import Optional 3 | 4 | from l4acados.models import ResidualModel 5 | from l4acados.models.pytorch_models.pytorch_feature_selector import ( 6 | PyTorchFeatureSelector, 7 | ) 8 | from l4acados.models.pytorch_models.pytorch_utils import to_numpy, to_tensor 9 | 10 | 11 | class PyTorchResidualModel(ResidualModel): 12 | """Basic PyTorch residual model class. 13 | 14 | Args: 15 | - model: A torch.nn.Module model. 16 | - feature_selector: Optional feature selector mapping controller states and inputs (x,u) to model inputs. 17 | If set to None, then no selection is performed. 18 | """ 19 | 20 | def __init__( 21 | self, 22 | model: torch.nn.Module, 23 | feature_selector: Optional[PyTorchFeatureSelector] = None, 24 | ): 25 | self.model = model 26 | self._feature_selector = ( 27 | feature_selector 28 | if feature_selector is not None 29 | else PyTorchFeatureSelector() 30 | ) 31 | self.device = next(model.parameters()).device 32 | self.to_numpy = lambda T: to_numpy(T, self.device.type) 33 | self.to_tensor = lambda X: to_tensor(X, self.device.type) 34 | 35 | def _predictions_fun_sum(self, y): 36 | """Helper function for jacobian computation 37 | 38 | sums up the mean predictions along the first dimension 39 | (i.e. along the horizon). 40 | """ 41 | self.evaluate(y, require_grad=True) 42 | return self.predictions.sum(dim=0) 43 | 44 | def evaluate(self, y, require_grad=False): 45 | y_tensor = self.to_tensor(y) 46 | if require_grad: 47 | self.predictions = self.model(self._feature_selector(y_tensor)) 48 | else: 49 | with torch.no_grad(): 50 | self.predictions = self.model(self._feature_selector(y_tensor)) 51 | 52 | self.current_prediction = self.to_numpy(self.predictions) 53 | return self.current_prediction 54 | 55 | def jacobian(self, y): 56 | y_tensor = self.to_tensor(y) 57 | mean_dy = torch.autograd.functional.jacobian( 58 | self._predictions_fun_sum, y_tensor 59 | ) 60 | return self.to_numpy(mean_dy) 61 | 62 | def value_and_jacobian(self, y): 63 | """Computes the necessary values for GPMPC 64 | 65 | Args: 66 | - x_input: (N, state_dim) tensor 67 | 68 | Returns: 69 | - mean: (N, residual_dim) tensor 70 | - mean_dy: (residual_dim, N, state_dim) tensor 71 | - covariance: (N, residual_dim) tensor 72 | """ 73 | self.current_prediction_dy = self.jacobian(y) 74 | return self.current_prediction, self.current_prediction_dy 75 | -------------------------------------------------------------------------------- /src/l4acados/models/pytorch_models/pytorch_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | 5 | def to_tensor(arr: np.ndarray, device: str) -> torch.Tensor: 6 | """ 7 | Converts a numpy array to a torch tensor on CPU/GPU. 8 | """ 9 | if device == "cuda": 10 | return torch.Tensor(arr).cuda() 11 | elif device == "cpu": 12 | return torch.Tensor(arr) 13 | else: 14 | raise RuntimeError(f"Device type {device} unkown.") 15 | 16 | 17 | def to_numpy(t: torch.Tensor, device: str) -> np.ndarray: 18 | """ 19 | Converts a torch tensor on CPU/GPU to a numpy array 20 | """ 21 | 22 | if device == "cuda": 23 | return t.cpu().detach().numpy() 24 | elif device == "cpu": 25 | return t.detach().numpy() 26 | else: 27 | raise RuntimeError(f"Device type {device} unkown.") 28 | -------------------------------------------------------------------------------- /src/l4acados/models/residual_model.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | import numpy as np 3 | 4 | 5 | class ResidualModel(ABC): 6 | @abstractmethod 7 | def value_and_jacobian(y: np.array): 8 | raise NotImplementedError 9 | -------------------------------------------------------------------------------- /test/acados_python_examples/getting_started/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/IntelligentControlSystems/l4acados/a56ec0184c52d832f6801ffe7e39c463d6fa1d37/test/acados_python_examples/getting_started/__init__.py -------------------------------------------------------------------------------- /test/acados_python_examples/getting_started/pendulum_model.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) The acados authors. 3 | # 4 | # This file is part of acados. 5 | # 6 | # The 2-Clause BSD License 7 | # 8 | # Redistribution and use in source and binary forms, with or without 9 | # modification, are permitted provided that the following conditions are met: 10 | # 11 | # 1. Redistributions of source code must retain the above copyright notice, 12 | # this list of conditions and the following disclaimer. 13 | # 14 | # 2. Redistributions in binary form must reproduce the above copyright notice, 15 | # this list of conditions and the following disclaimer in the documentation 16 | # and/or other materials provided with the distribution. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 21 | # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE 22 | # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 23 | # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 24 | # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 25 | # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 26 | # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 27 | # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 28 | # POSSIBILITY OF SUCH DAMAGE.; 29 | # 30 | 31 | from acados_template import AcadosModel 32 | from casadi import SX, vertcat, sin, cos 33 | 34 | 35 | def export_pendulum_ode_model(model_name="pendulum_ode") -> AcadosModel: 36 | 37 | # constants 38 | M = 1.0 # mass of the cart [kg] -> now estimated 39 | m = 0.1 # mass of the ball [kg] 40 | g = 9.81 # gravity constant [m/s^2] 41 | l = 0.8 # length of the rod [m] 42 | 43 | # set up states & controls 44 | x1 = SX.sym("x1") 45 | theta = SX.sym("theta") 46 | v1 = SX.sym("v1") 47 | dtheta = SX.sym("dtheta") 48 | 49 | x = vertcat(x1, theta, v1, dtheta) 50 | 51 | F = SX.sym("F") 52 | u = vertcat(F) 53 | 54 | # xdot 55 | x1_dot = SX.sym("x1_dot") 56 | theta_dot = SX.sym("theta_dot") 57 | v1_dot = SX.sym("v1_dot") 58 | dtheta_dot = SX.sym("dtheta_dot") 59 | 60 | xdot = vertcat(x1_dot, theta_dot, v1_dot, dtheta_dot) 61 | 62 | # dynamics 63 | cos_theta = cos(theta) 64 | sin_theta = sin(theta) 65 | denominator = M + m - m * cos_theta * cos_theta 66 | f_expl = vertcat( 67 | v1, 68 | dtheta, 69 | (-m * l * sin_theta * dtheta * dtheta + m * g * cos_theta * sin_theta + F) 70 | / denominator, 71 | ( 72 | -m * l * cos_theta * sin_theta * dtheta * dtheta 73 | + F * cos_theta 74 | + (M + m) * g * sin_theta 75 | ) 76 | / (l * denominator), 77 | ) 78 | 79 | f_impl = xdot - f_expl 80 | 81 | model = AcadosModel() 82 | 83 | model.f_impl_expr = f_impl 84 | model.f_expl_expr = f_expl 85 | model.x = x 86 | model.xdot = xdot 87 | model.u = u 88 | model.name = model_name 89 | 90 | # store meta information 91 | model.x_labels = [ 92 | "$x$ [m]", 93 | r"$\theta$ [rad]", 94 | "$v$ [m]", 95 | r"$\dot{\theta}$ [rad/s]", 96 | ] 97 | model.u_labels = ["$F$"] 98 | model.t_label = "$t$ [s]" 99 | 100 | return model 101 | -------------------------------------------------------------------------------- /test/acados_python_examples/getting_started/test_minimal_example_closed_loop.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) The acados authors. 2 | # 3 | # This file is part of acados. 4 | # 5 | # The 2-Clause BSD License 6 | # 7 | # Redistribution and use in source and binary forms, with or without 8 | # modification, are permitted provided that the following conditions are met: 9 | # 10 | # 1. Redistributions of source code must retain the above copyright notice, 11 | # this list of conditions and the following disclaimer. 12 | # 13 | # 2. Redistributions in binary form must reproduce the above copyright notice, 14 | # this list of conditions and the following disclaimer in the documentation 15 | # and/or other materials provided with the distribution. 16 | # 17 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 18 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 19 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 20 | # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE 21 | # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 22 | # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 23 | # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 24 | # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 25 | # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 26 | # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 27 | # POSSIBILITY OF SUCH DAMAGE.; 28 | # 29 | import numpy as np 30 | import scipy.linalg 31 | from casadi import vertcat 32 | from l4acados.controllers import ResidualLearningMPC 33 | import matplotlib.pyplot as plt 34 | from acados_template import AcadosOcp, AcadosOcpSolver, AcadosSimSolver 35 | 36 | from acados_python_examples.getting_started.pendulum_model import ( 37 | export_pendulum_ode_model, 38 | ) 39 | from acados_python_examples.getting_started.utils import plot_pendulum 40 | 41 | 42 | def create_ocp( 43 | x0, 44 | Fmax, 45 | N_horizon, 46 | Tf, 47 | RTI=False, 48 | model_name="pendulum_ode", 49 | rti_log_residuals=True, 50 | ): 51 | # create ocp object to formulate the OCP 52 | ocp = AcadosOcp() 53 | 54 | # set model 55 | model = export_pendulum_ode_model(model_name=model_name) 56 | ocp.model = model 57 | 58 | nx = model.x.rows() 59 | nu = model.u.rows() 60 | ny = nx + nu 61 | ny_e = nx 62 | 63 | ocp.dims.N = N_horizon 64 | 65 | # set cost module 66 | ocp.cost.cost_type = "NONLINEAR_LS" 67 | ocp.cost.cost_type_e = "NONLINEAR_LS" 68 | 69 | Q_mat = 2 * np.diag([1e3, 1e3, 1e-2, 1e-2]) 70 | R_mat = 2 * np.diag([1e-2]) 71 | 72 | ocp.cost.W = scipy.linalg.block_diag(Q_mat, R_mat) 73 | ocp.cost.W_e = Q_mat 74 | 75 | ocp.model.cost_y_expr = vertcat(model.x, model.u) 76 | ocp.model.cost_y_expr_e = model.x 77 | ocp.cost.yref = np.zeros((ny,)) 78 | ocp.cost.yref_e = np.zeros((ny_e,)) 79 | 80 | # set constraints 81 | ocp.constraints.lbu = np.array([-Fmax]) 82 | ocp.constraints.ubu = np.array([+Fmax]) 83 | 84 | ocp.constraints.x0 = x0 85 | ocp.constraints.idxbu = np.array([0]) 86 | 87 | ocp.solver_options.qp_solver = "PARTIAL_CONDENSING_HPIPM" # FULL_CONDENSING_QPOASES 88 | ocp.solver_options.hessian_approx = "GAUSS_NEWTON" 89 | ocp.solver_options.integrator_type = "IRK" 90 | ocp.solver_options.sim_method_newton_iter = 10 91 | ocp.solver_options.qp_tol = 1e-8 92 | ocp.solver_options.print_level = 0 93 | ocp.solver_options.tol = 1e-6 94 | 95 | if RTI: 96 | ocp.solver_options.nlp_solver_type = "SQP_RTI" 97 | ocp.solver_options.nlp_solver_max_iter = 1 98 | if rti_log_residuals: 99 | ocp.solver_options.rti_log_residuals = 1 100 | ocp.solver_options.rti_log_only_available_residuals = 1 101 | else: 102 | ocp.solver_options.nlp_solver_type = "SQP" 103 | # NOTE: globalization currently not supported by l4acados 104 | # ocp.solver_options.globalization = ( 105 | # "MERIT_BACKTRACKING" # turns on globalization 106 | # ) 107 | ocp.solver_options.nlp_solver_max_iter = 150 108 | # ocp.solver_options.store_iterates = True 109 | 110 | ocp.solver_options.qp_solver_cond_N = N_horizon 111 | 112 | # set prediction horizon 113 | ocp.solver_options.tf = Tf 114 | 115 | return ocp 116 | 117 | 118 | def setup_acados(ocp, json_file): 119 | acados_ocp_solver = AcadosOcpSolver(ocp, json_file=json_file) 120 | 121 | # create an integrator with the same settings as used in the OCP solver. 122 | acados_integrator = AcadosSimSolver(ocp, json_file=json_file) 123 | return acados_ocp_solver, acados_integrator 124 | 125 | 126 | def run( 127 | use_RTI=False, 128 | use_l4acados=False, 129 | plot=True, 130 | max_sqp_iter=150, 131 | use_cython=False, 132 | rti_log_residuals=True, 133 | ): 134 | 135 | x0 = np.array([0.0, np.pi, 0.0, 0.0]) 136 | Fmax = 80 137 | 138 | Tf = 0.8 139 | N_horizon = 40 140 | 141 | model_name = "pendulum_ode" 142 | solve_kwargs = {} 143 | if use_l4acados: 144 | model_name += f"_l4a" 145 | if not use_RTI: 146 | solve_kwargs = {"acados_sqp_mode": True} 147 | if use_RTI: 148 | max_sqp_iter = 1 149 | model_name += f"_rti_{max_sqp_iter}" 150 | else: 151 | model_name += f"_sqp_{max_sqp_iter}" 152 | 153 | ocp = create_ocp( 154 | x0, Fmax, N_horizon, Tf, RTI=use_RTI, rti_log_residuals=rti_log_residuals 155 | ) 156 | ocp.solver_options.nlp_solver_max_iter = max_sqp_iter 157 | 158 | solver_json = "acados_ocp_" + ocp.model.name + ".json" 159 | 160 | if use_l4acados: 161 | solver_json_ocp = "l4acados_ocp_" + ocp.model.name + ".json" 162 | solver_json_sim = "l4acados_sim_" + ocp.model.name + ".json" 163 | l4acados_solver = ResidualLearningMPC( 164 | ocp, 165 | use_cython=use_cython, 166 | path_json_ocp=solver_json_ocp, 167 | path_json_sim=solver_json_sim, 168 | ) 169 | ocp_solver = l4acados_solver 170 | integrator = AcadosSimSolver(ocp, json_file=solver_json) 171 | else: 172 | ocp_solver, integrator = setup_acados(ocp, solver_json) 173 | 174 | nx = ocp.dims.nx 175 | nu = ocp.dims.nu 176 | 177 | Nsim = 10 178 | simX = np.zeros((Nsim + 1, nx)) 179 | simU = np.zeros((Nsim, nu)) 180 | 181 | simX[0, :] = x0 182 | 183 | t_total = np.zeros((Nsim,)) 184 | nlp_residuals = np.zeros((Nsim, max_sqp_iter + 1, 4)) 185 | 186 | # do some initial iterations to start with a good initial guess 187 | # if use_RTI or max_sqp_iter == 1: 188 | # num_iter_initial = 5 189 | # for _ in range(num_iter_initial): 190 | # ocp_solver.solve_for_x0(x0_bar=x0, fail_on_nonzero_status=False) 191 | # print(ocp_solver.get_residuals()) 192 | 193 | # closed loop 194 | for i in range(Nsim): 195 | # set initial state 196 | ocp_solver.set(0, "lbx", simX[i, :]) 197 | ocp_solver.set(0, "ubx", simX[i, :]) 198 | 199 | if use_RTI and not use_l4acados: 200 | status = ocp_solver.solve(**solve_kwargs) 201 | else: 202 | status = ocp_solver.solve(**solve_kwargs) 203 | 204 | t_total[i] = ocp_solver.get_stats("time_tot") 205 | simU[i, :] = ocp_solver.get(0, "u") 206 | 207 | if rti_log_residuals: 208 | if use_RTI: 209 | res_stat = ocp_solver.get_stats("res_stat_all")[[0]] 210 | res_eq = ocp_solver.get_stats("res_eq_all")[[0]] 211 | res_ineq = ocp_solver.get_stats("res_ineq_all")[[0]] 212 | res_comp = ocp_solver.get_stats("res_comp_all")[[0]] 213 | else: 214 | res_stat = ocp_solver.get_stats("res_stat_all") 215 | res_eq = ocp_solver.get_stats("res_eq_all") 216 | res_ineq = ocp_solver.get_stats("res_ineq_all") 217 | res_comp = ocp_solver.get_stats("res_comp_all") 218 | 219 | nlp_residuals[i, : len(res_stat), 0] = res_stat 220 | nlp_residuals[i, : len(res_eq), 1] = res_eq 221 | nlp_residuals[i, : len(res_ineq), 2] = res_ineq 222 | nlp_residuals[i, : len(res_comp), 3] = res_comp 223 | # nlp_residuals[i, :] = ocp_solver.get_initial_residuals() 224 | 225 | # simulate system 226 | simX[i + 1, :] = integrator.simulate(x=simX[i, :], u=simU[i, :]) 227 | 228 | # scale to milliseconds 229 | t_total *= 1000 230 | print( 231 | f"Computation time (total) in ms: min {np.min(t_total):.3f} median {np.median(t_total):.3f} max {np.max(t_total):.3f}" 232 | ) 233 | 234 | timings = {"total": t_total} 235 | 236 | # plot results 237 | model = ocp.model 238 | 239 | if plot: 240 | plot_pendulum( 241 | np.linspace(0, (Tf / N_horizon) * Nsim, Nsim + 1), 242 | Fmax, 243 | simU, 244 | simX, 245 | latexify=False, 246 | time_label=model.t_label, 247 | x_labels=model.x_labels, 248 | u_labels=model.u_labels, 249 | plt_show=False, 250 | ) 251 | 252 | del ocp_solver 253 | 254 | return simX, simU, timings, nlp_residuals 255 | 256 | 257 | def test_minimal_example_closed_loop(): 258 | simX_l4acados, simU_l4acados, _, nlp_res_l4acados = run( 259 | use_RTI=False, use_l4acados=True 260 | ) 261 | simX_acados, simU_acados, _, nlp_res_acados = run(use_RTI=False, use_l4acados=False) 262 | simX_l4acados_1sqp, simU_l4acados_1sqp, _, nlp_res_l4acados_1sqp = run( 263 | use_RTI=False, use_l4acados=True, max_sqp_iter=1 264 | ) 265 | simX_acados_1sqp, simU_acados_1sqp, _, nlp_res_acados_1sqp = run( 266 | use_RTI=False, use_l4acados=False, max_sqp_iter=1 267 | ) 268 | simX_l4acados_rti, simU_l4acados_rti, _, nlp_res_l4acados_rti = run( 269 | use_RTI=True, use_l4acados=True 270 | ) 271 | simX_acados_rti, simU_acados_rti, _, nlp_res_acados_rti = run( 272 | use_RTI=True, use_l4acados=False 273 | ) 274 | # plt.show() 275 | 276 | atol = 1e-10 277 | rtol = 1e-6 278 | assert np.allclose(simX_l4acados, simX_acados, atol=atol, rtol=rtol) 279 | assert np.allclose(simU_l4acados, simU_acados, atol=atol, rtol=rtol) 280 | assert np.allclose(simX_l4acados_rti, simX_acados_rti, atol=atol, rtol=rtol) 281 | assert np.allclose(simU_l4acados_rti, simU_acados_rti, atol=atol, rtol=rtol) 282 | 283 | # acados check: *initial* residuals after 1 sqp iteration vs. after RTI 284 | assert np.allclose( 285 | nlp_res_acados_1sqp[:, 0, :], nlp_res_acados_rti[:, 0, :], atol=atol, rtol=rtol 286 | ) 287 | assert np.allclose(nlp_res_l4acados, nlp_res_acados, atol=atol, rtol=rtol) 288 | assert np.allclose(nlp_res_l4acados_rti, nlp_res_acados_rti, atol=atol, rtol=rtol) 289 | 290 | 291 | def test_cython_sqp_mode(): 292 | simX_l4acados, simU_l4acados, _, _ = run( 293 | use_RTI=False, 294 | use_l4acados=True, 295 | ) 296 | simX_l4acados_cy, simU_l4acados_cy, _, _ = run( 297 | use_RTI=False, 298 | use_l4acados=True, 299 | use_cython=True, 300 | rti_log_residuals=False, 301 | ) 302 | 303 | atol = 1e-10 304 | rtol = 1e-6 305 | assert np.allclose(simX_l4acados, simX_l4acados_cy, atol=atol, rtol=rtol) 306 | assert np.allclose(simU_l4acados, simU_l4acados_cy, atol=atol, rtol=rtol) 307 | # NOTE: cannot check residuals since not supports in acados Cython interface 308 | # assert np.allclose(nlp_res_l4acados, nlp_res_l4acados_cy, atol=atol, rtol=rtol) 309 | 310 | 311 | # TODO: call tests from subprocess to avoid conflicts with multiple cython solvers 312 | # def test_cython_sqp_mode_sqp1(): 313 | # simX_l4acados_sqp1, simU_l4acados_sqp1, _, _ = run_as_subprocess( 314 | # use_RTI=False, 315 | # use_l4acados=True, 316 | # max_sqp_iter=1, 317 | # ) 318 | # simX_l4acados_sqp1_cy, simU_l4acados_sqp1_cy, _, _ = run_as_subprocess( 319 | # use_RTI=False, 320 | # use_l4acados=True, 321 | # use_cython=True, 322 | # rti_log_residuals=False, 323 | # max_sqp_iter=1, 324 | # ) 325 | # atol = 1e-10 326 | # rtol = 1e-6 327 | # assert np.allclose(simX_l4acados_sqp1, simX_l4acados_sqp1_cy, atol=atol, rtol=rtol) 328 | # assert np.allclose(simU_l4acados_sqp1, simU_l4acados_sqp1_cy, atol=atol, rtol=rtol) 329 | 330 | 331 | if __name__ == "__main__": 332 | test_minimal_example_closed_loop() 333 | test_cython_sqp_mode() 334 | # test_cython_sqp_mode_sqp1() 335 | -------------------------------------------------------------------------------- /test/acados_python_examples/getting_started/utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) The acados authors. 3 | # 4 | # This file is part of acados. 5 | # 6 | # The 2-Clause BSD License 7 | # 8 | # Redistribution and use in source and binary forms, with or without 9 | # modification, are permitted provided that the following conditions are met: 10 | # 11 | # 1. Redistributions of source code must retain the above copyright notice, 12 | # this list of conditions and the following disclaimer. 13 | # 14 | # 2. Redistributions in binary form must reproduce the above copyright notice, 15 | # this list of conditions and the following disclaimer in the documentation 16 | # and/or other materials provided with the distribution. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 21 | # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE 22 | # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 23 | # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 24 | # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 25 | # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 26 | # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 27 | # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 28 | # POSSIBILITY OF SUCH DAMAGE.; 29 | # 30 | 31 | import matplotlib.pyplot as plt 32 | import numpy as np 33 | from acados_template import latexify_plot 34 | 35 | 36 | def plot_pendulum( 37 | t, 38 | u_max, 39 | U, 40 | X_true, 41 | latexify=False, 42 | plt_show=True, 43 | time_label="$t$", 44 | x_labels=None, 45 | u_labels=None, 46 | ): 47 | """ 48 | Params: 49 | t: time values of the discretization 50 | u_max: maximum absolute value of u 51 | U: arrray with shape (N_sim-1, nu) or (N_sim, nu) 52 | X_true: arrray with shape (N_sim, nx) 53 | latexify: latex style plots 54 | """ 55 | 56 | if latexify: 57 | latexify_plot() 58 | 59 | nx = X_true.shape[1] 60 | fig, axes = plt.subplots(nx + 1, 1, sharex=True) 61 | 62 | for i in range(nx): 63 | axes[i].plot(t, X_true[:, i]) 64 | axes[i].grid() 65 | if x_labels is not None: 66 | axes[i].set_ylabel(x_labels[i]) 67 | else: 68 | axes[i].set_ylabel(f"$x_{i}$") 69 | 70 | axes[-1].step(t, np.append([U[0]], U)) 71 | 72 | if u_labels is not None: 73 | axes[-1].set_ylabel(u_labels[0]) 74 | else: 75 | axes[-1].set_ylabel("$u$") 76 | 77 | axes[-1].hlines(u_max, t[0], t[-1], linestyles="dashed", alpha=0.7) 78 | axes[-1].hlines(-u_max, t[0], t[-1], linestyles="dashed", alpha=0.7) 79 | axes[-1].set_ylim([-1.2 * u_max, 1.2 * u_max]) 80 | axes[-1].set_xlim(t[0], t[-1]) 81 | axes[-1].set_xlabel(time_label) 82 | axes[-1].grid() 83 | 84 | plt.subplots_adjust(left=None, bottom=None, right=None, top=None, hspace=0.4) 85 | 86 | fig.align_ylabels() 87 | 88 | if plt_show: 89 | plt.show() 90 | -------------------------------------------------------------------------------- /test/gpytorch_model/test_gpytorch_learning_model.py: -------------------------------------------------------------------------------- 1 | import os 2 | from time import perf_counter, sleep 3 | from typing import Optional, Tuple 4 | 5 | import numpy as np 6 | import torch 7 | import gpytorch 8 | 9 | from l4acados.models import PyTorchFeatureSelector, GPyTorchResidualModel 10 | from l4acados.models.pytorch_models.gpytorch_models.gpytorch_data_processing_strategy import ( 11 | OnlineLearningStrategy, 12 | RecordDataStrategy, 13 | ) 14 | from l4acados.models.pytorch_models.gpytorch_models import gpytorch_gp 15 | 16 | 17 | def count_lines(filepath: str) -> int: 18 | with open(filepath, "r") as file: 19 | return sum(1 for line in file) 20 | 21 | 22 | def test_unconditioned_gp(num_tests: int = 6) -> None: 23 | """Test initalization, querying and updating of unconditioned gp""" 24 | add_data_times_ms = [] 25 | eval_times_ms = [] 26 | current_dir = os.getcwd() 27 | x_data_path = os.path.join(current_dir, "tmp_x_data.txt") 28 | y_data_path = os.path.join(current_dir, "tmp_y_data.txt") 29 | for curr_test_num in range(num_tests): 30 | print(f"\n\nPerforming run {curr_test_num}/{num_tests}") 31 | # Test random input size and random jacobian 32 | state_dimension = np.random.randint(1, 10) 33 | residual_dimension = np.random.randint(1, 5) 34 | input_feature_selection = np.random.randint(0, 2, size=(state_dimension)) 35 | input_feature_selection[0] = 1 # Make sure input dimension to GP is at least 1 36 | input_dimension = np.sum(input_feature_selection) 37 | 38 | if (curr_test_num % 2) == 0: 39 | input_feature_selection = None 40 | input_dimension = state_dimension 41 | 42 | enable_ard = curr_test_num > num_tests / 2 43 | 44 | hyperparameters = ( 45 | { 46 | "likelihood.task_noises": torch.rand(residual_dimension) + 1e-4, 47 | "covar_module.base_kernel.lengthscale": 10 48 | * torch.rand(residual_dimension).reshape(residual_dimension, 1, 1) 49 | + 1, 50 | "covar_module.outputscale": torch.rand(residual_dimension) + 1e-4, 51 | } 52 | if not enable_ard 53 | else { 54 | "likelihood.task_noises": torch.rand(residual_dimension) + 1e-4, 55 | "covar_module.base_kernel.lengthscale": 10 56 | * torch.rand((residual_dimension, 1, input_dimension)) 57 | + 1, 58 | "covar_module.outputscale": torch.rand(residual_dimension) + 1e-4, 59 | } 60 | ) 61 | 62 | input_selection = PyTorchFeatureSelector(input_feature_selection) 63 | 64 | gpytorch_gp_model = gpytorch_gp.BatchIndependentMultitaskGPModel( 65 | train_x=None, 66 | train_y=None, 67 | likelihood=gpytorch.likelihoods.MultitaskGaussianLikelihood( 68 | num_tasks=residual_dimension, 69 | ), 70 | residual_dimension=residual_dimension, 71 | input_dimension=input_dimension, 72 | use_ard=enable_ard, 73 | ) 74 | 75 | gpytorch_gp_model.initialize(**hyperparameters) 76 | 77 | gpytorch_gp_model.eval() 78 | gpytorch_gp_model.likelihood.eval() 79 | 80 | num_record_steps = 20 81 | gp = GPyTorchResidualModel( 82 | gp_model=gpytorch_gp_model, 83 | feature_selector=input_selection, 84 | data_processing_strategy=RecordDataStrategy( 85 | x_data_path, 86 | y_data_path, 87 | buffer_size=num_record_steps, 88 | # NOTE: num_record_steps needs to be divisible by buffer size; otherwise wrong file size 89 | ), 90 | ) 91 | 92 | for _ in range(num_record_steps): 93 | x_rand_data = np.random.rand(state_dimension) 94 | y_rand_data = np.random.rand(residual_dimension) 95 | 96 | time_before = perf_counter() 97 | gp.record_datapoint(x_rand_data, y_rand_data) 98 | add_data_times_ms.append((perf_counter() - time_before) * 1e3) 99 | 100 | num_query_points = np.random.randint(1, 40) 101 | x_rand_eval = torch.Tensor( 102 | np.random.rand(num_query_points, state_dimension) 103 | ) 104 | 105 | time_before = perf_counter() 106 | val, jac = gp.value_and_jacobian(x_rand_eval) 107 | assert val.shape == torch.Size([num_query_points, residual_dimension]) 108 | assert jac.shape == torch.Size( 109 | [residual_dimension, num_query_points, state_dimension] 110 | ) 111 | eval_times_ms.append((perf_counter() - time_before) * 1e3) 112 | 113 | print( 114 | f"data adding avg over {len(add_data_times_ms)} " 115 | f"trials: {np.average(add_data_times_ms):.4f} ms." 116 | ) 117 | print( 118 | f"gp eval times over {len(eval_times_ms)} " 119 | f"trials: {np.average(eval_times_ms):.4f} ms." 120 | ) 121 | 122 | # wait for threads to finish 123 | sleep(5) 124 | 125 | num_lines_x_data = count_lines(x_data_path) 126 | num_lines_y_data = count_lines(y_data_path) 127 | 128 | assert num_lines_x_data == num_tests * num_record_steps 129 | assert num_lines_y_data == num_tests * num_record_steps 130 | 131 | os.remove(x_data_path) 132 | os.remove(y_data_path) 133 | 134 | 135 | def generate_fake_data( 136 | filepath_x: str, 137 | filepath_y: str, 138 | state_size: int, 139 | residual_size: int, 140 | num_datapoints: Optional[int] = 800, 141 | ) -> None: 142 | os.remove(filepath_x) if os.path.exists(filepath_x) else None 143 | os.remove(filepath_y) if os.path.exists(filepath_y) else None 144 | x_data = np.random.rand(num_datapoints, state_size) 145 | y_data = np.random.rand(num_datapoints, residual_size) 146 | np.savetxt(filepath_x, x_data, delimiter=",") 147 | np.savetxt(filepath_y, y_data, delimiter=",") 148 | 149 | 150 | def load_data(x_data_path: str, y_data_path: str) -> Tuple[torch.Tensor, torch.Tensor]: 151 | try: 152 | with open(x_data_path, "r") as file: 153 | train_x_tensor = torch.from_numpy(np.genfromtxt(file, delimiter=",")).type( 154 | torch.float32 155 | ) 156 | except FileNotFoundError: 157 | print(f"File {x_data_path} does not exist.") 158 | except Exception as e: 159 | raise e 160 | 161 | # If the data only has one feature, we need to add a batch dimension 162 | if train_x_tensor.dim() == 1: 163 | train_x_tensor = torch.unsqueeze(train_x_tensor, -1) 164 | 165 | try: 166 | with open(y_data_path, "r") as file: 167 | train_y_tensor = torch.from_numpy(np.genfromtxt(file, delimiter=",")).type( 168 | torch.float32 169 | ) 170 | except FileNotFoundError: 171 | print(f"File {y_data_path} does not exist.") 172 | except Exception as e: 173 | raise e 174 | 175 | # If the data only has one feature, we need to add a batch dimension 176 | if train_y_tensor.dim() == 1: 177 | train_y_tensor = torch.unsqueeze(train_y_tensor, -1) 178 | 179 | return train_x_tensor, train_y_tensor 180 | 181 | 182 | def test_load_gp_from_file(num_tests: int = 6) -> None: 183 | """Test initialization of the GP from a file""" 184 | add_data_times_ms = [] 185 | eval_times_ms = [] 186 | current_dir = os.getcwd() 187 | x_data_path = os.path.join(current_dir, "tmp_x_data.txt") 188 | y_data_path = os.path.join(current_dir, "tmp_y_data.txt") 189 | 190 | for curr_test_num in range(num_tests): 191 | print(f"\n\nPerforming run {curr_test_num}/{num_tests}") 192 | # Test random input size and random jacobian 193 | state_dimension = np.random.randint(1, 10) 194 | residual_dimension = np.random.randint(1, 5) 195 | input_feature_selection = np.random.randint(0, 2, size=(state_dimension)) 196 | input_feature_selection[0] = 1 # Make sure input dimension to GP is at least 1 197 | input_dimension = np.sum(input_feature_selection) 198 | 199 | if (curr_test_num % 2) == 0: 200 | input_feature_selection = None 201 | input_dimension = state_dimension 202 | 203 | enable_ard = curr_test_num > num_tests / 2 204 | hyperparameters = ( 205 | { 206 | "likelihood.task_noises": torch.rand(residual_dimension) + 1e-4, 207 | "covar_module.base_kernel.lengthscale": 10 208 | * torch.rand(residual_dimension).reshape(residual_dimension, 1, 1) 209 | + 1, 210 | "covar_module.outputscale": torch.rand(residual_dimension) + 1e-4, 211 | } 212 | if not enable_ard 213 | else { 214 | "likelihood.task_noises": torch.rand(residual_dimension) + 1e-4, 215 | "covar_module.base_kernel.lengthscale": 10 216 | * torch.rand((residual_dimension, 1, input_dimension)) 217 | + 1, 218 | "covar_module.outputscale": torch.rand(residual_dimension) + 1e-4, 219 | } 220 | ) 221 | 222 | num_fake_datapoints = np.random.randint(100, 1000) 223 | 224 | generate_fake_data( 225 | x_data_path, 226 | y_data_path, 227 | state_dimension, 228 | residual_dimension, 229 | num_fake_datapoints, 230 | ) 231 | 232 | max_num_datapoints = ( 233 | None 234 | if np.random.randint(2) == 0 235 | else np.random.randint(50, num_fake_datapoints) 236 | ) 237 | print( 238 | f"Setting max num datapoints to {max_num_datapoints} out of " 239 | f"{num_fake_datapoints} datapoints." 240 | ) 241 | 242 | train_x_tensor, train_y_tensor = load_data(x_data_path, y_data_path) 243 | 244 | # Generate a random permutation of indices 245 | if np.random.randint(2): 246 | print("Permuting dataset") 247 | # Set manual seed for reproducibility 248 | indices = torch.randperm(train_x_tensor.size(0)) 249 | 250 | train_x_tensor = train_x_tensor[indices] 251 | train_y_tensor = train_y_tensor[indices] 252 | 253 | if max_num_datapoints is not None: 254 | train_x_tensor = train_x_tensor[:max_num_datapoints, :] 255 | train_y_tensor = train_y_tensor[:max_num_datapoints, :] 256 | 257 | print( 258 | f"Loaded training data sucessfully with x and y shapes " 259 | f"{train_x_tensor.shape} {train_y_tensor.shape}" 260 | ) 261 | 262 | input_selection = PyTorchFeatureSelector(input_feature_selection) 263 | 264 | gpytorch_model = gpytorch_gp.BatchIndependentMultitaskGPModel( 265 | input_selection(train_x_tensor), 266 | train_y_tensor, 267 | gpytorch.likelihoods.MultitaskGaussianLikelihood( 268 | num_tasks=train_y_tensor.shape[1], 269 | has_task_noise=True, 270 | has_global_noise=False, 271 | ), 272 | use_ard=enable_ard, 273 | ) 274 | 275 | gpytorch_model.initialize(**hyperparameters) 276 | 277 | gpytorch_model.eval() 278 | gpytorch_model.likelihood.eval() 279 | 280 | gp = GPyTorchResidualModel( 281 | gp_model=gpytorch_model, 282 | feature_selector=input_selection, 283 | data_processing_strategy=RecordDataStrategy(x_data_path, y_data_path), 284 | ) 285 | 286 | for _ in range(20): 287 | x_rand_data = np.random.rand(state_dimension) 288 | y_rand_data = np.random.rand(residual_dimension) 289 | 290 | time_before = perf_counter() 291 | gp.record_datapoint(x_rand_data, y_rand_data) 292 | add_data_times_ms.append((perf_counter() - time_before) * 1e3) 293 | 294 | num_query_points = np.random.randint(1, 40) 295 | x_rand_eval = torch.Tensor( 296 | np.random.rand(num_query_points, state_dimension) 297 | ) 298 | 299 | time_before = perf_counter() 300 | val, jac = gp.value_and_jacobian(x_rand_eval) 301 | assert val.shape == torch.Size([num_query_points, residual_dimension]) 302 | assert jac.shape == torch.Size( 303 | [residual_dimension, num_query_points, state_dimension] 304 | ) 305 | eval_times_ms.append((perf_counter() - time_before) * 1e3) 306 | 307 | print( 308 | f"data adding avg over {len(add_data_times_ms)} " 309 | f"trials: {np.average(add_data_times_ms):.4f} ms." 310 | ) 311 | print( 312 | f"gp eval times over {len(eval_times_ms)} " 313 | f"trials: {np.average(eval_times_ms):.4f} ms." 314 | ) 315 | 316 | # wait for threads to finish 317 | sleep(5) 318 | 319 | os.remove(x_data_path) 320 | os.remove(y_data_path) 321 | 322 | 323 | def test_incorporate_new_data(num_tests: int = 6): 324 | add_data_times_ms = [] 325 | eval_times_ms = [] 326 | 327 | max_points_online = 20 328 | sim_steps_list = np.linspace( 329 | max_points_online - 5, max_points_online + 5, num_tests, dtype=int 330 | ) 331 | 332 | for i in range(num_tests): 333 | # Test random input size and random jacobian 334 | state_dimension = np.random.randint(1, 10) 335 | residual_dimension = np.random.randint(1, 5) 336 | input_feature_selection = np.random.randint(0, 2, size=(state_dimension)) 337 | input_feature_selection[0] = 1 # Make sure input dimension to GP is at least 1 338 | input_dimension = np.sum(input_feature_selection) 339 | 340 | hyperparameters = { 341 | "likelihood.task_noises": torch.rand(residual_dimension) + 1e-4, 342 | "covar_module.base_kernel.lengthscale": 10 343 | * torch.rand(residual_dimension).reshape(residual_dimension, 1, 1) 344 | + 1, 345 | "covar_module.outputscale": torch.rand(residual_dimension) + 1e-4, 346 | } 347 | 348 | input_selection = PyTorchFeatureSelector(input_feature_selection) 349 | 350 | gpytorch_gp_model = gpytorch_gp.BatchIndependentMultitaskGPModel( 351 | train_x=None, 352 | train_y=None, 353 | likelihood=gpytorch.likelihoods.MultitaskGaussianLikelihood( 354 | num_tasks=residual_dimension, 355 | ), 356 | input_dimension=input_dimension, 357 | residual_dimension=residual_dimension, 358 | ) 359 | 360 | gpytorch_gp_model.initialize(**hyperparameters) 361 | 362 | gpytorch_gp_model.eval() 363 | gpytorch_gp_model.likelihood.eval() 364 | 365 | gp = GPyTorchResidualModel( 366 | gp_model=gpytorch_gp_model, 367 | feature_selector=input_selection, 368 | data_processing_strategy=OnlineLearningStrategy(max_points_online), 369 | ) 370 | 371 | gp.value_and_jacobian(torch.rand(1, state_dimension)) 372 | 373 | for _ in range(sim_steps_list[i]): 374 | x_rand_data = np.random.rand(state_dimension) 375 | y_rand_data = np.random.rand(residual_dimension) 376 | 377 | time_before = perf_counter() 378 | gp.record_datapoint(x_rand_data, y_rand_data) 379 | add_data_times_ms.append((perf_counter() - time_before) * 1e3) 380 | 381 | num_query_points = np.random.randint(1, 40) 382 | x_rand_eval = torch.Tensor( 383 | np.random.rand(num_query_points, state_dimension) 384 | ) 385 | 386 | time_before = perf_counter() 387 | val, jac = gp.value_and_jacobian(x_rand_eval) 388 | assert val.shape == torch.Size([num_query_points, residual_dimension]) 389 | assert jac.shape == torch.Size( 390 | [residual_dimension, num_query_points, state_dimension] 391 | ) 392 | eval_times_ms.append((perf_counter() - time_before) * 1e3) 393 | 394 | assert gp.gp_model.train_inputs[0].shape == torch.Size( 395 | [min(sim_steps_list[i], max_points_online), sum(input_feature_selection)] 396 | ) 397 | assert gp.gp_model.train_targets.shape == torch.Size( 398 | [min(sim_steps_list[i], max_points_online), residual_dimension] 399 | ) 400 | 401 | print( 402 | f"data adding avg over {len(add_data_times_ms)} " 403 | f"trials: {np.average(add_data_times_ms):.4f} ms." 404 | ) 405 | print( 406 | f"gp eval times over {len(eval_times_ms)} " 407 | f"trials: {np.average(eval_times_ms):.4f} ms." 408 | ) 409 | 410 | 411 | def train_gp_model( 412 | gp_model, torch_seed=None, training_iterations=10, learning_rate=1.0 413 | ): 414 | if torch_seed is not None: 415 | torch.manual_seed(torch_seed) 416 | 417 | likelihood = gp_model.likelihood 418 | train_x = gp_model.train_inputs[0] 419 | train_y = gp_model.train_targets 420 | 421 | # Find optimal model hyperparameters 422 | gp_model.train() 423 | likelihood.train() 424 | 425 | # Use the adam optimizer 426 | optimizer = torch.optim.Adam( 427 | filter(lambda p: p.requires_grad, gp_model.parameters()), lr=learning_rate 428 | ) # Includes GaussianLikelihood parameters 429 | 430 | # "Loss" for GPs - the marginal log likelihood 431 | mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp_model) 432 | 433 | prev_loss = np.inf 434 | num_loss_below_threshold = 0 435 | 436 | for i in range(training_iterations): 437 | optimizer.zero_grad() 438 | output = likelihood(gp_model(train_x)) 439 | loss = -mll(output, train_y.reshape((train_y.numel(),))) 440 | loss = loss.sum() 441 | 442 | num_loss_below_threshold = ( 443 | num_loss_below_threshold + 1 444 | if torch.any(torch.abs(loss - prev_loss)) < 5e-5 445 | else 0 446 | ) 447 | 448 | if num_loss_below_threshold > 5: 449 | print(f"stopping GP optimization early after {i} iterations.") 450 | break 451 | 452 | prev_loss = loss 453 | 454 | loss.backward() 455 | if (i + 1) % 20 == 0: 456 | print("Iter %d/%d - Loss: %.3f" % (i + 1, training_iterations, loss.item())) 457 | 458 | optimizer.step() 459 | 460 | gp_model.eval() 461 | likelihood.eval() 462 | 463 | 464 | def test_inducing_point_gp_from_file(num_tests: int = 6) -> None: 465 | """Test initialization of the GP from a file""" 466 | add_data_times_ms = [] 467 | eval_times_ms = [] 468 | current_dir = os.getcwd() 469 | x_data_path = os.path.join(current_dir, "tmp_x_data.txt") 470 | y_data_path = os.path.join(current_dir, "tmp_y_data.txt") 471 | 472 | for curr_test_num in range(num_tests): 473 | print(f"\n\nPerforming run {curr_test_num}/{num_tests}") 474 | # Test random input size and random jacobian 475 | state_dimension = np.random.randint(1, 10) 476 | residual_dimension = np.random.randint(1, 5) 477 | input_feature_selection = np.random.randint(0, 2, size=(state_dimension)) 478 | input_feature_selection[0] = 1 # Make sure input dimension to GP is at least 1 479 | input_dimension = np.sum(input_feature_selection) 480 | 481 | if (curr_test_num % 2) == 0: 482 | input_feature_selection = None 483 | input_dimension = state_dimension 484 | 485 | num_fake_datapoints = np.random.randint(100, 1000) 486 | 487 | generate_fake_data( 488 | x_data_path, 489 | y_data_path, 490 | state_dimension, 491 | residual_dimension, 492 | num_fake_datapoints, 493 | ) 494 | 495 | max_num_datapoints = ( 496 | None 497 | if np.random.randint(2) == 0 498 | else np.random.randint(50, num_fake_datapoints) 499 | ) 500 | print( 501 | f"Setting max num datapoints to {max_num_datapoints} out of " 502 | f"{num_fake_datapoints} datapoints." 503 | ) 504 | 505 | num_iducing_points = 50 506 | 507 | train_x_tensor, train_y_tensor = load_data(x_data_path, y_data_path) 508 | 509 | # Generate a random permutation of indices 510 | if np.random.randint(2): 511 | print("Permuting dataset") 512 | # Set manual seed for reproducibility 513 | indices = torch.randperm(train_x_tensor.size(0)) 514 | 515 | train_x_tensor = train_x_tensor[indices] 516 | train_y_tensor = train_y_tensor[indices] 517 | 518 | if max_num_datapoints is not None: 519 | train_x_tensor = train_x_tensor[:max_num_datapoints, :] 520 | train_y_tensor = train_y_tensor[:max_num_datapoints, :] 521 | 522 | print( 523 | f"Loaded training data sucessfully with x and y shapes " 524 | f"{train_x_tensor.shape} {train_y_tensor.shape}" 525 | ) 526 | 527 | input_selection = PyTorchFeatureSelector(input_feature_selection) 528 | 529 | gpytorch_model = gpytorch_gp.BatchIndependentInducingPointGpModel( 530 | input_selection(train_x_tensor), 531 | train_y_tensor, 532 | gpytorch.likelihoods.MultitaskGaussianLikelihood( 533 | num_tasks=train_y_tensor.shape[1], 534 | has_task_noise=True, 535 | has_global_noise=False, 536 | ), 537 | num_iducing_points, 538 | use_ard=False, 539 | ) 540 | 541 | gpytorch_model.train() 542 | gpytorch_model.likelihood.train() 543 | train_gp_model(gpytorch_model) 544 | gpytorch_model.eval() 545 | gpytorch_model.likelihood.eval() 546 | 547 | print(gpytorch_model.covar_module.inducing_points) 548 | print(gpytorch_model.covar_module.base_kernel.outputscale) 549 | print(gpytorch_model.covar_module.base_kernel.base_kernel.lengthscale) 550 | 551 | gp = GPyTorchResidualModel( 552 | gp_model=gpytorch_model, 553 | feature_selector=input_selection, 554 | data_processing_strategy=RecordDataStrategy(x_data_path, y_data_path), 555 | ) 556 | 557 | for _ in range(20): 558 | x_rand_data = np.random.rand(state_dimension) 559 | y_rand_data = np.random.rand(residual_dimension) 560 | 561 | time_before = perf_counter() 562 | gp.record_datapoint(x_rand_data, y_rand_data) 563 | add_data_times_ms.append((perf_counter() - time_before) * 1e3) 564 | 565 | num_query_points = np.random.randint(1, 40) 566 | x_rand_eval = torch.Tensor( 567 | np.random.rand(num_query_points, state_dimension) 568 | ) 569 | 570 | time_before = perf_counter() 571 | val, jac = gp.value_and_jacobian(x_rand_eval) 572 | assert val.shape == torch.Size([num_query_points, residual_dimension]) 573 | assert jac.shape == torch.Size( 574 | [residual_dimension, num_query_points, state_dimension] 575 | ) 576 | eval_times_ms.append((perf_counter() - time_before) * 1e3) 577 | 578 | print( 579 | f"data adding avg over {len(add_data_times_ms)} " 580 | f"trials: {np.average(add_data_times_ms):.4f} ms." 581 | ) 582 | print( 583 | f"gp eval times over {len(eval_times_ms)} " 584 | f"trials: {np.average(eval_times_ms):.4f} ms." 585 | ) 586 | 587 | os.remove(x_data_path) 588 | os.remove(y_data_path) 589 | 590 | 591 | if __name__ == "__main__": 592 | torch.random.manual_seed(42) 593 | np.random.seed(42) 594 | test_unconditioned_gp() 595 | print(5 * "\n") 596 | test_load_gp_from_file() 597 | print(5 * "\n") 598 | test_incorporate_new_data() 599 | print(5 * "\n") 600 | test_inducing_point_gp_from_file() 601 | -------------------------------------------------------------------------------- /test/l4casadi_comparison/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/IntelligentControlSystems/l4acados/a56ec0184c52d832f6801ffe7e39c463d6fa1d37/test/l4casadi_comparison/__init__.py -------------------------------------------------------------------------------- /test/l4casadi_comparison/l4casadi_with_acados.py: -------------------------------------------------------------------------------- 1 | import casadi as cs 2 | import numpy as np 3 | import torch 4 | import l4casadi as l4c 5 | from acados_template import AcadosOcpSolver, AcadosOcp, AcadosModel 6 | import time 7 | 8 | COST = "LINEAR_LS" # NONLINEAR_LS 9 | 10 | 11 | class MultiLayerPerceptron(torch.nn.Module): 12 | def __init__(self, n_inputs=2, hidden_layers=2, hidden_size=512, n_outputs=1): 13 | super().__init__() 14 | 15 | self.input_layer = torch.nn.Linear(n_inputs, hidden_size) 16 | 17 | all_hidden_layers = [] 18 | for i in range(hidden_layers): 19 | all_hidden_layers.append(torch.nn.Linear(hidden_size, hidden_size)) 20 | 21 | self.hidden_layer = torch.nn.ModuleList(all_hidden_layers) 22 | self.out_layer = torch.nn.Linear(hidden_size, n_outputs) 23 | 24 | # Model is not trained -- setting output to zero 25 | with torch.no_grad(): 26 | self.out_layer.bias.fill_(0.0) 27 | self.out_layer.weight.fill_(0.0) 28 | 29 | def forward(self, x): 30 | x = self.input_layer(x) 31 | for layer in self.hidden_layer: 32 | x = torch.tanh(layer(x)) 33 | x = self.out_layer(x) 34 | return x 35 | 36 | 37 | def train_nn(nn_params_path="nn_params.pt"): 38 | # torch set seed 39 | torch.manual_seed(12345) 40 | 41 | train_inputs = np.linspace(-2 * np.pi, 2 * np.pi, int(1e5)) 42 | train_targets = 10 * np.sin(0.25 * train_inputs) ** 2 + 0.01 * np.random.randn( 43 | *train_inputs.shape 44 | ) 45 | 46 | train_inputs = torch.tensor(train_inputs, dtype=torch.float32).unsqueeze(1) 47 | train_targets = torch.tensor(train_targets, dtype=torch.float32).unsqueeze(1) 48 | 49 | train_set = torch.utils.data.DataLoader( 50 | torch.utils.data.TensorDataset(train_inputs, train_targets), 51 | batch_size=500, 52 | shuffle=True, 53 | ) 54 | 55 | nn_fun = MultiLayerPerceptron( 56 | n_inputs=1, hidden_layers=2, hidden_size=64, n_outputs=1 57 | ) 58 | nn_fun.train() 59 | 60 | # train the neural network 61 | optimizer = torch.optim.SGD(nn_fun.parameters()) 62 | loss_fn = torch.nn.MSELoss() 63 | 64 | for epoch in range(50): 65 | 66 | for i, data in enumerate(train_set): 67 | train_inputs_batch, train_targets_batch = data 68 | optimizer.zero_grad() 69 | outputs_batch = nn_fun(train_inputs_batch) 70 | loss = loss_fn(outputs_batch, train_targets_batch) 71 | loss.backward() 72 | optimizer.step() 73 | 74 | print(f"Epoch {epoch}, loss {loss.item()}") 75 | 76 | nn_fun.eval() 77 | 78 | test_inputs = torch.tensor( 79 | np.linspace(-2 * np.pi, 2 * np.pi, int(1e3)), dtype=torch.float32 80 | ).unsqueeze(1) 81 | 82 | test_outputs = nn_fun(test_inputs) 83 | 84 | print( 85 | f"Test loss: {loss_fn(test_outputs, 10 * torch.sin(0.25 * test_inputs) ** 2)}" 86 | ) 87 | 88 | nn_params = nn_fun.state_dict() 89 | torch.save(nn_params, nn_params_path) 90 | 91 | 92 | class DoubleIntegratorWithLearnedDynamics: 93 | def __init__(self, learned_dyn=None, name="wr"): 94 | self.learned_dyn = learned_dyn 95 | self.name = name 96 | 97 | def model(self): 98 | s = cs.MX.sym("s", 1) 99 | s_dot = cs.MX.sym("s_dot", 1) 100 | s_dot_dot = cs.MX.sym("s_dot_dot", 1) 101 | u = cs.MX.sym("u", 1) 102 | x = cs.vertcat(s, s_dot) 103 | x_dot = cs.vertcat(s_dot, s_dot_dot) 104 | x_l4casadi = s 105 | 106 | if self.learned_dyn is None: 107 | res_model = cs.MX.zeros(1, 1) 108 | else: 109 | res_model = self.learned_dyn(x_l4casadi) 110 | 111 | B = cs.DM([[1.0], [1.0]]) 112 | 113 | # f_expl = cs.vertcat(s_dot, u - cs.sin(3.1 * s) ** 2) + B.T @ res_model 114 | f_expl = cs.vertcat(s_dot, u) + B @ res_model 115 | 116 | x_start = np.zeros((2,)) 117 | 118 | # store to struct 119 | model = cs.types.SimpleNamespace() 120 | model.x = x 121 | model.xdot = x_dot 122 | model.u = u 123 | model.z = cs.vertcat([]) 124 | model.p = cs.vertcat([]) 125 | model.f_expl = f_expl 126 | model.x_start = x_start 127 | model.constraints = cs.vertcat([]) 128 | model.name = self.name 129 | 130 | return model 131 | 132 | 133 | class MPC: 134 | def __init__( 135 | self, 136 | model, 137 | N, 138 | external_shared_lib_dir=None, 139 | external_shared_lib_name=None, 140 | num_threads_acados_openmp=1, 141 | ): 142 | self.N = N 143 | self.model = model 144 | self.external_shared_lib_dir = external_shared_lib_dir 145 | self.external_shared_lib_name = external_shared_lib_name 146 | self.num_threads_acados_openmp = num_threads_acados_openmp 147 | 148 | @property 149 | def solver(self): 150 | return AcadosOcpSolver(self.ocp()) 151 | 152 | def ocp(self): 153 | model = self.model 154 | 155 | t_horizon = 1.0 156 | N = self.N 157 | 158 | # Get model 159 | model_ac = self.acados_model(model=model) 160 | model_ac.p = model.p 161 | 162 | # Dimensions 163 | nx = 2 164 | nu = 1 165 | ny = 1 166 | 167 | # Create OCP object to formulate the optimization 168 | ocp = AcadosOcp() 169 | ocp.model = model_ac 170 | 171 | ocp.dims.N = N 172 | ocp.dims.nx = nx 173 | ocp.dims.nu = nu 174 | ocp.dims.ny = ny 175 | ocp.solver_options.tf = t_horizon 176 | 177 | if COST == "LINEAR_LS": 178 | # Initialize cost function 179 | ocp.cost.cost_type = "LINEAR_LS" 180 | ocp.cost.cost_type_e = "LINEAR_LS" 181 | 182 | ocp.cost.W = np.array([[1.0]]) 183 | 184 | ocp.cost.Vx = np.zeros((ny, nx)) 185 | ocp.cost.Vx[0, 0] = 1.0 186 | ocp.cost.Vu = np.zeros((ny, nu)) 187 | ocp.cost.Vz = np.array([[]]) 188 | ocp.cost.Vx_e = np.zeros((ny, nx)) 189 | 190 | l4c_y_expr = None 191 | else: 192 | ocp.cost.cost_type = "NONLINEAR_LS" 193 | ocp.cost.cost_type_e = "NONLINEAR_LS" 194 | 195 | x = ocp.model.x 196 | 197 | ocp.cost.W = np.array([[1.0]]) 198 | 199 | # Trivial PyTorch index 0 200 | l4c_y_expr = l4c.L4CasADi( 201 | lambda x: x[0], name="y_expr", model_expects_batch_dim=False 202 | ) 203 | 204 | ocp.model.cost_y_expr = l4c_y_expr(x) 205 | ocp.model.cost_y_expr_e = x[0] 206 | 207 | ocp.cost.W_e = np.array([[0.0]]) 208 | ocp.cost.yref_e = np.array([0.0]) 209 | 210 | # Initial reference trajectory (will be overwritten) 211 | ocp.cost.yref = np.zeros(1) 212 | 213 | # Initial state (will be overwritten) 214 | ocp.constraints.x0 = model.x_start 215 | 216 | # Set constraints 217 | a_max = 10 218 | ocp.constraints.lbu = np.array([-a_max]) 219 | ocp.constraints.ubu = np.array([a_max]) 220 | ocp.constraints.idxbu = np.array([0]) 221 | 222 | # Solver options 223 | ocp.solver_options.qp_solver = "FULL_CONDENSING_HPIPM" 224 | ocp.solver_options.hessian_approx = "GAUSS_NEWTON" 225 | ocp.solver_options.integrator_type = "ERK" 226 | ocp.solver_options.nlp_solver_type = "SQP_RTI" 227 | ocp.solver_options.sim_method_num_stages = 1 # Euler 228 | ocp.solver_options.sim_method_num_steps = 1 229 | ocp.solver_options.nlp_solver_max_iter = 1 230 | # ocp.solver_options.num_threads_in_batch_solve = self.num_threads_acados_openmp 231 | 232 | if self.external_shared_lib_dir is not None: 233 | ocp.solver_options.model_external_shared_lib_dir = ( 234 | self.external_shared_lib_dir 235 | ) 236 | 237 | if COST == "LINEAR_LS": 238 | ocp.solver_options.model_external_shared_lib_name = ( 239 | self.external_shared_lib_name 240 | ) 241 | else: 242 | ocp.solver_options.model_external_shared_lib_name = ( 243 | self.external_shared_lib_name + " -l" + l4c_y_expr.name 244 | ) 245 | 246 | return ocp 247 | 248 | def acados_model(self, model): 249 | model_ac = AcadosModel() 250 | model_ac.f_impl_expr = model.xdot - model.f_expl 251 | model_ac.f_expl_expr = model.f_expl 252 | model_ac.x = model.x 253 | model_ac.xdot = model.xdot 254 | model_ac.u = model.u 255 | model_ac.name = model.name 256 | return model_ac 257 | 258 | 259 | def run(): 260 | N = 10 261 | 262 | nn_fun_new = MultiLayerPerceptron( 263 | n_inputs=1, hidden_layers=2, hidden_size=64, n_outputs=1 264 | ) 265 | nn_fun_new.load_state_dict(torch.load("nn_params.pt")) 266 | nn_fun_new.eval() 267 | 268 | learned_dyn_model = l4c.L4CasADi(nn_fun_new, name="learned_dyn") 269 | 270 | model = DoubleIntegratorWithLearnedDynamics(learned_dyn_model) 271 | 272 | print(learned_dyn_model.shared_lib_dir) 273 | print(learned_dyn_model.name) 274 | 275 | solver = MPC( 276 | model=model.model(), 277 | N=N, 278 | external_shared_lib_dir=learned_dyn_model.shared_lib_dir, 279 | external_shared_lib_name=learned_dyn_model.name, 280 | ).solver 281 | 282 | x = [] 283 | x_ref = [] 284 | ts = 1.0 / N 285 | xt = np.array([1.0, 0.0]) 286 | opt_times = [] 287 | 288 | for i in range(50): 289 | now = time.time() 290 | t = np.linspace(i * ts, i * ts + 1.0, 10) 291 | yref = np.sin(0.5 * t + np.pi / 2) 292 | x_ref.append(yref[0]) 293 | for t, ref in enumerate(yref): 294 | solver.set(t, "yref", ref) 295 | solver.set(0, "lbx", xt) 296 | solver.set(0, "ubx", xt) 297 | solver.solve() 298 | xt = solver.get(1, "x") 299 | x.append(xt) 300 | 301 | x_l = [] 302 | for i in range(N): 303 | x_l.append(solver.get(i, "x")) 304 | 305 | elapsed = time.time() - now 306 | opt_times.append(elapsed) 307 | 308 | print( 309 | f"Mean iteration time: {1000*np.mean(opt_times):.1f}ms -- {1/np.mean(opt_times):.0f}Hz)" 310 | ) 311 | 312 | 313 | if __name__ == "__main__": 314 | run() 315 | -------------------------------------------------------------------------------- /test/l4casadi_comparison/nn_params.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/IntelligentControlSystems/l4acados/a56ec0184c52d832f6801ffe7e39c463d6fa1d37/test/l4casadi_comparison/nn_params.pt -------------------------------------------------------------------------------- /test/l4casadi_comparison/test_l4casadi_l4acados_comparison.py: -------------------------------------------------------------------------------- 1 | import l4casadi as l4c 2 | import numpy as np 3 | import time 4 | import l4acados as l4a 5 | from typing import Optional, Union 6 | import torch 7 | import casadi as cs 8 | from l4acados.controllers import ResidualLearningMPC 9 | from l4acados.models import PyTorchFeatureSelector, PyTorchResidualModel 10 | from l4acados.controllers.zoro_acados_utils import setup_sim_from_ocp 11 | import argparse 12 | import os, shutil, re 13 | import subprocess 14 | import matplotlib.pyplot as plt 15 | 16 | from l4casadi_comparison.l4casadi_with_acados import ( 17 | run, 18 | MultiLayerPerceptron, 19 | DoubleIntegratorWithLearnedDynamics, 20 | MPC, 21 | train_nn, 22 | ) 23 | 24 | 25 | def init_l4casadi( 26 | nn_model: torch.nn.Module, 27 | N: int, 28 | device="cpu", 29 | ): 30 | learned_dyn_model = l4c.L4CasADi( 31 | nn_model, 32 | name="learned_dyn", 33 | device=device, 34 | ) 35 | 36 | model = DoubleIntegratorWithLearnedDynamics(learned_dyn_model) 37 | 38 | mpc_obj = MPC( 39 | model=model.model(), 40 | N=N, 41 | external_shared_lib_dir=learned_dyn_model.shared_lib_dir, 42 | external_shared_lib_name=learned_dyn_model.name, 43 | ) 44 | solver = mpc_obj.solver 45 | 46 | ocp = mpc_obj.ocp() 47 | return solver 48 | 49 | 50 | def init_l4acados( 51 | nn_model: torch.nn.Module, 52 | N: int, 53 | device="cpu", 54 | use_cython=False, 55 | ): 56 | feature_selector = PyTorchFeatureSelector([1, 0, 0], device=device) 57 | residual_model = PyTorchResidualModel( 58 | nn_model, 59 | feature_selector, 60 | ) 61 | B_proj = (1.0 / N) * np.array([[1.0], [1.0]]) # NOTE: Ts = 1.0 hard coded 62 | model_new = DoubleIntegratorWithLearnedDynamics(None, name="wr_new") 63 | mpc_obj_nolib = MPC( 64 | model=model_new.model(), 65 | N=N, 66 | ) 67 | solver_nolib = mpc_obj_nolib.solver 68 | ocp_nolib = mpc_obj_nolib.ocp() 69 | sim_nolib = setup_sim_from_ocp(ocp_nolib) 70 | 71 | solver_l4acados = ResidualLearningMPC( 72 | ocp=ocp_nolib, 73 | B=B_proj, 74 | residual_model=residual_model, 75 | use_cython=use_cython, 76 | ) 77 | 78 | return solver_l4acados 79 | 80 | 81 | def run_timing_experiment(N, solver, solve_call, solve_steps=1e3): 82 | x = [] 83 | u = [] 84 | xt = np.array([1.0, 0.0]) 85 | T = 1.0 86 | ts = T / N 87 | opt_times = [] 88 | 89 | for i in range(solve_steps): 90 | # t = np.linspace(i * ts, (i + 1) * ts, N) # TODO: this ts should be T IMO, maybe with lower frequency? 91 | t = np.linspace( 92 | i * T, (i + 1) * T, N 93 | ) # TODO: this ts should be T IMO, maybe with lower frequency? 94 | yref = np.sin(0.1 * t + np.pi / 2) 95 | for t, ref in enumerate(yref): 96 | solver.set(t, "yref", np.array([ref])) 97 | solver.set(0, "lbx", xt) 98 | solver.set(0, "ubx", xt) 99 | 100 | # now = time.time() 101 | elapsed = solve_call() 102 | 103 | xt = solver.get(1, "x") 104 | 105 | opt_times.append(elapsed) 106 | x.append(xt) 107 | 108 | print(f"Running timing experiment: {i}/{solve_steps}") 109 | 110 | return x, opt_times 111 | 112 | 113 | def time_fun_call(fun): 114 | now = time.perf_counter() 115 | fun() 116 | return time.perf_counter() - now 117 | 118 | 119 | def delete_file_by_pattern(dir_path, pattern): 120 | for f in os.listdir(dir_path): 121 | if re.search(pattern, f): 122 | os.remove(os.path.join(dir_path, f)) 123 | 124 | 125 | def run( 126 | N, 127 | solve_steps, 128 | device="cpu", 129 | save_data=False, 130 | save_dir="data", 131 | nn_params_path=None, 132 | ): 133 | 134 | nn_fun = MultiLayerPerceptron( 135 | n_inputs=1, hidden_layers=2, hidden_size=64, n_outputs=1 136 | ) 137 | if nn_params_path is not None: 138 | nn_fun.load_state_dict(torch.load(nn_params_path)) 139 | nn_fun.eval() 140 | 141 | solver_l4casadi = init_l4casadi( 142 | nn_fun, 143 | N, 144 | device=device, 145 | ) 146 | x_l4casadi, opt_times_l4casadi = run_timing_experiment( 147 | N, 148 | solver_l4casadi, 149 | lambda: time_fun_call(solver_l4casadi.solve), 150 | solve_steps=solve_steps, 151 | ) 152 | 153 | shutil.rmtree("c_generated_code") 154 | shutil.rmtree("_l4c_generated") 155 | delete_file_by_pattern("./", r".*[ocp|sim].*\.json") 156 | 157 | solver_l4acados = init_l4acados( 158 | nn_fun, 159 | N, 160 | device=device, 161 | use_cython=False, 162 | ) 163 | x_l4acados, opt_times_l4acados = run_timing_experiment( 164 | N, 165 | solver_l4acados.ocp_solver, 166 | lambda: time_fun_call(lambda: solver_l4acados.solve()), 167 | solve_steps=solve_steps, 168 | ) 169 | 170 | shutil.rmtree("c_generated_code") 171 | delete_file_by_pattern("./", r".*[ocp|sim].*\.json") 172 | 173 | del solver_l4casadi, solver_l4acados 174 | 175 | return x_l4casadi, opt_times_l4casadi, x_l4acados, opt_times_l4acados 176 | 177 | 178 | def test_l4acados_equal_l4casadi(): 179 | # check if nn_params.pt is in the root directory 180 | test_l4casadi_dir = os.path.dirname(os.path.abspath(__file__)) 181 | nn_params_path = os.path.join("nn_params.pt") 182 | 183 | if not os.path.exists(nn_params_path): 184 | train_nn(nn_params_path=nn_params_path) 185 | 186 | x_l4casadi, opt_times_l4casadi, x_l4acados, opt_times_l4acados = run( 187 | N=10, 188 | solve_steps=100, 189 | device="cpu", 190 | save_data=False, 191 | nn_params_path=nn_params_path, 192 | ) 193 | 194 | x_l4casadi = np.array(x_l4casadi) 195 | x_l4acados = np.array(x_l4acados) 196 | 197 | ( 198 | x_l4casadi_zero, 199 | opt_times_l4casadi_zero, 200 | x_l4acados_zero, 201 | opt_times_l4acados_zero, 202 | ) = run( 203 | N=10, 204 | solve_steps=100, 205 | device="cpu", 206 | save_data=False, 207 | nn_params_path=None, 208 | ) 209 | 210 | x_l4casadi_zero = np.array(x_l4casadi_zero) 211 | x_l4acados_zero = np.array(x_l4acados_zero) 212 | 213 | print(f"Normdiff casadi: {np.linalg.norm(x_l4casadi - x_l4casadi_zero)}") 214 | print(f"Normdiff acados: {np.linalg.norm(x_l4acados - x_l4acados_zero)}") 215 | print(f"Normdiff casadi vs acados: {np.linalg.norm(x_l4casadi - x_l4acados)}") 216 | print( 217 | f"Normdiff casadi vs acados zero: {np.linalg.norm(x_l4casadi_zero - x_l4acados_zero)}" 218 | ) 219 | 220 | assert not np.allclose(x_l4casadi, x_l4casadi_zero) 221 | assert not np.allclose(x_l4acados, x_l4acados_zero) 222 | 223 | assert np.allclose(x_l4casadi_zero, x_l4acados_zero, atol=1e-5, rtol=1e-4) 224 | assert np.allclose(x_l4casadi, x_l4acados, atol=1e-5, rtol=1e-4) 225 | 226 | 227 | if __name__ == "__main__": 228 | test_l4acados_equal_l4casadi() 229 | --------------------------------------------------------------------------------