Repository: TUM-AAS/neural-mpc Branch: main Commit: 5f15661c2d1a Files: 51 Total size: 460.3 KB Directory structure: gitextract__1r96zfd/ ├── .gitignore ├── .gitmodules ├── LICENSE ├── README.md ├── requirements.txt └── ros_dd_mpc/ ├── CMakeLists.txt ├── config/ │ ├── agisim_simulation_run.yaml │ ├── arena_limits.yaml │ ├── arena_run.yaml │ ├── basic.world │ ├── circle_and_lemniscate_options.yaml │ ├── configuration_parameters.py │ ├── flyingroom_limits.yaml │ ├── ground_effect_limits.yaml │ ├── kingfisher.yaml │ └── simulation_run.yaml ├── launch/ │ └── dd_mpc_wrapper.launch ├── msg/ │ └── ReferenceTrajectory.msg ├── nodes/ │ ├── dd_mpc_node.py │ └── reference_publisher_node.py ├── package.xml └── src/ ├── __init__.py ├── experiments/ │ ├── __init__.py │ ├── comparative_experiment.py │ ├── point_tracking_and_record.py │ └── trajectory_test.py ├── model_fitting/ │ ├── __init__.py │ ├── gp.py │ ├── gp_common.py │ ├── gp_fitting.py │ ├── gp_visualization.py │ ├── mlp_common.py │ ├── mlp_fitting.py │ ├── mlp_quad_res_fitting.py │ ├── process_neurobem_dataset.py │ ├── rdrv_fitting.py │ └── system_identification.py ├── quad_mpc/ │ ├── __init__.py │ ├── create_ros_dd_mpc.py │ ├── quad_3d.py │ ├── quad_3d_mpc.py │ └── quad_3d_optimizer.py └── utils/ ├── __init__.py ├── animator.py ├── ground_map.py ├── keyframe_3d_gen.py ├── quad_3d_opt_utils.py ├── trajectories.py ├── trajectory_generator.py ├── utils.py └── visualization.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: .gitignore ================================================ venv/* data/ results/ .idea/ *.pyc ros_gp_mpc/acados_models/* c_generated_code/ *.egg-info/ dist/ build/ acados_models/ ================================================ FILE: .gitmodules ================================================ [submodule "ml-casadi"] path = ml-casadi url = https://github.com/TUM-AAS/ml-casadi ================================================ FILE: LICENSE ================================================ GNU GENERAL PUBLIC LICENSE Version 3, 29 June 2007 Copyright (C) 2007 Free Software Foundation, Inc. Everyone is permitted to copy and distribute verbatim copies of this license document, but changing it is not allowed. Preamble The GNU General Public License is a free, copyleft license for software and other kinds of works. The licenses for most software and other practical works are designed to take away your freedom to share and change the works. By contrast, the GNU General Public License is intended to guarantee your freedom to share and change all versions of a program--to make sure it remains free software for all its users. We, the Free Software Foundation, use the GNU General Public License for most of our software; it applies also to any other work released this way by its authors. You can apply it to your programs, too. When we speak of free software, we are referring to freedom, not price. Our General Public Licenses are designed to make sure that you have the freedom to distribute copies of free software (and charge for them if you wish), that you receive source code or can get it if you want it, that you can change the software or use pieces of it in new free programs, and that you know you can do these things. To protect your rights, we need to prevent others from denying you these rights or asking you to surrender the rights. Therefore, you have certain responsibilities if you distribute copies of the software, or if you modify it: responsibilities to respect the freedom of others. For example, if you distribute copies of such a program, whether gratis or for a fee, you must pass on to the recipients the same freedoms that you received. You must make sure that they, too, receive or can get the source code. And you must show them these terms so they know their rights. Developers that use the GNU GPL protect your rights with two steps: (1) assert copyright on the software, and (2) offer you this License giving you legal permission to copy, distribute and/or modify it. For the developers' and authors' protection, the GPL clearly explains that there is no warranty for this free software. For both users' and authors' sake, the GPL requires that modified versions be marked as changed, so that their problems will not be attributed erroneously to authors of previous versions. Some devices are designed to deny users access to install or run modified versions of the software inside them, although the manufacturer can do so. This is fundamentally incompatible with the aim of protecting users' freedom to change the software. The systematic pattern of such abuse occurs in the area of products for individuals to use, which is precisely where it is most unacceptable. Therefore, we have designed this version of the GPL to prohibit the practice for those products. If such problems arise substantially in other domains, we stand ready to extend this provision to those domains in future versions of the GPL, as needed to protect the freedom of users. Finally, every program is threatened constantly by software patents. States should not allow patents to restrict development and use of software on general-purpose computers, but in those that do, we wish to avoid the special danger that patents applied to a free program could make it effectively proprietary. To prevent this, the GPL assures that patents cannot be used to render the program non-free. The precise terms and conditions for copying, distribution and modification follow. TERMS AND CONDITIONS 0. Definitions. "This License" refers to version 3 of the GNU General Public License. "Copyright" also means copyright-like laws that apply to other kinds of works, such as semiconductor masks. "The Program" refers to any copyrightable work licensed under this License. Each licensee is addressed as "you". "Licensees" and "recipients" may be individuals or organizations. To "modify" a work means to copy from or adapt all or part of the work in a fashion requiring copyright permission, other than the making of an exact copy. The resulting work is called a "modified version" of the earlier work or a work "based on" the earlier work. A "covered work" means either the unmodified Program or a work based on the Program. To "propagate" a work means to do anything with it that, without permission, would make you directly or secondarily liable for infringement under applicable copyright law, except executing it on a computer or modifying a private copy. Propagation includes copying, distribution (with or without modification), making available to the public, and in some countries other activities as well. To "convey" a work means any kind of propagation that enables other parties to make or receive copies. Mere interaction with a user through a computer network, with no transfer of a copy, is not conveying. An interactive user interface displays "Appropriate Legal Notices" to the extent that it includes a convenient and prominently visible feature that (1) displays an appropriate copyright notice, and (2) tells the user that there is no warranty for the work (except to the extent that warranties are provided), that licensees may convey the work under this License, and how to view a copy of this License. If the interface presents a list of user commands or options, such as a menu, a prominent item in the list meets this criterion. 1. Source Code. The "source code" for a work means the preferred form of the work for making modifications to it. "Object code" means any non-source form of a work. A "Standard Interface" means an interface that either is an official standard defined by a recognized standards body, or, in the case of interfaces specified for a particular programming language, one that is widely used among developers working in that language. The "System Libraries" of an executable work include anything, other than the work as a whole, that (a) is included in the normal form of packaging a Major Component, but which is not part of that Major Component, and (b) serves only to enable use of the work with that Major Component, or to implement a Standard Interface for which an implementation is available to the public in source code form. A "Major Component", in this context, means a major essential component (kernel, window system, and so on) of the specific operating system (if any) on which the executable work runs, or a compiler used to produce the work, or an object code interpreter used to run it. The "Corresponding Source" for a work in object code form means all the source code needed to generate, install, and (for an executable work) run the object code and to modify the work, including scripts to control those activities. However, it does not include the work's System Libraries, or general-purpose tools or generally available free programs which are used unmodified in performing those activities but which are not part of the work. For example, Corresponding Source includes interface definition files associated with source files for the work, and the source code for shared libraries and dynamically linked subprograms that the work is specifically designed to require, such as by intimate data communication or control flow between those subprograms and other parts of the work. The Corresponding Source need not include anything that users can regenerate automatically from other parts of the Corresponding Source. The Corresponding Source for a work in source code form is that same work. 2. Basic Permissions. All rights granted under this License are granted for the term of copyright on the Program, and are irrevocable provided the stated conditions are met. This License explicitly affirms your unlimited permission to run the unmodified Program. The output from running a covered work is covered by this License only if the output, given its content, constitutes a covered work. This License acknowledges your rights of fair use or other equivalent, as provided by copyright law. You may make, run and propagate covered works that you do not convey, without conditions so long as your license otherwise remains in force. You may convey covered works to others for the sole purpose of having them make modifications exclusively for you, or provide you with facilities for running those works, provided that you comply with the terms of this License in conveying all material for which you do not control copyright. Those thus making or running the covered works for you must do so exclusively on your behalf, under your direction and control, on terms that prohibit them from making any copies of your copyrighted material outside their relationship with you. Conveying under any other circumstances is permitted solely under the conditions stated below. Sublicensing is not allowed; section 10 makes it unnecessary. 3. Protecting Users' Legal Rights From Anti-Circumvention Law. No covered work shall be deemed part of an effective technological measure under any applicable law fulfilling obligations under article 11 of the WIPO copyright treaty adopted on 20 December 1996, or similar laws prohibiting or restricting circumvention of such measures. When you convey a covered work, you waive any legal power to forbid circumvention of technological measures to the extent such circumvention is effected by exercising rights under this License with respect to the covered work, and you disclaim any intention to limit operation or modification of the work as a means of enforcing, against the work's users, your or third parties' legal rights to forbid circumvention of technological measures. 4. Conveying Verbatim Copies. You may convey verbatim copies of the Program's source code as you receive it, in any medium, provided that you conspicuously and appropriately publish on each copy an appropriate copyright notice; keep intact all notices stating that this License and any non-permissive terms added in accord with section 7 apply to the code; keep intact all notices of the absence of any warranty; and give all recipients a copy of this License along with the Program. You may charge any price or no price for each copy that you convey, and you may offer support or warranty protection for a fee. 5. Conveying Modified Source Versions. You may convey a work based on the Program, or the modifications to produce it from the Program, in the form of source code under the terms of section 4, provided that you also meet all of these conditions: a) The work must carry prominent notices stating that you modified it, and giving a relevant date. b) The work must carry prominent notices stating that it is released under this License and any conditions added under section 7. This requirement modifies the requirement in section 4 to "keep intact all notices". c) You must license the entire work, as a whole, under this License to anyone who comes into possession of a copy. This License will therefore apply, along with any applicable section 7 additional terms, to the whole of the work, and all its parts, regardless of how they are packaged. This License gives no permission to license the work in any other way, but it does not invalidate such permission if you have separately received it. d) If the work has interactive user interfaces, each must display Appropriate Legal Notices; however, if the Program has interactive interfaces that do not display Appropriate Legal Notices, your work need not make them do so. A compilation of a covered work with other separate and independent works, which are not by their nature extensions of the covered work, and which are not combined with it such as to form a larger program, in or on a volume of a storage or distribution medium, is called an "aggregate" if the compilation and its resulting copyright are not used to limit the access or legal rights of the compilation's users beyond what the individual works permit. Inclusion of a covered work in an aggregate does not cause this License to apply to the other parts of the aggregate. 6. Conveying Non-Source Forms. You may convey a covered work in object code form under the terms of sections 4 and 5, provided that you also convey the machine-readable Corresponding Source under the terms of this License, in one of these ways: a) Convey the object code in, or embodied in, a physical product (including a physical distribution medium), accompanied by the Corresponding Source fixed on a durable physical medium customarily used for software interchange. b) Convey the object code in, or embodied in, a physical product (including a physical distribution medium), accompanied by a written offer, valid for at least three years and valid for as long as you offer spare parts or customer support for that product model, to give anyone who possesses the object code either (1) a copy of the Corresponding Source for all the software in the product that is covered by this License, on a durable physical medium customarily used for software interchange, for a price no more than your reasonable cost of physically performing this conveying of source, or (2) access to copy the Corresponding Source from a network server at no charge. c) Convey individual copies of the object code with a copy of the written offer to provide the Corresponding Source. This alternative is allowed only occasionally and noncommercially, and only if you received the object code with such an offer, in accord with subsection 6b. d) Convey the object code by offering access from a designated place (gratis or for a charge), and offer equivalent access to the Corresponding Source in the same way through the same place at no further charge. You need not require recipients to copy the Corresponding Source along with the object code. If the place to copy the object code is a network server, the Corresponding Source may be on a different server (operated by you or a third party) that supports equivalent copying facilities, provided you maintain clear directions next to the object code saying where to find the Corresponding Source. Regardless of what server hosts the Corresponding Source, you remain obligated to ensure that it is available for as long as needed to satisfy these requirements. e) Convey the object code using peer-to-peer transmission, provided you inform other peers where the object code and Corresponding Source of the work are being offered to the general public at no charge under subsection 6d. A separable portion of the object code, whose source code is excluded from the Corresponding Source as a System Library, need not be included in conveying the object code work. A "User Product" is either (1) a "consumer product", which means any tangible personal property which is normally used for personal, family, or household purposes, or (2) anything designed or sold for incorporation into a dwelling. In determining whether a product is a consumer product, doubtful cases shall be resolved in favor of coverage. For a particular product received by a particular user, "normally used" refers to a typical or common use of that class of product, regardless of the status of the particular user or of the way in which the particular user actually uses, or expects or is expected to use, the product. A product is a consumer product regardless of whether the product has substantial commercial, industrial or non-consumer uses, unless such uses represent the only significant mode of use of the product. "Installation Information" for a User Product means any methods, procedures, authorization keys, or other information required to install and execute modified versions of a covered work in that User Product from a modified version of its Corresponding Source. The information must suffice to ensure that the continued functioning of the modified object code is in no case prevented or interfered with solely because modification has been made. If you convey an object code work under this section in, or with, or specifically for use in, a User Product, and the conveying occurs as part of a transaction in which the right of possession and use of the User Product is transferred to the recipient in perpetuity or for a fixed term (regardless of how the transaction is characterized), the Corresponding Source conveyed under this section must be accompanied by the Installation Information. But this requirement does not apply if neither you nor any third party retains the ability to install modified object code on the User Product (for example, the work has been installed in ROM). The requirement to provide Installation Information does not include a requirement to continue to provide support service, warranty, or updates for a work that has been modified or installed by the recipient, or for the User Product in which it has been modified or installed. Access to a network may be denied when the modification itself materially and adversely affects the operation of the network or violates the rules and protocols for communication across the network. Corresponding Source conveyed, and Installation Information provided, in accord with this section must be in a format that is publicly documented (and with an implementation available to the public in source code form), and must require no special password or key for unpacking, reading or copying. 7. Additional Terms. "Additional permissions" are terms that supplement the terms of this License by making exceptions from one or more of its conditions. Additional permissions that are applicable to the entire Program shall be treated as though they were included in this License, to the extent that they are valid under applicable law. If additional permissions apply only to part of the Program, that part may be used separately under those permissions, but the entire Program remains governed by this License without regard to the additional permissions. When you convey a copy of a covered work, you may at your option remove any additional permissions from that copy, or from any part of it. (Additional permissions may be written to require their own removal in certain cases when you modify the work.) You may place additional permissions on material, added by you to a covered work, for which you have or can give appropriate copyright permission. Notwithstanding any other provision of this License, for material you add to a covered work, you may (if authorized by the copyright holders of that material) supplement the terms of this License with terms: a) Disclaiming warranty or limiting liability differently from the terms of sections 15 and 16 of this License; or b) Requiring preservation of specified reasonable legal notices or author attributions in that material or in the Appropriate Legal Notices displayed by works containing it; or c) Prohibiting misrepresentation of the origin of that material, or requiring that modified versions of such material be marked in reasonable ways as different from the original version; or d) Limiting the use for publicity purposes of names of licensors or authors of the material; or e) Declining to grant rights under trademark law for use of some trade names, trademarks, or service marks; or f) Requiring indemnification of licensors and authors of that material by anyone who conveys the material (or modified versions of it) with contractual assumptions of liability to the recipient, for any liability that these contractual assumptions directly impose on those licensors and authors. All other non-permissive additional terms are considered "further restrictions" within the meaning of section 10. If the Program as you received it, or any part of it, contains a notice stating that it is governed by this License along with a term that is a further restriction, you may remove that term. If a license document contains a further restriction but permits relicensing or conveying under this License, you may add to a covered work material governed by the terms of that license document, provided that the further restriction does not survive such relicensing or conveying. If you add terms to a covered work in accord with this section, you must place, in the relevant source files, a statement of the additional terms that apply to those files, or a notice indicating where to find the applicable terms. Additional terms, permissive or non-permissive, may be stated in the form of a separately written license, or stated as exceptions; the above requirements apply either way. 8. Termination. You may not propagate or modify a covered work except as expressly provided under this License. Any attempt otherwise to propagate or modify it is void, and will automatically terminate your rights under this License (including any patent licenses granted under the third paragraph of section 11). However, if you cease all violation of this License, then your license from a particular copyright holder is reinstated (a) provisionally, unless and until the copyright holder explicitly and finally terminates your license, and (b) permanently, if the copyright holder fails to notify you of the violation by some reasonable means prior to 60 days after the cessation. Moreover, your license from a particular copyright holder is reinstated permanently if the copyright holder notifies you of the violation by some reasonable means, this is the first time you have received notice of violation of this License (for any work) from that copyright holder, and you cure the violation prior to 30 days after your receipt of the notice. Termination of your rights under this section does not terminate the licenses of parties who have received copies or rights from you under this License. If your rights have been terminated and not permanently reinstated, you do not qualify to receive new licenses for the same material under section 10. 9. Acceptance Not Required for Having Copies. You are not required to accept this License in order to receive or run a copy of the Program. Ancillary propagation of a covered work occurring solely as a consequence of using peer-to-peer transmission to receive a copy likewise does not require acceptance. However, nothing other than this License grants you permission to propagate or modify any covered work. These actions infringe copyright if you do not accept this License. Therefore, by modifying or propagating a covered work, you indicate your acceptance of this License to do so. 10. Automatic Licensing of Downstream Recipients. Each time you convey a covered work, the recipient automatically receives a license from the original licensors, to run, modify and propagate that work, subject to this License. You are not responsible for enforcing compliance by third parties with this License. An "entity transaction" is a transaction transferring control of an organization, or substantially all assets of one, or subdividing an organization, or merging organizations. If propagation of a covered work results from an entity transaction, each party to that transaction who receives a copy of the work also receives whatever licenses to the work the party's predecessor in interest had or could give under the previous paragraph, plus a right to possession of the Corresponding Source of the work from the predecessor in interest, if the predecessor has it or can get it with reasonable efforts. You may not impose any further restrictions on the exercise of the rights granted or affirmed under this License. For example, you may not impose a license fee, royalty, or other charge for exercise of rights granted under this License, and you may not initiate litigation (including a cross-claim or counterclaim in a lawsuit) alleging that any patent claim is infringed by making, using, selling, offering for sale, or importing the Program or any portion of it. 11. Patents. A "contributor" is a copyright holder who authorizes use under this License of the Program or a work on which the Program is based. The work thus licensed is called the contributor's "contributor version". A contributor's "essential patent claims" are all patent claims owned or controlled by the contributor, whether already acquired or hereafter acquired, that would be infringed by some manner, permitted by this License, of making, using, or selling its contributor version, but do not include claims that would be infringed only as a consequence of further modification of the contributor version. For purposes of this definition, "control" includes the right to grant patent sublicenses in a manner consistent with the requirements of this License. Each contributor grants you a non-exclusive, worldwide, royalty-free patent license under the contributor's essential patent claims, to make, use, sell, offer for sale, import and otherwise run, modify and propagate the contents of its contributor version. In the following three paragraphs, a "patent license" is any express agreement or commitment, however denominated, not to enforce a patent (such as an express permission to practice a patent or covenant not to sue for patent infringement). To "grant" such a patent license to a party means to make such an agreement or commitment not to enforce a patent against the party. If you convey a covered work, knowingly relying on a patent license, and the Corresponding Source of the work is not available for anyone to copy, free of charge and under the terms of this License, through a publicly available network server or other readily accessible means, then you must either (1) cause the Corresponding Source to be so available, or (2) arrange to deprive yourself of the benefit of the patent license for this particular work, or (3) arrange, in a manner consistent with the requirements of this License, to extend the patent license to downstream recipients. "Knowingly relying" means you have actual knowledge that, but for the patent license, your conveying the covered work in a country, or your recipient's use of the covered work in a country, would infringe one or more identifiable patents in that country that you have reason to believe are valid. If, pursuant to or in connection with a single transaction or arrangement, you convey, or propagate by procuring conveyance of, a covered work, and grant a patent license to some of the parties receiving the covered work authorizing them to use, propagate, modify or convey a specific copy of the covered work, then the patent license you grant is automatically extended to all recipients of the covered work and works based on it. A patent license is "discriminatory" if it does not include within the scope of its coverage, prohibits the exercise of, or is conditioned on the non-exercise of one or more of the rights that are specifically granted under this License. You may not convey a covered work if you are a party to an arrangement with a third party that is in the business of distributing software, under which you make payment to the third party based on the extent of your activity of conveying the work, and under which the third party grants, to any of the parties who would receive the covered work from you, a discriminatory patent license (a) in connection with copies of the covered work conveyed by you (or copies made from those copies), or (b) primarily for and in connection with specific products or compilations that contain the covered work, unless you entered into that arrangement, or that patent license was granted, prior to 28 March 2007. Nothing in this License shall be construed as excluding or limiting any implied license or other defenses to infringement that may otherwise be available to you under applicable patent law. 12. No Surrender of Others' Freedom. If conditions are imposed on you (whether by court order, agreement or otherwise) that contradict the conditions of this License, they do not excuse you from the conditions of this License. If you cannot convey a covered work so as to satisfy simultaneously your obligations under this License and any other pertinent obligations, then as a consequence you may not convey it at all. For example, if you agree to terms that obligate you to collect a royalty for further conveying from those to whom you convey the Program, the only way you could satisfy both those terms and this License would be to refrain entirely from conveying the Program. 13. Use with the GNU Affero General Public License. Notwithstanding any other provision of this License, you have permission to link or combine any covered work with a work licensed under version 3 of the GNU Affero General Public License into a single combined work, and to convey the resulting work. The terms of this License will continue to apply to the part which is the covered work, but the special requirements of the GNU Affero General Public License, section 13, concerning interaction through a network will apply to the combination as such. 14. Revised Versions of this License. The Free Software Foundation may publish revised and/or new versions of the GNU General Public License from time to time. Such new versions will be similar in spirit to the present version, but may differ in detail to address new problems or concerns. Each version is given a distinguishing version number. If the Program specifies that a certain numbered version of the GNU General Public License "or any later version" applies to it, you have the option of following the terms and conditions either of that numbered version or of any later version published by the Free Software Foundation. If the Program does not specify a version number of the GNU General Public License, you may choose any version ever published by the Free Software Foundation. If the Program specifies that a proxy can decide which future versions of the GNU General Public License can be used, that proxy's public statement of acceptance of a version permanently authorizes you to choose that version for the Program. Later license versions may give you additional or different permissions. However, no additional obligations are imposed on any author or copyright holder as a result of your choosing to follow a later version. 15. Disclaimer of Warranty. THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING, REPAIR OR CORRECTION. 16. Limitation of Liability. IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. 17. Interpretation of Sections 15 and 16. If the disclaimer of warranty and limitation of liability provided above cannot be given local legal effect according to their terms, reviewing courts shall apply local law that most closely approximates an absolute waiver of all civil liability in connection with the Program, unless a warranty or assumption of liability accompanies a copy of the Program in return for a fee. END OF TERMS AND CONDITIONS How to Apply These Terms to Your New Programs If you develop a new program, and you want it to be of the greatest possible use to the public, the best way to achieve this is to make it free software which everyone can redistribute and change under these terms. To do so, attach the following notices to the program. It is safest to attach them to the start of each source file to most effectively state the exclusion of warranty; and each file should have at least the "copyright" line and a pointer to where the full notice is found. Copyright (C) This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . Also add information on how to contact you by electronic and paper mail. If the program does terminal interaction, make it output a short notice like this when it starts in an interactive mode: Copyright (C) This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'. This is free software, and you are welcome to redistribute it under certain conditions; type `show c' for details. The hypothetical commands `show w' and `show c' should show the appropriate parts of the General Public License. Of course, your program's commands might be different; for a GUI interface, you would use an "about box". You should also get your employer (if you work as a programmer) or school, if any, to sign a "copyright disclaimer" for the program, if necessary. For more information on this, and how to apply and follow the GNU GPL, see . The GNU General Public License does not permit incorporating your program into proprietary programs. If your program is a subroutine library, you may consider it more useful to permit linking proprietary applications with the library. If this is what you want to do, use the GNU Lesser General Public License instead of this License. But first, please read . ================================================ FILE: README.md ================================================ # Real-Time Neural MPC This repository contains the code for experiments associated to our paper ``` Real-time Neural-MPC: Deep Learning Model Predictive Control for Quadrotors and Agile Robotic Platforms ``` [Arxiv Link](https://arxiv.org/pdf/2203.07747) If you are looking for the ML-CasADi framework code you can find it [here](https://github.com/TUM-AAS/ml-casadi). ## Installation ### Checkout Submodules ``` git submodule update --init --recursive ``` ### Acados - Follow the [installation instructions](https://docs.acados.org/installation/index.html). - Install the [Python Interface](https://docs.acados.org/python_interface/index.html). - Ensure that `LD_LIBRARY_PATH` is set correctly (`DYLD_LIBRARY_PATH`on MacOS). - Ensure that `ACADOS_SOURCE_DIR` is set correctly. ### Further Requirements ``` pip install -r requirements.txt ``` Make sure the ML-CasADi framework is part of the python path. ``` export PYTHONPATH="${PYTHONPATH}:/ml-casadi" ``` Python 3.9 is recommended. # Experiments The provided code is based on the work of [Torrente et al.](https://github.com/uzh-rpg/data_driven_mpc) All functionality of the original code base is retained. Change the working directory to ``` cd ros_dd_mpc ``` ## Simulation ### Data Collection Run the following script to collect a few minutes of flight samples ``` python src/experiments/point_tracking_and_record.py --recording --dataset_name simplified_sim_dataset --simulation_time 300 ``` ### Fitting a MLP Model Edit the following variables of configuration file in `config/configuration_parameters.py` (class `ModelFitConfig`) so that the training script is referenced to the desired dataset. For redundancy, in order to identify the correct data file, we require to specify both the name of the dataset as well as the parameters used while acquiring the data. In other words, you must input the simulator options used while running the previous python script. If you did not modify these variables earlier, you don't need to change anything this time as the default setting will work: ``` # ## Dataset loading ## # ds_name = "simplified_sim_dataset" ds_metadata = { "noisy": True, "drag": True, "payload": False, "motor_noise": True } ``` The following command will train an MLP model with 4 hidden layers 64 neurons each to model the residual error on the velocities in x, y, and z direction (7, 8, 9 in the state). We assign a name to the model for future referencing, e.g.: `simple_sim_mlp` ``` python src/model_fitting/mlp_fitting.py --model_name simple_sim_mlp --hidden_size 64 --hidden_layers 4 --x 7 8 9 --y 7 8 9 --epochs 100 ``` The model will be saved under the directory `ros_dd_mpc/results/model_fitting//` ### Fitting GP and RDRv For instructions on how to fit a GP or RDRv model for comparison see the [here](https://github.com/uzh-rpg/data_driven_mpc) ### Test the Fitted Model ``` python src/experiments/trajectory_test.py --model_version --model_name simple_sim_mlp --model_typ mlp_approx ``` where the `model_type` argument can be one of `mlp_approx`(Real-time Neural MPC), `mlp` (Naive Integration), `gp` (Gaussian Process Model). For a baseline comparison result run the same script without model parameters: ``` python src/experiments/trajectory_test.py ``` Multiple models can be compared at once via ``` python src/experiments/comparative_experiment.py --model_version --model_name --model_type --fast ``` Results are saved in the `results/` folder. ### Citing If you use this code in an academic context, please cite the following publication: ``` @article{salzmann2023neural, title={Real-time Neural-MPC: Deep Learning Model Predictive Control for Quadrotors and Agile Robotic Platforms}, author={Salzmann, Tim and Kaufmann, Elia and Arrizabalaga, Jon and Pavone, Marco and Scaramuzza, Davide and Ryll, Markus}, journal={IEEE Robotics and Automation Letters}, doi={10.1109/LRA.2023.3246839}, year={2023} } ``` ================================================ FILE: requirements.txt ================================================ numpy==2.2 scipy==1.15 tqdm matplotlib==3.9 scikit-learn==1.6 casadi==3.6 pyquaternion joblib pandas PyYAML rospkg>=1.3 torch==2.6 -i https://rospypi.github.io/simple/ rospy ================================================ FILE: ros_dd_mpc/CMakeLists.txt ================================================ cmake_minimum_required(VERSION 2.8.3) project(ros_dd_mpc) ## Compile as C++11, supported in ROS Kinetic and newer # add_compile_options(-std=c++11) ## Find catkin macros and libraries ## if COMPONENTS list like find_package(catkin REQUIRED COMPONENTS xyz) ## is used, also find other catkin packages find_package(catkin REQUIRED COMPONENTS rospy std_msgs message_generation ) ## System dependencies are found with CMake's conventions # find_package(Boost REQUIRED COMPONENTS system) ## Uncomment this if the package has a setup.py. This macro ensures ## modules and global scripts declared therein get installed ## See http://ros.org/doc/api/catkin/html/user_guide/setup_dot_py.html # catkin_python_setup() ################################################ ## Declare ROS messages, services and actions ## ################################################ ## To declare and build messages, services or actions from within this ## package, follow these steps: ## * Let MSG_DEP_SET be the set of packages whose message types you use in ## your messages/services/actions (e.g. std_msgs, actionlib_msgs, ...). ## * In the file package.xml: ## * add a build_depend tag for "message_generation" ## * add a build_depend and a exec_depend tag for each package in MSG_DEP_SET ## * If MSG_DEP_SET isn't empty the following dependency has been pulled in ## but can be declared for certainty nonetheless: ## * add a exec_depend tag for "message_runtime" ## * In this file (CMakeLists.txt): ## * add "message_generation" and every package in MSG_DEP_SET to ## find_package(catkin REQUIRED COMPONENTS ...) ## * add "message_runtime" and every package in MSG_DEP_SET to ## catkin_package(CATKIN_DEPENDS ...) ## * uncomment the add_*_files sections below as needed ## and list every .msg/.srv/.action file to be processed ## * uncomment the generate_messages entry below ## * add every package in MSG_DEP_SET to generate_messages(DEPENDENCIES ...) ## Generate messages in the 'msg' folder add_message_files( FILES ReferenceTrajectory.msg ) ## Generate services in the 'srv' folder # add_service_files( # FILES # Service1.srv # Service2.srv # ) ## Generate actions in the 'action' folder # add_action_files( # FILES # Action1.action # Action2.action # ) ## Generate added messages and services with any dependencies listed here generate_messages( DEPENDENCIES std_msgs # Or other packages containing msgs ) ################################################ ## Declare ROS dynamic reconfigure parameters ## ################################################ ## To declare and build dynamic reconfigure parameters within this ## package, follow these steps: ## * In the file package.xml: ## * add a build_depend and a exec_depend tag for "dynamic_reconfigure" ## * In this file (CMakeLists.txt): ## * add "dynamic_reconfigure" to ## find_package(catkin REQUIRED COMPONENTS ...) ## * uncomment the "generate_dynamic_reconfigure_options" section below ## and list every .cfg file to be processed ## Generate dynamic reconfigure parameters in the 'cfg' folder # generate_dynamic_reconfigure_options( # cfg/DynReconf1.cfg # cfg/DynReconf2.cfg # ) ################################### ## catkin specific configuration ## ################################### ## The catkin_package macro generates cmake config files for your package ## Declare things to be passed to dependent projects ## INCLUDE_DIRS: uncomment this if your package contains header files ## LIBRARIES: libraries you create in this project that dependent projects also need ## CATKIN_DEPENDS: catkin_packages dependent projects also need ## DEPENDS: system dependencies of this project that dependent projects also need catkin_package( # INCLUDE_DIRS include # LIBRARIES ros_python3_issues CATKIN_DEPENDS message_runtime # DEPENDS system_lib ) ########### ## Build ## ########### ## Specify additional locations of header files ## Your package locations should be listed before other locations include_directories( # include ${catkin_INCLUDE_DIRS} ) ## Declare a C++ library # add_library(${PROJECT_NAME} # src/${PROJECT_NAME}/ros_python3_issues.cpp # ) ## Add cmake target dependencies of the library ## as an example, code may need to be generated before libraries ## either from message generation or dynamic reconfigure # add_dependencies(${PROJECT_NAME} ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS}) ## Declare a C++ executable ## With catkin_make all packages are built within a single CMake context ## The recommended prefix ensures that target names across packages don't collide # add_executable(${PROJECT_NAME}_node src/ros_python3_issues_node.cpp) ## Rename C++ executable without prefix ## The above recommended prefix causes long target names, the following renames the ## target back to the shorter version for ease of user use ## e.g. "rosrun someones_pkg node" instead of "rosrun someones_pkg someones_pkg_node" # set_target_properties(${PROJECT_NAME}_node PROPERTIES OUTPUT_NAME node PREFIX "") ## Add cmake target dependencies of the executable ## same as for the library above # add_dependencies(${PROJECT_NAME}_node ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS}) ## Specify libraries to link a library or executable target against # target_link_libraries(${PROJECT_NAME}_node # ${catkin_LIBRARIES} # ) ############# ## Install ## ############# # all install targets should use catkin DESTINATION variables # See http://ros.org/doc/api/catkin/html/adv_user_guide/variables.html ## Mark executable scripts (Python etc.) for installation ## in contrast to setup.py, you can choose the destination # install(PROGRAMS # scripts/my_python_script # DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} # ) ## Mark executables and/or libraries for installation # install(TARGETS ${PROJECT_NAME} ${PROJECT_NAME}_node # ARCHIVE DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION} # LIBRARY DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION} # RUNTIME DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} # ) ## Mark cpp header files for installation # install(DIRECTORY include/${PROJECT_NAME}/ # DESTINATION ${CATKIN_PACKAGE_INCLUDE_DESTINATION} # FILES_MATCHING PATTERN "*.h" # PATTERN ".svn" EXCLUDE # ) ## Mark other files for installation (e.g. launch and bag files, etc.) # install(FILES # # myfile1 # # myfile2 # DESTINATION ${CATKIN_PACKAGE_SHARE_DESTINATION} # ) ############# ## Testing ## ############# ## Add gtest based cpp test target and link libraries # catkin_add_gtest(${PROJECT_NAME}-test test/test_ros_python3_issues.cpp) # if(TARGET ${PROJECT_NAME}-test) # target_link_libraries(${PROJECT_NAME}-test ${PROJECT_NAME}) # endif() ## Add folders to be run by python nosetests # catkin_add_nosetests(test) # if (CATKIN_ENABLE_TESTING) # find_package(rostest REQUIRED) # add_rostest(test/test_issue_rosunit.launch) # endif() ================================================ FILE: ros_dd_mpc/config/agisim_simulation_run.yaml ================================================ quad_name: 'kingfisher' world_limits: None use_ekf_synchronization: False control_freq_factor: 5 environment: 'agisim' ================================================ FILE: ros_dd_mpc/config/arena_limits.yaml ================================================ x_min: -6.0 x_max: 12.0 y_min: -7.0 y_max: 7.0 z_min: 1.0 z_max: 4.0 ================================================ FILE: ros_dd_mpc/config/arena_run.yaml ================================================ quad_name: 'kingfisher' world_limits: arena_limits use_ekf_synchronization: False control_freq_factor: 5 environment: 'arena' ================================================ FILE: ros_dd_mpc/config/basic.world ================================================ model://ground_plane model://sun EARTH_WGS84 47.3667 8.5500 500.0 0 quick 1000 1.3 0 0.2 100 0.001 0.01 0.5 50 0 0 -9.8 ================================================ FILE: ros_dd_mpc/config/circle_and_lemniscate_options.yaml ================================================ # Parameters for the loop and lemniscate trajectories at successively increasing speed loop_z: 2.5 # Z position of loop plane [m] loop_r: 5.0 # Radius of loop [m] loop_v_max: 10.0 # Maximum speed achieved (approx.) [m/s] loop_lin_a: 0.25 # Linear acceleration and deceleration of trajectory [m/s^2] loop_clockwise: False # Rotation direction of the quad loop_yawing: False # Yaw quadrotor along the circle (only for loop) ================================================ FILE: ros_dd_mpc/config/configuration_parameters.py ================================================ """ Set of tunable parameters for the Simplified Simulator and model fitting. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import os class DirectoryConfig: """ Class for storing directories within the package """ _dir_path = os.path.dirname(os.path.realpath(__file__)) SAVE_DIR = _dir_path + '/../results/model_fitting' RESULTS_DIR = _dir_path + '/../results' CONFIG_DIR = _dir_path + '' DATA_DIR = _dir_path + '/../data' class SimpleSimConfig: """ Class for storing the Simplified Simulator configurations. """ # Set to True to show a real-time Matplotlib animation of the experiments for the Simplified Simulator. Execution # will be slower if the GUI is turned on. Note: setting to True may require some further library installation work. custom_sim_gui = False # Set to True to display a plot describing the trajectory tracking results after the execution. result_plots = True # Set to True to show the trajectory that will be executed before the execution time pre_run_debug_plots = True # Choice of disturbances modeled in our Simplified Simulator. For more details about the parameters used refer to # the script: src/quad_mpc/quad_3d.py. simulation_disturbances = { "noisy": True, # Thrust and torque gaussian noises "drag": True, # 2nd order polynomial aerodynamic drag effect "payload": False, # Payload force in the Z axis "motor_noise": True # Asymmetric voltage noise in the motors } class ModelFitConfig: """ Class for storing flags for the model fitting scripts. """ # Dataset loading ## # ds_name = "simplified_sim_dataset" ds_metadata = { "noisy": True, "drag": True, "payload": False, "motor_noise": True } # ds_name = "agisim_dataset" # ds_metadata = { # "agisim": "default", # } # ds_name = "arena_dataset" # ds_metadata = { # "arena": "default", # } # ds_name = "neurobem_dataset" # ds_metadata = { # "arena": "default", # } # ## Visualization ## # # Training mode visualize_training_result = True visualize_data = False # Visualization mode grid_sampling_viz = True x_viz = [7, 8, 9] u_viz = [] y_viz = [7, 8, 9] # ## Data post-processing ## # histogram_bins = 40 # Cluster data using histogram binning histogram_threshold = 0.001 # Remove bins where the total ratio of data is lower than this threshold velocity_cap = 16 # Also remove datasets point if abs(velocity) > x_cap # ############# Experimental ############# # # ## Use fit model to generate synthetic data ## # use_dense_model = False dense_model_version = "" dense_model_name = "" dense_training_points = 200 # ## Clustering for multidimensional models ## # clusters = 1 load_clusters = False class GroundEffectMapConfig: """ Class for storing parameters for the ground effect map. """ resolution = 0.1 origin = (-4, 9) horizon = ((-7, 7), (-7, 7)) box_min = (-4.25, 9.37) box_max = (-2.76, 10.13) box_height = 0.7 ================================================ FILE: ros_dd_mpc/config/flyingroom_limits.yaml ================================================ x_min: -0.4 x_max: 4.3 y_min: -2.4 y_max: 2.3 z_min: 0.5 z_max: 2.0 ================================================ FILE: ros_dd_mpc/config/ground_effect_limits.yaml ================================================ x_min: -4.25 x_max: -2.76 y_min: 9.37 y_max: 10.13 z_min: 0.8 # 0.95 z_max: 0.85 ================================================ FILE: ros_dd_mpc/config/kingfisher.yaml ================================================ mass: 0.752 # [kg] tbm_fr: [ 0.075, -0.10, 0.0] # [m] tbm_bl: [-0.075, 0.10, 0.0] # [m] tbm_br: [-0.075, -0.10, 0.0] # [m] tbm_fl: [ 0.075, 0.10, 0.0] # [m] inertia: [0.0025, 0.0021, 0.0043] # [kgm^2] motor_omega_min: 150.0 # [rad/s] motor_omega_max: 2800.0 # [rad/s] motor_tau: 0.033 # [s] omega_max: [10.0, 10.0, 4.0] # [rad/s] comm_delay: 0.04 # [s] thrust_map: [1.562522e-6, 0.0, 0.0] kappa: 0.022 # [Nm/N] thrust_min: 0.0 # [N] thrust_max: 8.5 # [N] per motor rotors_config: "cross" aero_coeff_1: [0.0, 0.0, 0.0] # [0.26, 0.28, 0.42] # [N/(v/m)] aero_coeff_3: [0.0, 0.0, 0.0] # [N/(v/m)] aero_coeff_h: 0.0 # 0.01 [N/(v^2/m^2)] ================================================ FILE: ros_dd_mpc/config/simulation_run.yaml ================================================ quad_name: 'kingfisher' world_limits: None use_ekf_synchronization: False control_freq_factor: 5 environment: 'gazebo' ================================================ FILE: ros_dd_mpc/launch/dd_mpc_wrapper.launch ================================================ ================================================ FILE: ros_dd_mpc/msg/ReferenceTrajectory.msg ================================================ int32 seq_len string traj_name float64 v_input float64[] trajectory float64[] dt float64[] inputs ================================================ FILE: ros_dd_mpc/nodes/dd_mpc_node.py ================================================ #!/usr/bin/env python3.6 """ ROS node for the data-augmented MPC, to use in the Gazebo simulator and real world experiments. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import json import os import threading import time import numpy as np import pandas as pd import rosbag import rospy import std_msgs.msg from agiros_msgs.msg import Command from gazebo_msgs.srv import GetPhysicsProperties from geometry_msgs.msg import PoseStamped from nav_msgs.msg import Odometry from ros_dd_mpc.msg import ReferenceTrajectory from std_msgs.msg import Bool, Empty from tqdm import tqdm from src.experiments.point_tracking_and_record import make_record_dict, get_record_file_and_dir, check_out_data from src.model_fitting.rdrv_fitting import load_rdrv from src.quad_mpc.create_ros_dd_mpc import ROSDDMPC from src.utils.utils import jsonify, interpol_mse, quaternion_state_mse, load_pickled_models, \ separate_variables, get_model_dir_and_file from src.utils.visualization import trajectory_tracking_results, mse_tracking_experiment_plot, \ load_past_experiments, get_experiment_files def odometry_parse(odom_msg): p = [odom_msg.pose.pose.position.x, odom_msg.pose.pose.position.y, odom_msg.pose.pose.position.z] q = [odom_msg.pose.pose.orientation.w, odom_msg.pose.pose.orientation.x, odom_msg.pose.pose.orientation.y, odom_msg.pose.pose.orientation.z] v = [odom_msg.twist.twist.linear.x, odom_msg.twist.twist.linear.y, odom_msg.twist.twist.linear.z] w = [odom_msg.twist.twist.angular.x, odom_msg.twist.twist.angular.y, odom_msg.twist.twist.angular.z] return p, q, v, w def state_parse(state_msg): p = [state_msg.pose.position.x, state_msg.pose.position.y, state_msg.pose.position.z] q = [state_msg.pose.orientation.w, state_msg.pose.orientation.x, state_msg.pose.orientation.y, state_msg.pose.orientation.z] v = [state_msg.velocity.linear.x, state_msg.velocity.linear.y, state_msg.velocity.linear.z] w = [state_msg.velocity.angular.x, state_msg.velocity.angular.y, state_msg.velocity.angular.z] omega = state_msg.motor_speeds return p, q, v, w, omega def make_raw_optitrack_dict(): rec_dict_raw = make_record_dict(state_dim=7) # Remove unnecessary entries keys = list(rec_dict_raw.keys()) for key in keys: if key not in ["state_in", "timestamp"]: rec_dict_raw.pop(key) return rec_dict_raw def make_state_record_dict(state_dim): blank_state_recording_dict = { "state_in": np.zeros((0, state_dim)), "state_ref": np.zeros((0, state_dim)), "input_in": np.zeros((0, 4)), "timestamp": np.zeros((0, 1)), } return blank_state_recording_dict def odometry_skipped_warning(last_seq, current_seq, stage): skip_msg = "Odometry skipped at %s step. Last: %d, current: %d" % (stage, last_seq, current_seq) rospy.logwarn(skip_msg) class DDMPCWrapper: def __init__(self, quad_name, environment="agisim", recording_options=None, load_options=None, rdrv=None, plot=False, reset_experiment=False): if recording_options is None: recording_options = {"recording": False} # If on a simulation environment, figure out if physics are slowed down if environment == "gazebo": try: get_gazebo_physics = rospy.ServiceProxy('/gazebo/get_physics_properties', GetPhysicsProperties) resp = get_gazebo_physics() physics_speed = resp.max_update_rate * resp.time_step rospy.loginfo("Physics running at %.2f normal speed" % physics_speed) except rospy.ServiceException as e: print("Service call failed: %s" % e) physics_speed = 1 else: physics_speed = 1 self.physics_speed = physics_speed self.environment = environment self.plot = plot self.recording_options = recording_options # Control at 50 hz. Use time horizon=1 and 10 nodes self.n_mpc_nodes = rospy.get_param('~n_nodes', default=10) self.t_horizon = rospy.get_param('~t_horizon', default=1.0) self.control_freq_factor = rospy.get_param('~control_freq_factor', default=5 if environment == "gazebo" else 6) self.opt_dt = self.t_horizon / (self.n_mpc_nodes * self.control_freq_factor) # Load trained model mlp_conf = None if load_options is not None: if load_options['model_type'] == 'gp': rospy.loginfo("Attempting to load GP model from:\n git: {}\n name: {}\n meta: {}".format( load_options["git"], load_options["model_name"], load_options["params"])) pre_trained_models = load_pickled_models(model_options=load_options) else: import torch import ml_casadi.torch as mc from src.model_fitting.mlp_common import NormalizedMLP, QuadResidualModel directory, file_name = get_model_dir_and_file(load_options) saved_dict = torch.load(os.path.join(directory, f"{file_name}.pt"), map_location="cpu") model_type = load_options['model_type'] if '_qres' in model_type: mlp_model = QuadResidualModel(saved_dict['hidden_size'], saved_dict['hidden_layers']) else: mlp_model = mc.nn.MultiLayerPerceptron(saved_dict['input_size'], saved_dict['hidden_size'], saved_dict['output_size'], saved_dict['hidden_layers'], 'Tanh') model = NormalizedMLP(mlp_model, torch.tensor(np.zeros((saved_dict['input_size'],))).float(), torch.tensor(np.zeros((saved_dict['input_size'],))).float(), torch.tensor(np.zeros((saved_dict['output_size'],))).float(), torch.tensor(np.zeros((saved_dict['output_size'],))).float()) model.load_state_dict(saved_dict['state_dict']) model.eval() if 'gpu' in model_type: model = model.to('cuda:0') pre_trained_models = model mlp_conf = {'approximated': False, 'v_inp': True, 'u_inp': False, 'ground_map_input': False, 'torque_output': False} print(model_type) if model_type.endswith('approx') or model_type.endswith('approx1'): mlp_conf['approximated'] = True mlp_conf['approx_order'] = 1 if model_type.endswith('approx2'): mlp_conf['approximated'] = True mlp_conf['approx_order'] = 2 if '_u' in model_type: mlp_conf['u_inp'] = True if '_ge' in model_type: mlp_conf['ground_map_input'] = True if '_T' in model_type: mlp_conf['torque_output'] = True if '_qres' in model_type: mlp_conf['u_inp'] = True mlp_conf['torque_output'] = True if pre_trained_models is None: rospy.logwarn("Model parameters specified did not match with any pre-trained model") else: pre_trained_models = None self.pre_trained_models = pre_trained_models self.git_v = load_options["git"] if self.pre_trained_models is not None: rospy.loginfo("Successfully loaded model") self.model_name = load_options["model_name"] elif rdrv is not None: self.model_name = "rdrv" else: self.model_name = "Nominal" # Initialize MPC for point tracking self.dd_mpc = ROSDDMPC(self.t_horizon, self.n_mpc_nodes, self.opt_dt, quad_name=quad_name, model_conf=mlp_conf, point_reference=False, models=pre_trained_models, rdrv=rdrv) # Last state obtained from odometry self.x = None # Elapsed time between two recordings self.last_update_time = time.time() # Last references. Use hovering activation as input reference self.last_x_ref = None self.last_u_ref = None # Reference trajectory variables self.x_ref = None self.t_ref = None self.u_ref = None self.current_idx = 0 self.skipped_idx = [] self.quad_trajectory = None self.quad_controls = None self.w_control = None # Provisional reference for "waiting_for_reference" hovering mode self.x_ref_prov = None # To measure optimization elapsed time self.optimization_dt = 0 self.optimization_steps = 0 # Thread for MPC optimization self.mpc_thread = threading.Thread() # Trajectory tracking experiment. Dims: seed x av_v x n_samples if reset_experiment: self.metadata_dict = {} else: self.metadata_dict, _, _, _ = load_past_experiments() self.mse_exp = np.zeros((0, 0, 0, 1)) self.t_opt = np.zeros((0, 0, 0)) self.mse_exp_v_max = np.zeros((0, 0)) self.ref_traj_name = "" self.ref_v = 0 self.run_traj_counter = 0 # Keep track of status of MPC object self.odom_available = False # Binary variable to run MPC only once every other odometry callback self.optimize_next = True # Binary variable to completely skip an odometry if in flying arena self.skip_next = False # Remember the sequence number of the last odometry message received. self.last_odom_seq_number = 0 # Measure if trajectory starting point is reached self.x_initial_reached = False # Variables for recording mode self.recording_warmup = True self.x_pred = None self.w_opt = None # Get recording file and directory blank_recording_dict = make_record_dict(state_dim=13) if recording_options["recording"]: record_raw_optitrack = recording_options["record_raw_optitrack"] overwrite = recording_options["overwrite"] metadata = {self.environment: "default"} rec_dict, rec_file = get_record_file_and_dir( blank_recording_dict, recording_options, simulation_setup=metadata, overwrite=overwrite) state_rec_dict = make_state_record_dict(state_dim=13) # If flying with the optitrack system, also record raw optitrack estimates if self.environment == "flying_room" or self.environment == 'arena' and record_raw_optitrack: rec_dict_raw = make_raw_optitrack_dict() metadata = {self.environment: "optitrack_raw"} rec_dict_raw, rec_file_raw = get_record_file_and_dir( rec_dict_raw, recording_options, simulation_setup=metadata, overwrite=overwrite) else: rec_dict_raw = rec_file_raw = None if recording_options["record_raw_state_control"]: record_raw_state_control = True raw_state_control_bag = rosbag.Bag(os.path.splitext(rec_file)[0] + '.bag', 'w') else: state_rec_dict = None record_raw_optitrack = False rec_dict = rec_file = None rec_dict_raw = rec_file_raw = None record_raw_state_control = False raw_state_control_bag = None self.state_rec_dict = state_rec_dict self.rec_dict = rec_dict self.rec_file = rec_file self.rec_dict_raw = rec_dict_raw self.rec_file_raw = rec_file_raw self.record_raw_state_control = record_raw_state_control self.raw_state_control_bag = raw_state_control_bag self.landing = False self.override_land = False self.ground_level = False self.controller_off = False # Setup node publishers and subscribers. The odometry (sub) and control (pub) topics will vary depending on # which environment is being used odom_topic = "/" + quad_name + "/agiros_pilot/odometry" raw_topic = None if self.environment == "arena": raw_topic = "/vicon/" + quad_name land_topic = "/" + quad_name + "/agiros_pilot/land" control_topic = "/" + quad_name + "/agiros_pilot/feedthrough_command" off_topic = "/" + quad_name + "/agiros_pilot/off" reference_topic = "reference" status_topic = "busy" # Publishers self.control_pub = rospy.Publisher(control_topic, Command, queue_size=1, tcp_nodelay=True) self.status_pub = rospy.Publisher(status_topic, Bool, queue_size=1) self.off_pub = rospy.Publisher(off_topic, Empty, queue_size=1) # Subscribers self.land_sub = rospy.Subscriber(land_topic, Empty, self.land_callback) self.ref_sub = rospy.Subscriber(reference_topic, ReferenceTrajectory, self.reference_callback) self.odom_sub = rospy.Subscriber(odom_topic, Odometry, self.odometry_callback, queue_size=1, tcp_nodelay=True) if raw_topic is not None and record_raw_optitrack: self.raw_sub = rospy.Subscriber(raw_topic, PoseStamped, self.raw_odometry_callback) rate = rospy.Rate(1) while not rospy.is_shutdown(): # Publish if MPC is busy with a current trajectory msg = Bool() msg.data = not (self.x_ref is None and self.odom_available) self.status_pub.publish(msg) rate.sleep() def land_callback(self, _): """ Trigger landing sequence. :param _: empty message """ rospy.loginfo("Controller disabled. Landing with RPG MPC") self.override_land = True def rest_state(self): """ Set quad reference to hover state at position (0, 0, 0.1) """ self.last_x_ref = [[self.x[0], self.x[1], self.x[2]], [1, 0, 0, 0], [0, 0, 0], [0, 0, 0]] self.last_u_ref = self.dd_mpc.quad.g[-1] * self.dd_mpc.quad.mass / (self.dd_mpc.quad.max_thrust * 4) self.last_u_ref = self.last_u_ref[0] * np.array([1, 1, 1, 1]) def create_command_msg(self, w_opt, x_opt): """ Creates Command Msg from MPC output. """ next_control = Command() next_control.header = std_msgs.msg.Header() next_control.header.stamp = rospy.Time.now() next_control.is_single_rotor_thrust = True if self.environment == 'agisim' else False next_control.collective_thrust = np.sum(w_opt[:4]) * self.dd_mpc.quad.max_thrust / self.dd_mpc.quad.mass next_control.bodyrates.x = x_opt[1, -3] next_control.bodyrates.y = x_opt[1, -2] next_control.bodyrates.z = x_opt[1, -1] next_control.thrusts = w_opt[:4] * self.dd_mpc.quad.max_thrust return next_control def run_mpc(self, odom, recording=True): """ :param odom: message from subscriber. :type odom: Odometry :param recording: If False, some messages were skipped between this and last optimization. Don't record any data during this optimization if in recording mode. """ if not self.odom_available: return # Measure time between initial state was checked in and now dt = odom.header.stamp.to_time() - self.last_update_time model_data, x_guess, u_guess = self.set_reference() # Run MPC and publish control try: tic = time.time() w_opt, x_opt = self.dd_mpc.optimize(model_data) next_control = self.create_command_msg(w_opt, x_opt) self.optimization_dt += time.time() - tic self.optimization_steps += 1 if time.time() - tic > 0.01: print("MPC thread. Seq: %d. Topt: %.4f" % (odom.header.seq, (time.time() - tic) * 1000)) # print("MPC thread. Seq: %d. Topt: %.4f" % (odom.header.seq, (time.time() - tic) * 1000)) self.control_pub.publish(next_control) if self.x_initial_reached and self.current_idx < self.w_control.shape[0]: self.w_control[self.current_idx, 0] = next_control.bodyrates.x self.w_control[self.current_idx, 1] = next_control.bodyrates.y self.w_control[self.current_idx, 2] = next_control.bodyrates.z except KeyError: self.recording_warmup = True # Should not happen anymore. rospy.logwarn("Tried to run an MPC optimization but MPC is not ready yet.") return if w_opt is not None: # Check out final states. self.recording_warmup can only be true in recording mode. if not self.recording_warmup and recording and self.x_initial_reached: x_out = np.array(self.x)[np.newaxis, :] self.rec_dict = check_out_data(self.rec_dict, x_out, None, self.w_opt, dt) if self.record_raw_state_control: self.raw_state_control_bag.write('control', next_control) self.w_opt = w_opt if self.x_initial_reached and self.current_idx < self.quad_controls.shape[0]: self.quad_controls[self.current_idx, :] = np.expand_dims(self.w_opt[:4], axis=0) def check_out_initial_state(self, odom): """ Add the initial state to the recording dictionary and start counting until next optimization :param odom: message from subscriber. :type odom: Odometry """ if self.w_opt is not None: self.last_update_time = odom.header.stamp.to_time() self.rec_dict["state_in"] = np.append(self.rec_dict["state_in"], np.array(self.x)[np.newaxis, :], 0) self.rec_dict["timestamp"] = np.append(self.rec_dict["timestamp"], odom.header.stamp.to_time()) if self.current_idx < self.x_ref.shape[0]: self.rec_dict["state_ref"] = np.append(self.rec_dict["state_ref"], self.x_ref[np.newaxis, self.current_idx, :], 0) self.recording_warmup = False def reference_callback(self, msg): """ Callback for receiving a reference trajectory :param msg: message from subscriber :type msg: ReferenceTrajectory """ seq_len = msg.seq_len if seq_len == 0: # Hover-in-place mode self.x_ref = self.x[:7] self.u_ref = None self.t_ref = None off_msg = Empty() self.off_pub.publish(off_msg) self.controller_off = True # If this is the end of a reference tracking experiment, generate the results plot if self.mse_exp.shape[0] != 0 and self.run_traj_counter > 0: self.plot_tracking_mse_experiment() self.landing = False rospy.loginfo("No more references will be received") if self.raw_state_control_bag is not None: self.record_raw_state_control = False self.raw_state_control_bag.close() return # Save reference name self.ref_traj_name = msg.traj_name self.ref_v = msg.v_input # Save reference trajectory, relative times and inputs self.x_ref = np.array(msg.trajectory).reshape(seq_len, -1) self.t_ref = np.array(msg.dt) self.u_ref = np.array(msg.inputs).reshape(seq_len, -1) self.quad_trajectory = np.zeros((len(self.t_ref), len(self.x))) self.quad_controls = np.zeros((len(self.t_ref), 4)) self.w_control = np.zeros((len(self.t_ref), 3)) rospy.loginfo("New trajectory received. Time duration: %.2f s" % self.t_ref[-1]) def odometry_callback(self, msg): """ Callback function for Odometry subscriber :param msg: message from subscriber. :type msg: Odometry """ if self.controller_off: return p, q, v, w = odometry_parse(msg) self.x = p + q + v + w try: # Update the state estimate of the quad self.dd_mpc.set_state(self.x) except AttributeError: # The DD MPC object instantiation is still not finished return if self.override_land: return # If the above try passed then MPC is ready self.odom_available = True if self.record_raw_state_control: self.raw_state_control_bag.write('state', msg) # We only optimize once every two odometry messages if not self.optimize_next: self.mpc_thread.join() # If currently on trajectory tracking, pay close attention to any skipped messages. if self.x_initial_reached: # Count how many messages were skipped (ideally 0) skipped_messages = int(msg.header.seq - self.last_odom_seq_number - 1) if skipped_messages > 0: warn_msg = "Recording time skipped messages: %d" % skipped_messages rospy.logwarn(warn_msg) # Adjust current index in trajectory for i in range(divmod(skipped_messages, 2)[0]): if self.current_idx + i < self.x_ref.shape[0]: self.skipped_idx.append(self.current_idx + i) self.current_idx += divmod(skipped_messages, 2)[0] # If odd number of skipped messages, do optimization if skipped_messages > 0 and skipped_messages % 2 == 1: if self.recording_options["recording"]: self.check_out_initial_state(msg) # Run MPC now self.run_mpc(msg) self.last_odom_seq_number = msg.header.seq self.optimize_next = False return self.optimize_next = True if self.recording_options["recording"] and self.x_initial_reached: self.check_out_initial_state(msg) return # Run MPC if msg.header.seq > self.last_odom_seq_number + 2 and self.last_odom_seq_number > 0 and self.x_initial_reached: # If one message was skipped at this point, then the reference is already late. Compensate by # optimizing twice in a row and hope to do it fast... if self.recording_options["recording"] and self.x_initial_reached: self.check_out_initial_state(msg) self.run_mpc(msg) self.optimize_next = True self.last_odom_seq_number = msg.header.seq odometry_skipped_warning(self.last_odom_seq_number, msg.header.seq, "optimization") return def _thread_func(): self.run_mpc(msg) self.mpc_thread = threading.Thread(target=_thread_func(), args=(), daemon=True) self.mpc_thread.start() self.last_odom_seq_number = msg.header.seq self.optimize_next = False def raw_odometry_callback(self, msg): """ Callback function for the raw Optitrack subscriber. Adds the data to the raw data dictionary. :param msg: Raw data from Optitrack estimator :type msg: PoseStamped """ if not self.recording_options["recording"] or not self.x_initial_reached: return x = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z, msg.pose.orientation.w, msg.pose.orientation.x, msg.pose.orientation.y, msg.pose.orientation.z]) self.rec_dict_raw["state_in"] = np.append(self.rec_dict_raw["state_in"], x[np.newaxis, :], 0) self.rec_dict_raw["timestamp"] = np.append(self.rec_dict_raw["timestamp"], msg.header.stamp.to_time()) def hover_here(self, x): self.rest_state() x_ref = [x[:3], x[3:7], [0, 0, 0], [0, 0, 0]] u_ref = self.last_u_ref x_guess = np.tile(np.concatenate(x_ref)[np.newaxis, :], (self.n_mpc_nodes, 1)) u_guess = np.tile(self.last_u_ref[np.newaxis, :], (self.n_mpc_nodes, 1)) return self.dd_mpc.set_reference(x_ref, u_ref), x_guess, u_guess def set_reference(self): if self.environment == "gazebo" or self.environment== "agisim": th = 0.1 else: th = 0.2 mask = [1] * 9 + [0] * 3 x_ref = self.last_x_ref u_ref = self.last_u_ref x_guess = None u_guess = None if not self.odom_available: return # Check if landing mode if self.landing: dz = np.sign(0.3 - self.x[2]) dz = dz * 0.1 if self.environment != "agisim" else dz * 0.5 x_ref[0][2] = min(0.1, self.x[2] + dz) if dz > 0 else max(0.1, self.x[2] + dz) # Check if z position is close to target. if abs(self.x[2] - 0.2) < 0.05: executed_x_ref = self.x_ref executed_u_ref = self.u_ref executed_t_ref = self.t_ref self.x_ref = None self.u_ref = None self.t_ref = None self.x_initial_reached = False if self.recording_options["recording"]: self.save_recording_data() # Calculate MSE of position tracking and maximum axial velocity achieved rmse = interpol_mse(np.delete(executed_t_ref, self.skipped_idx, axis=0), np.delete(executed_x_ref[:, :3], self.skipped_idx, axis=0), np.delete(executed_t_ref, self.skipped_idx, axis=0), np.delete(self.quad_trajectory[:, :3], self.skipped_idx, axis=0)) self.optimization_dt /= self.optimization_steps if self.ref_traj_name in self.metadata_dict.keys(): if self.model_name in self.metadata_dict[self.ref_traj_name].keys(): self.metadata_dict[self.ref_traj_name][self.model_name][self.ref_v] = [rmse, self.optimization_dt] else: self.metadata_dict[self.ref_traj_name][self.model_name] = { self.ref_v: [rmse, self.optimization_dt]} else: self.metadata_dict[self.ref_traj_name] = { self.model_name: {self.ref_v: [rmse, self.optimization_dt]}} n_trajectories = len(self.metadata_dict.keys()) n_models = len(self.metadata_dict[self.ref_traj_name].keys()) n_vel = len(self.metadata_dict[self.ref_traj_name][self.model_name].keys()) # Figure out dimensions of data so far self.mse_exp = np.zeros((n_trajectories, n_vel, n_models, 1)) self.t_opt = np.zeros((n_trajectories, n_vel, n_models)) self.mse_exp_v_max = np.zeros((n_trajectories, n_vel)) # Add data to array # Dimensions of mse_exp: n_trajectories x n_average_speeds x n_models x n_sim_options for traj_id, traj_type in enumerate(self.metadata_dict.keys()): for model_id, model_type in enumerate(self.metadata_dict[traj_type].keys()): for vel_id, vel in enumerate(self.metadata_dict[traj_type][model_type].keys()): self.mse_exp[traj_id, vel_id, model_id, 0] = self.metadata_dict[traj_type][model_type][vel][0] self.t_opt[traj_id, vel_id, model_id] = self.optimization_dt self.mse_exp_v_max[traj_id, vel_id] = vel v_max = np.max(self.quad_trajectory[:, 7:10]) rospy.loginfo("Tracking complete. Total RMSE: %.5f m. Max axial vel: %.3f. " "Mean optimization time: %.3f ms" % (rmse, v_max, self.optimization_dt * 1000)) self.current_idx = 0 if self.plot: with_gp = ' + GP ' if self.pre_trained_models is not None else ' - GP ' tit = r'$v_{max}$=%.2f m/s | RMSE: %.4f | %s ' % (v_max, float(rmse), with_gp) trajectory_tracking_results(np.delete(executed_t_ref, self.skipped_idx, axis=0), np.delete(executed_x_ref, self.skipped_idx, axis=0), np.delete(self.quad_trajectory, self.skipped_idx, axis=0), np.delete(executed_u_ref, self.skipped_idx, axis=0), np.delete(self.quad_controls, self.skipped_idx, axis=0), w_control=np.delete(self.w_control, self.skipped_idx, axis=0), title=tit) rospy.loginfo('Saved Plot!') # Stop landing. Quad is close to ground level self.landing = False self.ground_level = True return self.dd_mpc.set_reference(x_ref, u_ref), x_guess, u_guess # Check if reference trajectory is set up. If not, pick current position and keep hover if self.x_ref is None: self.ground_level = False # We are waiting for a new reference. Set in provisional hover mode at current position if self.x_ref_prov is None: rospy.loginfo("Entering provisional hovering mode while to reference is available at: ") self.x_ref_prov = self.x rospy.loginfo(self.x_ref_prov) # Provisional hovering mode return self.hover_here(self.x_ref_prov) if self.x_ref_prov is not None: self.x_ref_prov = None rospy.loginfo("Abandoning provisional hovering mode.") # Check if reference is hovering mode if isinstance(self.x_ref, list): return self.hover_here(self.x_ref) # Trajectory tracking mode. Check if target reached if quaternion_state_mse(np.array(self.x), self.x_ref[0, :], mask) < th and not self.x_initial_reached: if self.record_raw_state_control: self.raw_state_control_bag.write('recording_ctrl', std_msgs.msg.Bool(True)) # Initial position of trajectory has been reached self.x_initial_reached = True self.odom_available = False self.optimization_dt = 0 rospy.loginfo("Reached initial position of trajectory.") model_data = self.dd_mpc.set_reference(separate_variables(self.x_ref[:1, :]), self.u_ref[:1, :]) return model_data, x_guess, u_guess # Raise the drone towards the initial position of the trajectory if not self.x_initial_reached: dx = 0.3 * np.sign(self.x_ref[0, 0] - self.x[0]) dy = 0.3 * np.sign(self.x_ref[0, 1] - self.x[1]) dz = 0.3 * np.sign(self.x_ref[0, 2] - self.x[2]) x_ref[0][0] = min(self.x_ref[0, 0], self.x[0] + dx) if dx > 0 else max(self.x_ref[0, 0], self.x[0] + dx) x_ref[0][1] = min(self.x_ref[0, 1], self.x[1] + dy) if dy > 0 else max(self.x_ref[0, 1], self.x[1] + dy) x_ref[0][2] = min(self.x_ref[0, 2], self.x[2] + dz) if dz > 0 else max(self.x_ref[0, 2], self.x[2] + dz) elif self.current_idx < self.x_ref.shape[0]: self.quad_trajectory[self.current_idx, :] = np.expand_dims(self.x, axis=0) # Trajectory tracking ref_traj = self.x_ref[self.current_idx:self.current_idx + self.n_mpc_nodes * self.control_freq_factor, :] ref_u = self.u_ref[self.current_idx:self.current_idx + self.n_mpc_nodes * self.control_freq_factor, :] # Indices for down-sampling the reference to number of MPC nodes downsample_ref_ind = np.arange(0, min(self.control_freq_factor * self.n_mpc_nodes, ref_traj.shape[0]), self.control_freq_factor, dtype=int) # Sparser references (same dt as node separation) x_ref = ref_traj[downsample_ref_ind, :] u_ref = ref_u[downsample_ref_ind, :] # Initial guesses u_guess = u_ref x_guess = x_ref while u_guess.shape[0] < self.n_mpc_nodes: x_guess = np.concatenate((x_guess, x_guess[-1:, :]), axis=0) u_guess = np.concatenate((u_guess, u_guess[-1:, :]), axis=0) x_ref = separate_variables(x_ref) self.current_idx += 1 # End of reference reached elif self.current_idx == self.x_ref.shape[0]: rospy.loginfo("Finished trajectory - Landing.") # Add one to the completed trajectory counter self.run_traj_counter += 1 # Lower drone to a safe height self.landing = True self.rest_state() x_ref = self.last_x_ref u_ref = self.last_u_ref # Stop recording self.x_initial_reached = False self.recording_warmup = True if self.record_raw_state_control: self.raw_state_control_bag.write('recording_ctrl', std_msgs.msg.Bool(False)) self.last_x_ref = x_ref self.last_u_ref = u_ref return self.dd_mpc.set_reference(x_ref, u_ref), x_guess, u_guess def plot_tracking_mse_experiment(self): metadata_file, _, _, _ = get_experiment_files() # Save data for reload with open(metadata_file, 'w') as json_file: json.dump(self.metadata_dict, json_file, indent=4) # Sort seeds dictionary by value traj_type_labels = [k for k in self.metadata_dict.keys()] model_type_labels = [k for k in self.metadata_dict[traj_type_labels[0]].keys()] mse_tracking_experiment_plot(v_max=self.mse_exp_v_max, mse=self.mse_exp, traj_type_vec=traj_type_labels, train_samples_vec=model_type_labels, legends=model_type_labels, y_labels=["RotorS"], t_opt=self.t_opt) def save_recording_data(self): # Remove exceeding data entry if needed if len(self.rec_dict['state_in']) > len(self.rec_dict['input_in']): self.rec_dict['state_in'] = self.rec_dict['state_in'][:-1] self.rec_dict['timestamp'] = self.rec_dict['timestamp'][:-1] # Compute predictions offline to avoid extra overhead while in trajectory tracking control rospy.loginfo("Filling in dataset and saving...") for i in tqdm(range(len(self.rec_dict['input_in']))): x_0 = self.rec_dict['state_in'][i] x_f = self.rec_dict['state_out'][i] u = self.rec_dict['input_in'][i] dt = self.rec_dict['dt'][i] x_pred, _ = self.dd_mpc.quad_mpc.forward_prop(x_0, u, t_horizon=dt) x_pred = x_pred[-1, np.newaxis, :] self.rec_dict['state_pred'] = np.append(self.rec_dict['state_pred'], x_pred, axis=0) self.rec_dict['error'] = np.append(self.rec_dict['error'], x_f - x_pred, axis=0) # Save datasets x_dim = self.rec_dict["state_in"].shape[1] for key in self.rec_dict.keys(): print(key, " ", self.rec_dict[key].shape) self.rec_dict[key] = jsonify(self.rec_dict[key]) df = pd.DataFrame(self.rec_dict) df.to_csv(self.rec_file, index=True, mode='a', header=False) if self.rec_dict_raw is not None: data_len = min(self.rec_dict_raw["state_in"].shape[0], len(self.rec_dict_raw["timestamp"])) # To ensure same length of all entries for key in self.rec_dict_raw.keys(): self.rec_dict_raw[key] = self.rec_dict_raw[key][:data_len] print(key, " ", self.rec_dict_raw[key].shape) self.rec_dict_raw[key] = jsonify(self.rec_dict_raw[key]) df = pd.DataFrame(self.rec_dict_raw) df.to_csv(self.rec_file_raw, index=True, mode='a', header=False) self.rec_dict_raw = make_raw_optitrack_dict() # Reset recording dictionaries self.rec_dict = make_record_dict(x_dim) self.state_rec_dict = make_state_record_dict(x_dim) def main(): rospy.init_node("dd_mpc") # Recording parameters recording_options = { "recording": rospy.get_param('~recording', default=True), "dataset_name": "deleteme", "training_split": True, "overwrite": True, "record_raw_optitrack": True, "record_raw_state_control": True, } dataset_name = rospy.get_param('~dataset_name', default=None) overwrite = rospy.get_param('~overwrite', default=None) training = rospy.get_param('~training_split', default=None) raw_optitrack = rospy.get_param('~record_raw_optitrack', default=None) raw_state_control = rospy.get_param('~record_raw_state_control', default=None) if dataset_name is not None: recording_options["dataset_name"] = dataset_name if overwrite is not None: recording_options["overwrite"] = overwrite if training is not None: recording_options["training_split"] = training if raw_optitrack is not None: recording_options["record_raw_optitrack"] = raw_optitrack if raw_state_control is not None: recording_options["record_raw_state_control"] = raw_state_control # Model loading parameters load_options = { "git": "b6e73a5", "model_name": "", "params": None } git_id = rospy.get_param('~model_version', default=None) model_name = rospy.get_param('~model_name', default=None) model_type = rospy.get_param('~model_type', default="gp") if git_id is not None: load_options["git"] = git_id if model_name is not None: load_options["model_name"] = str(model_name) if model_type is not None: load_options["model_type"] = str(model_type) plot = False plot = rospy.get_param('~plot', default=None) if rospy.get_param('~plot', default=None) is not None else plot env = rospy.get_param('~environment', default='gazebo') default_quad = "hummingbird" if env == "gazebo" else "kingfisher" load_options["params"] = {env: "default"} if model_type == "rdrv": rdrv = load_rdrv(model_options=load_options) else: rdrv = None quad_name = rospy.get_param('~quad_name', default=None) quad_name = quad_name if quad_name is not None else default_quad # Change if needed. This is currently the supported combination. if env == "gazebo": assert quad_name == "hummingbird" elif env == "agisim": assert quad_name == "kingfisher" # Reset experiments switch reset = rospy.get_param('~reset_experiment', default=True) DDMPCWrapper(quad_name, env, recording_options, load_options, rdrv=rdrv, plot=plot, reset_experiment=reset) if __name__ == "__main__": main() ================================================ FILE: ros_dd_mpc/nodes/reference_publisher_node.py ================================================ #!/usr/bin/env python3.6 """ Node wrapper for publishing trajectories for the MPC pipeline to track. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ from std_msgs.msg import Bool from ros_dd_mpc.msg import ReferenceTrajectory from src.quad_mpc.create_ros_dd_mpc import custom_quad_param_loader from src.utils.trajectories import loop_trajectory, random_trajectory, lemniscate_trajectory, flyover_trajectory, \ flyover_trajectory_collect import numpy as np import rospy class ReferenceGenerator: def __init__(self): self.gp_mpc_busy = True rospy.init_node("reference_generator") plot = rospy.get_param('~plot', default=True) quad_name = rospy.get_param('~quad_name', default='hummingbird') quad = custom_quad_param_loader(quad_name) # Configuration for random flight mode n_seeds = rospy.get_param('~n_seeds', default=1) v_list = rospy.get_param('~v_list', default=[3.5]) if isinstance(v_list, str): v_list = v_list.split('[')[1].split(']')[0] v_list = [float(v.strip()) for v in v_list.split(',')] # Select if generate "random" trajectories, "hover" mode or increasing speed "loop" mode mode = rospy.get_param('~mode', default="random") if mode != "random" and mode != "flyover": n_seeds = 1 # Load parameters of loop trajectory loop_r = rospy.get_param('~loop_r', default=1.5) loop_z = rospy.get_param('~loop_z', default=1) loop_v_max = rospy.get_param('~loop_v_max', default=3) loop_a = rospy.get_param('~loop_lin_a', default=0.075) loop_cc = rospy.get_param('~loop_clockwise', default=True) loop_yawing = rospy.get_param('~loop_yawing', default=True) # Load world limits if any map_limits = rospy.get_param('~world_limits', default=None) # Control at 50 hz. Use time horizon=1 and 10 nodes n_mpc_nodes = rospy.get_param('~n_nodes', default=10) t_horizon = rospy.get_param('~t_horizon', default=1.0) control_freq_factor = rospy.get_param('~control_freq_factor', default=5 if quad_name == "hummingbird" else 6) opt_dt = t_horizon / (n_mpc_nodes * control_freq_factor) reference_topic = "reference" status_topic = "busy" reference_pub = rospy.Publisher(reference_topic, ReferenceTrajectory, queue_size=1) rospy.Subscriber(status_topic, Bool, self.status_callback) v_ind = 0 seed = 0 # Calculate total number of trajectories n_trajectories = n_seeds * len(v_list) curr_trajectory_ind = 0 rate = rospy.Rate(0.2) while not rospy.is_shutdown(): if not self.gp_mpc_busy and mode == "hover": rospy.loginfo("Sending hover-in-place command") msg = ReferenceTrajectory() reference_pub.publish(msg) rospy.signal_shutdown("All trajectories were sent to the MPC") break if not self.gp_mpc_busy and curr_trajectory_ind == n_trajectories: msg = ReferenceTrajectory() reference_pub.publish(msg) rospy.signal_shutdown("All trajectories were sent to the MPC") break if not self.gp_mpc_busy and mode == "loop": rospy.loginfo("Sending increasing speed loop trajectory") x_ref, t_ref, u_ref = loop_trajectory(quad, opt_dt, v_max=loop_v_max, radius=loop_r, z=loop_z, lin_acc=loop_a, clockwise=loop_cc, map_name=map_limits, yawing=loop_yawing, plot=plot) msg = ReferenceTrajectory() msg.traj_name = "circle" msg.v_input = loop_v_max msg.seq_len = x_ref.shape[0] msg.trajectory = np.reshape(x_ref, (-1,)).tolist() msg.dt = t_ref.tolist() msg.inputs = np.reshape(u_ref, (-1,)).tolist() reference_pub.publish(msg) curr_trajectory_ind += 1 self.gp_mpc_busy = True elif not self.gp_mpc_busy and mode == "lemniscate": rospy.loginfo("Sending increasing speed lemniscate trajectory") x_ref, t_ref, u_ref = lemniscate_trajectory(quad, opt_dt, v_max=loop_v_max, radius=loop_r, z=loop_z, lin_acc=loop_a, clockwise=loop_cc, map_name=map_limits, yawing=loop_yawing, plot=plot) msg = ReferenceTrajectory() msg.traj_name = "lemniscate" msg.v_input = loop_v_max msg.seq_len = x_ref.shape[0] msg.trajectory = np.reshape(x_ref, (-1,)).tolist() msg.dt = t_ref.tolist() msg.inputs = np.reshape(u_ref, (-1,)).tolist() reference_pub.publish(msg) curr_trajectory_ind += 1 self.gp_mpc_busy = True elif not self.gp_mpc_busy and mode == "random": speed = v_list[v_ind] log_msg = "Random trajectory generator %d/%d. Seed: %d. Mean vel: %.3f m/s" % \ (curr_trajectory_ind + 1, n_trajectories, seed, speed) rospy.loginfo(log_msg) x_ref, t_ref, u_ref = random_trajectory(quad, opt_dt, seed=seed, speed=speed, map_name=map_limits, plot=plot) msg = ReferenceTrajectory() msg.traj_name = "random_" + str(seed) msg.v_input = speed msg.seq_len = x_ref.shape[0] msg.trajectory = np.reshape(x_ref, (-1, )).tolist() msg.dt = t_ref.tolist() msg.inputs = np.reshape(u_ref, (-1, )).tolist() reference_pub.publish(msg) curr_trajectory_ind += 1 self.gp_mpc_busy = True if v_ind + 1 < len(v_list): v_ind += 1 else: seed += 1 v_ind = 0 elif not self.gp_mpc_busy and mode == "flyover_collect": speed = v_list[v_ind] log_msg = "Flyover trajectory generator %d/%d. Seed: %d. Mean vel: %.3f m/s" % \ (curr_trajectory_ind + 1, n_trajectories, seed, speed) rospy.loginfo(log_msg) x_ref, t_ref, u_ref = flyover_trajectory_collect(quad, opt_dt, seed=seed, speed=speed, flyover_box_name=map_limits) msg = ReferenceTrajectory() msg.traj_name = "flyover_" + str(seed) msg.v_input = speed msg.seq_len = x_ref.shape[0] msg.trajectory = np.reshape(x_ref, (-1, )).tolist() msg.dt = t_ref.tolist() msg.inputs = np.reshape(u_ref, (-1, )).tolist() reference_pub.publish(msg) curr_trajectory_ind += 1 self.gp_mpc_busy = True if v_ind + 1 < len(v_list): v_ind += 1 else: seed += 1 v_ind = 0 elif not self.gp_mpc_busy and mode == "flyover": speed = v_list[v_ind] log_msg = "Flyover trajectory generator %d/%d. Seed: %d. Mean vel: %.3f m/s" % \ (curr_trajectory_ind + 1, n_trajectories, seed, speed) rospy.loginfo(log_msg) x_ref, t_ref, u_ref = flyover_trajectory(quad, opt_dt, seed=seed, speed=speed, flyover_box_name=map_limits) msg = ReferenceTrajectory() msg.traj_name = "flyover_" + str(seed) msg.v_input = speed msg.seq_len = x_ref.shape[0] msg.trajectory = np.reshape(x_ref, (-1, )).tolist() msg.dt = t_ref.tolist() msg.inputs = np.reshape(u_ref, (-1, )).tolist() reference_pub.publish(msg) curr_trajectory_ind += 1 self.gp_mpc_busy = True if v_ind + 1 < len(v_list): v_ind += 1 else: seed += 1 v_ind = 0 elif not self.gp_mpc_busy: raise ValueError("Unknown trajectory type: %s" % mode) rate.sleep() def status_callback(self, msg): """ Callback function for tracking if the dd_mpc node is busy :param msg: Message from the subscriber :type msg: Bool """ self.gp_mpc_busy = msg.data if __name__ == "__main__": ReferenceGenerator() ================================================ FILE: ros_dd_mpc/package.xml ================================================ ros_dd_mpc 0.0.0 The ros_dd_mpc package Guillem Torrente TODO message_generation message_runtime catkin rospy rospy rospy agiros_msgs deep_casadi ================================================ FILE: ros_dd_mpc/src/__init__.py ================================================ ================================================ FILE: ros_dd_mpc/src/experiments/__init__.py ================================================ ================================================ FILE: ros_dd_mpc/src/experiments/comparative_experiment.py ================================================ """ Runs the experimental setup to compare different data-learned models for the MPC on the Simplified Simulator. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import os.path import time import argparse import numpy as np import torch import ml_casadi.torch as mc from config.configuration_parameters import SimpleSimConfig from src.model_fitting.mlp_common import NormalizedMLP from src.quad_mpc.quad_3d_mpc import Quad3DMPC from src.quad_mpc.quad_3d import Quadrotor3D from src.utils.quad_3d_opt_utils import get_reference_chunk from src.utils.utils import load_pickled_models, interpol_mse, separate_variables, get_model_dir_and_file from src.utils.visualization import initialize_drone_plotter, draw_drone_simulation, trajectory_tracking_results, \ get_experiment_files from src.utils.visualization import mse_tracking_experiment_plot from src.utils.trajectories import random_trajectory, lemniscate_trajectory, loop_trajectory from src.model_fitting.rdrv_fitting import load_rdrv global model_num def prepare_quadrotor_mpc(simulation_options, version=None, name=None, reg_type="gp", quad_name=None, t_horizon=1.0, q_diagonal=None, r_diagonal=None, q_mask=None): """ Creates a Quad3DMPC for the custom simulator. @param simulation_options: Parameters for the Quadrotor3D object. @param version: loading version for the GP/RDRv model. @param name: name to load for the GP/RDRv model. @param reg_type: either `gp` or `rdrv`. @param quad_name: Name for the quadrotor. Default name will be used if not specified. @param t_horizon: Time horizon of MPC in seconds. @param q_diagonal: 12-dimensional diagonal of the Q matrix (p_xyz, a_xyz, v_xyz, w_xyz) @param r_diagonal: 4-dimensional diagonal of the R matrix (motor inputs 1-4) @param q_mask: State variable weighting mask (boolean). Which state variables compute towards state loss function? @return: A Quad3DMPC wrapper for the custom simulator. @rtype: Quad3DMPC """ # Default Q and R matrix for LQR cost if q_diagonal is None: q_diagonal = np.array([10, 10, 10, 0.1, 0.1, 0.1, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05]) if r_diagonal is None: r_diagonal = np.array([0.1, 0.1, 0.1, 0.1]) if q_mask is None: q_mask = np.array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]).T # Simulation integration step (the smaller the more "continuous"-like simulation. simulation_dt = 5e-4 # Number of MPC optimization nodes n_mpc_nodes = 10 # Calculate time between two MPC optimization nodes [s] node_dt = t_horizon / n_mpc_nodes # Quadrotor simulator my_quad = Quadrotor3D(**simulation_options) model_name = quad_name # Configuration for MLP Model mlp_conf = None if version is not None and name is not None: load_ops = {"params": simulation_options} load_ops.update({"git": version, "model_name": name}) # Load trained GP model if reg_type == "gp": pre_trained_models = load_pickled_models(model_options=load_ops) rdrv_d = None elif 'mlp' in reg_type: mlp_conf = {'approximated': False, 'v_inp': True, 'u_inp': False, 'T_out': False, 'ground_map_input': False, 'torque_output': False, 'two_step_rti': False} directory, file_name = get_model_dir_and_file(load_ops) saved_dict = torch.load(os.path.join(directory, f"{file_name}.pt")) mlp_model = mc.nn.MultiLayerPerceptron(saved_dict['input_size'], saved_dict['hidden_size'], saved_dict['output_size'], saved_dict['hidden_layers'], 'Tanh') model = NormalizedMLP(mlp_model, torch.tensor(np.zeros((saved_dict['input_size'],))).float(), torch.tensor(np.zeros((saved_dict['input_size'],))).float(), torch.tensor(np.zeros((saved_dict['output_size'],))).float(), torch.tensor(np.zeros((saved_dict['output_size'],))).float()) model.load_state_dict(saved_dict['state_dict']) model.eval() pre_trained_models = model rdrv_d = None if reg_type.endswith('approx2'): mlp_conf['approximated'] = True mlp_conf['approx_order'] = 2 elif reg_type.endswith('approx') or reg_type.endswith('approx_1'): mlp_conf['approximated'] = True mlp_conf['approx_order'] = 1 if '_u' in reg_type: mlp_conf['u_inp'] = True if '_T' in reg_type: mlp_conf['T_out'] = True else: rdrv_d = load_rdrv(model_options=load_ops) pre_trained_models = None else: pre_trained_models = rdrv_d = None if model_name is None: model_name = "my_quad_" + str(globals()['model_num']) globals()['model_num'] += 1 # Initialize quad MPC quad_mpc = Quad3DMPC(my_quad, t_horizon=t_horizon, optimization_dt=node_dt, simulation_dt=simulation_dt, q_cost=q_diagonal, r_cost=r_diagonal, n_nodes=n_mpc_nodes, pre_trained_models=pre_trained_models, model_name=model_name, q_mask=q_mask, rdrv_d_mat=rdrv_d, model_conf=mlp_conf) return quad_mpc def main(quad_mpc, av_speed, reference_type=None, plot=False): """ :param quad_mpc: :type quad_mpc: Quad3DMPC :param av_speed: :param reference_type: :param plot: :return: """ # Recover some necessary variables from the MPC object my_quad = quad_mpc.quad n_mpc_nodes = quad_mpc.n_nodes simulation_dt = quad_mpc.simulation_dt t_horizon = quad_mpc.t_horizon reference_over_sampling = 5 mpc_period = t_horizon / (n_mpc_nodes * reference_over_sampling) # Choose the reference trajectory: if reference_type == "loop": # Circular trajectory reference_traj, reference_timestamps, reference_u = loop_trajectory( quad=my_quad, discretization_dt=mpc_period, radius=5, z=1, lin_acc=av_speed * 0.25, clockwise=True, yawing=False, v_max=av_speed, map_name=None, plot=plot) elif reference_type == "lemniscate": # Lemniscate trajectory reference_traj, reference_timestamps, reference_u = lemniscate_trajectory( quad=my_quad, discretization_dt=mpc_period, radius=5, z=1, lin_acc=av_speed * 0.25, clockwise=True, yawing=False, v_max=av_speed, map_name=None, plot=plot) else: # Get a random smooth position trajectory reference_traj, reference_timestamps, reference_u = random_trajectory( quad=my_quad, discretization_dt=mpc_period, seed=reference_type["random"], speed=av_speed, plot=plot) # Set quad initial state equal to the initial reference trajectory state quad_current_state = reference_traj[0, :].tolist() my_quad.set_state(quad_current_state) real_time_artists = None if plot: # Initialize real time plot stuff world_radius = np.max(np.abs(reference_traj[:, :2])) * 1.2 real_time_artists = initialize_drone_plotter(n_props=n_mpc_nodes, quad_rad=my_quad.length, world_rad=world_radius, full_traj=reference_traj) start_time = time.time() max_simulation_time = 10000 ref_u = reference_u[0, :] quad_trajectory = np.zeros((len(reference_timestamps), len(quad_current_state))) u_optimized_seq = np.zeros((len(reference_timestamps), 4)) # Sliding reference trajectory initial index current_idx = 0 # Measure the MPC optimization time mean_opt_time = 0.0 # Measure total simulation time total_sim_time = 0.0 while (time.time() - start_time) < max_simulation_time and current_idx < reference_traj.shape[0]: quad_current_state = my_quad.get_state(quaternion=True, stacked=True) quad_trajectory[current_idx, :] = np.expand_dims(quad_current_state, axis=0) u_optimized_seq[current_idx, :] = np.reshape(ref_u, (1, -1)) # ##### Optimization runtime (outer loop) ##### # # Get the chunk of trajectory required for the current optimization ref_traj_chunk, ref_u_chunk = get_reference_chunk( reference_traj, reference_u, current_idx, n_mpc_nodes, reference_over_sampling) model_ind = quad_mpc.set_reference(x_reference=separate_variables(ref_traj_chunk), u_reference=ref_u_chunk) # Optimize control input to reach pre-set target t_opt_init = time.time() w_opt, x_pred = quad_mpc.optimize(use_model=model_ind, return_x=True) mean_opt_time += time.time() - t_opt_init # Select first input (one for each motor) - MPC applies only first optimized input to the plant ref_u = np.squeeze(np.array(w_opt[:4])) if len(quad_trajectory) > 0 and plot and current_idx > 0: draw_drone_simulation(real_time_artists, quad_trajectory[:current_idx, :], my_quad, targets=None, targets_reached=None, pred_traj=x_pred, x_pred_cov=None) simulation_time = 1e-8 # ##### Simulation runtime (inner loop) ##### # while simulation_time < mpc_period: simulation_time += simulation_dt total_sim_time += simulation_dt quad_mpc.simulate(ref_u) u_optimized_seq[current_idx, :] = np.reshape(ref_u, (1, -1)) current_idx += 1 quad_current_state = my_quad.get_state(quaternion=True, stacked=True) quad_trajectory[-1, :] = np.expand_dims(quad_current_state, axis=0) u_optimized_seq[-1, :] = np.reshape(ref_u, (1, -1)) # Average elapsed time per optimization mean_opt_time /= current_idx rmse = interpol_mse(reference_timestamps, reference_traj[:, :3], reference_timestamps, quad_trajectory[:, :3]) max_vel = np.max(np.sqrt(np.sum(reference_traj[:, 7:10] ** 2, 1))) with_gp = ' + GP ' if quad_mpc.gp_ensemble is not None else ' - GP ' title = r'$v_{max}$=%.2f m/s | RMSE: %.4f m | %s ' % (max_vel, float(rmse), with_gp) if plot: trajectory_tracking_results(reference_timestamps, reference_traj, quad_trajectory, reference_u, u_optimized_seq, title) return rmse, max_vel, mean_opt_time if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--model_version", type=str, default="", nargs="+", help="Versions to load for the regression models. By default it is an 8 digit git hash." "Must specify the version for each model separated by spaces.") parser.add_argument("--model_name", type=str, default="", nargs="+", help="Name of the regression models within the specified folders. " "Must specify the names for all models separated by spaces.") parser.add_argument("--model_type", type=str, default="", nargs="+", help="Type of regression models (GP or RDRv linear). " "Must be specified for all models separated by spaces.") parser.add_argument("--fast", dest="fast", action="store_true", help="Set to True to run a fast experiment with less velocity samples.") parser.add_argument("--print_results", dest="print_results", action="store_true", help="Print the results data frame.") parser.set_defaults(print_results=False) input_arguments = parser.parse_args() globals()['model_num'] = 0 # Trajectory options traj_type_vec = [{"random": 1}, "loop", "lemniscate"] traj_type_labels = ["Random", "Circle", "Lemniscate"] if input_arguments.fast: av_speed_vec = [[2.0, 3.5], [2.0, 12.0], [2.0, 12.0]] else: av_speed_vec = [[2.0, 2.7, 3.0, 3.2, 3.5], [2.0, 4.5, 7.0, 9.5, 12.0], [2.0, 4.5, 7.0, 9.5, 12.0]] # Model options git_list = input_arguments.model_version name_list = input_arguments.model_name type_list = input_arguments.model_type assert len(git_list) == len(name_list) == len(type_list) # Simulation options plot_sim = SimpleSimConfig.custom_sim_gui noisy_sim_options = SimpleSimConfig.simulation_disturbances perfect_sim_options = {"payload": False, "drag": False, "noisy": False, "motor_noise": False} model_vec = [ {"simulation_options": perfect_sim_options, "model": None}, {"simulation_options": noisy_sim_options, "model": None}] legends = ['perfect', 'nominal'] for git, m_name, gp_or_rdrv in zip(git_list, name_list, type_list): model_vec += [{"simulation_options": noisy_sim_options, "model": {"version": git, "name": m_name, "reg_type": gp_or_rdrv}}] legends += [gp_or_rdrv + ": " + m_name] y_label = "RMSE [m]" # Define result vectors mse = np.zeros((len(traj_type_vec), len(av_speed_vec[0]), len(model_vec))) v_max = np.zeros((len(traj_type_vec), len(av_speed_vec[0]))) t_opt = np.zeros((len(traj_type_vec), len(av_speed_vec[0]), len(model_vec))) for n_train_id, model_type in enumerate(model_vec): for traj_id, traj_type in enumerate(traj_type_vec): for v_id, speed in enumerate(av_speed_vec[traj_id]): if model_type["model"] is not None: custom_mpc = prepare_quadrotor_mpc(model_type["simulation_options"], **model_type["model"]) else: custom_mpc = prepare_quadrotor_mpc(model_type["simulation_options"]) traj_params = {"av_speed": speed, "reference_type": traj_type, "plot": plot_sim} mse[traj_id, v_id, n_train_id], traj_v, opt_dt = main(custom_mpc, **traj_params) t_opt[traj_id, v_id, n_train_id] += opt_dt if v_max[traj_id, v_id] == 0: v_max[traj_id, v_id] = traj_v _, err_file, v_file, t_file = get_experiment_files() np.save(err_file, mse) np.save(v_file, v_max) np.save(t_file, t_opt) # from src.utils.visualization import load_past_experiments # _, mse, v_max, t_opt = load_past_experiments() mse_tracking_experiment_plot(v_max, mse, traj_type_labels, model_vec, legends, [y_label], t_opt=t_opt, font_size=12) if input_arguments.print_results: import pandas as pd pd.options.display.float_format = "{:,.2f}".format track_dfs = [] for i, track in enumerate(traj_type_labels): index = pd.MultiIndex.from_arrays([[track]*len(av_speed_vec[i]), av_speed_vec[i], v_max[i]], names=['Track', 'v_avg', 'v_max']) track_df = pd.DataFrame(mse[i].T*1000, columns=index, index=legends) track_dfs.append(track_df) track_dfs = pd.concat(track_dfs, axis=1) print(track_dfs.to_string()) pd.options.display.float_format = "{:,.3f}".format t_df = pd.DataFrame(np.mean(np.mean(t_opt, axis=0), axis=0)*1000, columns=['avg'], index=legends) print(t_df.to_string()) ================================================ FILE: ros_dd_mpc/src/experiments/point_tracking_and_record.py ================================================ """ Executes aggressive maneuvers for collecting flight data on the Simplified Simulator to later train models on. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import os import sys import time import copy import argparse import itertools import pandas as pd import numpy as np import casadi as cs from src.utils.visualization import draw_drone_simulation, initialize_drone_plotter from src.experiments.comparative_experiment import prepare_quadrotor_mpc from src.utils.utils import safe_mknode_recursive, jsonify, euclidean_dist, get_data_dir_and_file from config.configuration_parameters import SimpleSimConfig def get_record_file_and_dir(record_dict_template, recording_options, simulation_setup, overwrite=True): dataset_name = recording_options["dataset_name"] training_split = recording_options["training_split"] # Directory and file name for data recording rec_file_dir, rec_file_name = get_data_dir_and_file(dataset_name, training_split, simulation_setup) overwritten = safe_mknode_recursive(rec_file_dir, rec_file_name, overwrite=overwrite) rec_dict = copy.deepcopy(record_dict_template) rec_file = os.path.join(rec_file_dir, rec_file_name) if overwrite or (not overwrite and not overwritten): for key in rec_dict.keys(): rec_dict[key] = jsonify(rec_dict[key]) df = pd.DataFrame(rec_dict) df.to_csv(rec_file, index=False, header=True) rec_dict = copy.deepcopy(record_dict_template) return rec_dict, rec_file def make_record_dict(state_dim): blank_recording_dict = { "state_in": np.zeros((0, state_dim)), "state_ref": np.zeros((0, state_dim)), "error": np.zeros((0, state_dim)), "input_in": np.zeros((0, 4)), "state_out": np.zeros((0, state_dim)), "state_pred": np.zeros((0, state_dim)), "timestamp": np.zeros((0, 1)), "dt": np.zeros((0, 1)), } return blank_recording_dict def check_out_data(rec_dict, state_final, x_pred, w_opt, dt): rec_dict["dt"] = np.append(rec_dict["dt"], dt) rec_dict["input_in"] = np.append(rec_dict["input_in"], w_opt[np.newaxis, :4], axis=0) rec_dict["state_out"] = np.append(rec_dict["state_out"], state_final, 0) if x_pred is not None: err = state_final - x_pred rec_dict["error"] = np.append(rec_dict["error"], err, axis=0) rec_dict["state_pred"] = np.append(rec_dict["state_pred"], x_pred[np.newaxis, :], axis=0) return rec_dict def sample_random_target(x_current, world_radius, aggressive=True): """ Creates a new target point to reach. :param x_current: current position of the quad. Only used if aggressive=True :param world_radius: radius of the area where points are sampled from :param aggressive: if aggressive=True, points will be sampled away from the current position. If False, then points will be sampled uniformly from the world area. :return: new sampled target point. A 3-dimensional numpy array. """ if aggressive: # Polar 3D coordinates theta = np.random.uniform(0, 2 * np.pi, 1) psi = np.random.uniform(0, 2 * np.pi, 1) r = 1 * world_radius + np.random.uniform(-0.5, 0.5, 1) * world_radius # Transform to cartesian x = r * np.sin(theta) * np.cos(psi) y = r * np.sin(theta) * np.sin(psi) z = r * np.cos(theta) return x_current + np.array([x, y, z]).reshape((1, 3)) else: return np.random.uniform(-world_radius, world_radius, (1, 3)) def main(model_options, recording_options, simulation_options, parameters): world_radius = 3 if parameters["initial_state"] is None: initial_state = [0.0, 0.0, 0.0] + [1, 0, 0, 0] + [0, 0, 0] + [0, 0, 0] else: initial_state = parameters["initial_state"] sim_starting_pos = initial_state quad_current_state = sim_starting_pos if parameters["preset_targets"] is not None: targets = parameters["preset_targets"] else: targets = sample_random_target(np.array(initial_state[:3]), world_radius, aggressive=recording_options["recording"]) quad_mpc = prepare_quadrotor_mpc(simulation_options, **model_options, t_horizon=0.5, q_mask=np.array([1, 1, 1, 0.01, 0.01, 0.01, 1, 1, 1, 0, 0, 0])) # Recover some necessary variables from the MPC object my_quad = quad_mpc.quad n_mpc_nodes = quad_mpc.n_nodes t_horizon = quad_mpc.t_horizon simulation_dt = quad_mpc.simulation_dt reference_over_sampling = 3 control_period = t_horizon / (n_mpc_nodes * reference_over_sampling) my_quad.set_state(quad_current_state) # Real time plot params n_forward_props = n_mpc_nodes plot_sim_traj = False x_pred = None w_opt = None initial_guess = None # The optimization should be faster or equal than the duration of the optimization time step assert control_period <= t_horizon / n_mpc_nodes state = quad_mpc.get_state() # ####### Recording mode code ####### # recording = recording_options["recording"] state_dim = state.shape[0] blank_recording_dict = make_record_dict(state_dim) # Get recording file and directory if recording: if parameters["real_time_plot"]: parameters["real_time_plot"] = False print("Turned off real time plot during recording mode.") rec_dict, rec_file = get_record_file_and_dir(blank_recording_dict, recording_options, simulation_options) else: rec_dict = rec_file = None # Generate necessary art pack for real time plot if parameters["real_time_plot"]: real_time_art_pack = initialize_drone_plotter(n_props=n_forward_props, quad_rad=my_quad.x_f, world_rad=world_radius) else: real_time_art_pack = None start_time = time.time() simulation_time = 0.0 # Simulation tracking stuff targets_reached = np.array([False for _ in targets]) quad_trajectory = np.array(quad_current_state).reshape(1, -1) n_iteration_count = 0 print("Targets reached: ", end='') # All targets loop while False in targets_reached and (time.time() - start_time) < parameters["max_sim_time"]: current_target_i = np.where(targets_reached == False)[0][0] current_target = targets[current_target_i] current_target_reached = False quad_target_state = [list(current_target), [1, 0, 0, 0], [0, 0, 0], [0, 0, 0]] model_ind = quad_mpc.set_reference(quad_target_state) # Provide an initial guess without the uncertainty prop. if initial_guess is None: initial_guess = quad_mpc.optimize(use_model=model_ind) initial_guess = quad_mpc.reshape_input_sequence(initial_guess) # MPC loop while not current_target_reached and (time.time() - start_time) < parameters["max_sim_time"]: # Emergency recovery (quad controller gone out of control lol) if np.any(state[7:10] > 14) or n_iteration_count > 100: n_iteration_count = 0 my_quad.set_state(list(itertools.chain.from_iterable(quad_target_state))) state = quad_mpc.get_state() if recording: rec_dict["state_in"] = np.append(rec_dict["state_in"], state.T, 0) rec_dict["timestamp"] = np.append(rec_dict["timestamp"], time.time() - start_time) stacked_ref = np.array(list(itertools.chain.from_iterable(quad_target_state)))[np.newaxis, :] rec_dict["state_ref"] = np.append(rec_dict["state_ref"], stacked_ref, 0) if simulation_time != 0.0: rec_dict = check_out_data(rec_dict, state.T, x_pred, w_opt, simulation_time) simulation_time = 0.0 # Optimize control input to reach pre-set target w_opt, x_pred_horizon = quad_mpc.optimize(use_model=model_ind, return_x=True) if np.any(w_opt > (my_quad.max_input_value + 0.01)) or np.any(w_opt < (my_quad.min_input_value - 0.01)): print("MPC constraints were violated") initial_guess = quad_mpc.reshape_input_sequence(w_opt) # Save initial guess for future optimization. It is a time-shift of the current optimized variables initial_guess = np.array(cs.vertcat(initial_guess[1:, :], cs.DM.zeros(4).T)) # Select first input (one for each motor) - MPC applies only first optimized input to the plant ref_u = np.squeeze(np.array(w_opt[:4])) if recording: # Integrate first input. Will be used as nominal model prediction during next save x_pred, _ = quad_mpc.forward_prop(np.squeeze(state), w_opt=w_opt[:4], t_horizon=control_period, use_gp=False) x_pred = x_pred[-1, :] if parameters["real_time_plot"]: prop_params = {"x_0": np.squeeze(state), "w_opt": w_opt, "use_model": model_ind, "t_horizon": t_horizon} x_int, _ = quad_mpc.forward_prop(**prop_params, use_gp=False) if plot_sim_traj: x_sim = quad_mpc.simulate_plant(quad_mpc.reshape_input_sequence(w_opt)) else: x_sim = None draw_drone_simulation(real_time_art_pack, quad_trajectory, my_quad, targets, targets_reached, x_sim, x_int, x_pred_horizon, follow_quad=False) while simulation_time < control_period: # Simulation runtime (inner loop) simulation_time += simulation_dt quad_mpc.simulate(ref_u) quad_current_state = quad_mpc.get_state() # Target is reached if euclidean_dist(current_target[0:3], quad_current_state[0:3, 0], thresh=0.05): print("*", end='') sys.stdout.flush() n_iteration_count = 0 # Check out data immediately as new target will be optimized in next step if recording and len(rec_dict['state_in']) > len(rec_dict['input_in']): x_pred, _ = quad_mpc.forward_prop(np.squeeze(state), w_opt=w_opt[:4], t_horizon=simulation_time, use_gp=False) x_pred = x_pred[-1, :] rec_dict = check_out_data(rec_dict, quad_mpc.get_state().T, x_pred, w_opt, simulation_time) # Reset optimization time -> Ask for new optimization for next target in next dt simulation_time = 0.0 # Mark current target as reached current_target_reached = True targets_reached[current_target_i] = True # Remove initial guess initial_guess = None if parameters["preset_targets"] is None: new_target = sample_random_target(quad_current_state[:3], world_radius, aggressive=recording) targets = np.append(targets, new_target, axis=0) targets_reached = np.append(targets_reached, False) # Reset PID integral and past errors quad_mpc.reset() break n_iteration_count += 1 if parameters["real_time_plot"]: quad_trajectory = np.append(quad_trajectory, quad_current_state.T, axis=0) if len(quad_trajectory) > 300: quad_trajectory = np.delete(quad_trajectory, 0, 0) # Current target was reached. Remove incomplete recordings if recording: if len(rec_dict['state_in']) > len(rec_dict['input_in']): rec_dict["state_in"] = rec_dict["state_in"][:-1] rec_dict["timestamp"] = rec_dict["timestamp"][:-1] rec_dict["state_ref"] = rec_dict["state_ref"][:-1] for key in rec_dict.keys(): rec_dict[key] = jsonify(rec_dict[key]) df = pd.DataFrame(rec_dict) df.to_csv(rec_file, index=True, mode='a', header=False) rec_dict = copy.deepcopy(blank_recording_dict) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--model_version", type=str, default="", help="Version to load for the regression models. By default it is an 8 digit git hash.") parser.add_argument("--model_name", type=str, default="", help="Name of the regression model within the specified folder.") parser.add_argument("--model_type", type=str, default="gp", choices=["gp", "rdrv"], help="Type of regression model (GP or RDRv linear)") parser.add_argument("--recording", dest="recording", action="store_true", help="Set to True to enable recording mode.") parser.set_defaults(recording=False) parser.add_argument("--dataset_name", type=str, default="simplified_sim_dataset", help="Name for the generated dataset.") parser.add_argument("--test_split", dest="test_split", action="store_true", help="If the data set is the training or test split.") parser.set_defaults(test_split=False) parser.add_argument("--simulation_time", type=float, default=300, help="Total duration of the simulation in seconds.") args = parser.parse_args() np.random.seed(123 if args.test_split else 456) acados_config = { "solver_type": "SQP", "terminal_cost": True } run_options = { "model_options": { "version": args.model_version, "name": args.model_name, "reg_type": args.model_type, "quad_name": "my_quad" }, "recording_options": { "recording": args.recording, "dataset_name": args.dataset_name, "training_split": not args.test_split, }, "simulation_options": SimpleSimConfig.simulation_disturbances, "parameters": { "real_time_plot": SimpleSimConfig.custom_sim_gui, "max_sim_time": args.simulation_time, "preset_targets": None, "initial_state": None, "acados_options": acados_config } } main(**run_options) ================================================ FILE: ros_dd_mpc/src/experiments/trajectory_test.py ================================================ """ Tracks a specified trajectory on the simplified simulator using the data-augmented MPC. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import time import json import argparse import numpy as np from tqdm import tqdm from src.utils.utils import separate_variables from src.utils.quad_3d_opt_utils import get_reference_chunk from src.utils.trajectories import loop_trajectory, lemniscate_trajectory, check_trajectory from src.utils.visualization import initialize_drone_plotter, draw_drone_simulation, trajectory_tracking_results from src.experiments.comparative_experiment import prepare_quadrotor_mpc from config.configuration_parameters import SimpleSimConfig def main(args): params = { "version": args.model_version, "name": args.model_name, "reg_type": args.model_type, "quad_name": "my_quad" } # Load the disturbances for the custom offline simulator. simulation_options = SimpleSimConfig.simulation_disturbances debug_plots = SimpleSimConfig.pre_run_debug_plots tracking_results_plot = SimpleSimConfig.result_plots sim_gui = SimpleSimConfig.custom_sim_gui quad_mpc = prepare_quadrotor_mpc(simulation_options, **params) # Recover some necessary variables from the MPC object my_quad = quad_mpc.quad n_mpc_nodes = quad_mpc.n_nodes t_horizon = quad_mpc.t_horizon simulation_dt = quad_mpc.simulation_dt reference_over_sampling = 5 control_period = t_horizon / (n_mpc_nodes * reference_over_sampling) if args.trajectory == "loop": reference_traj, reference_timestamps, reference_u = loop_trajectory( my_quad, control_period, radius=args.trajectory_radius, z=1, lin_acc=args.acceleration, clockwise=True, yawing=False, v_max=args.max_speed, map_name=None, plot=debug_plots) elif args.trajectory == "lemniscate": reference_traj, reference_timestamps, reference_u = lemniscate_trajectory( my_quad, control_period, radius=args.trajectory_radius, z=1, lin_acc=args.acceleration, clockwise=True, yawing=False, v_max=args.max_speed, map_name=None, plot=debug_plots) else: raise ValueError("Unknown trajectory {}. Options are `lemniscate` and `loop`".format(args.trajectory)) if not check_trajectory(reference_traj, reference_u, reference_timestamps, debug_plots): return # Set quad initial state equal to the initial reference trajectory state quad_current_state = reference_traj[0, :].tolist() my_quad.set_state(quad_current_state) real_time_artists = None if sim_gui: # Initialize real time plot stuff world_radius = np.max(np.abs(reference_traj[:, :2])) * 1.2 real_time_artists = initialize_drone_plotter(n_props=n_mpc_nodes, quad_rad=my_quad.length, world_rad=world_radius, full_traj=reference_traj) ref_u = reference_u[0, :] quad_trajectory = np.zeros((len(reference_timestamps), len(quad_current_state))) u_optimized_seq = np.zeros((len(reference_timestamps), 4)) # Sliding reference trajectory initial index current_idx = 0 # Measure the MPC optimization time mean_opt_time = 0.0 # Measure total simulation time total_sim_time = 0.0 print("\nRunning simulation...") for current_idx in tqdm(range(reference_traj.shape[0])): quad_current_state = my_quad.get_state(quaternion=True, stacked=True) quad_trajectory[current_idx, :] = np.expand_dims(quad_current_state, axis=0) # ##### Optimization runtime (outer loop) ##### # # Get the chunk of trajectory required for the current optimization ref_traj_chunk, ref_u_chunk = get_reference_chunk( reference_traj, reference_u, current_idx, n_mpc_nodes, reference_over_sampling) # Set the reference for the OCP model_ind = quad_mpc.set_reference(x_reference=separate_variables(ref_traj_chunk), u_reference=ref_u_chunk) # Optimize control input to reach pre-set target t_opt_init = time.time() w_opt, x_pred = quad_mpc.optimize(use_model=model_ind, return_x=True) mean_opt_time += time.time() - t_opt_init # Select first input (one for each motor) - MPC applies only first optimized input to the plant ref_u = np.squeeze(np.array(w_opt[:4])) u_optimized_seq[current_idx, :] = np.reshape(ref_u, (1, -1)) if len(quad_trajectory) > 0 and sim_gui and current_idx > 0: draw_drone_simulation(real_time_artists, quad_trajectory[:current_idx, :], my_quad, targets=None, targets_reached=None, pred_traj=x_pred, x_pred_cov=None) simulation_time = 0.0 # ##### Simulation runtime (inner loop) ##### # while simulation_time < control_period: simulation_time += simulation_dt total_sim_time += simulation_dt quad_mpc.simulate(ref_u) u_optimized_seq[current_idx, :] = np.reshape(ref_u, (1, -1)) quad_current_state = my_quad.get_state(quaternion=True, stacked=True) quad_trajectory[-1, :] = np.expand_dims(quad_current_state, axis=0) u_optimized_seq[-1, :] = np.reshape(ref_u, (1, -1)) # Average elapsed time per optimization mean_opt_time = mean_opt_time / current_idx * 1000 tracking_rmse = np.mean(np.sqrt(np.sum((reference_traj[:, :3] - quad_trajectory[:, :3]) ** 2, axis=1))) if tracking_results_plot: v_max = np.max(reference_traj[:, 7:10]) with_gp = ' + GP ' if quad_mpc.gp_ensemble is not None else ' - GP ' title = r'$v_{max}$=%.2f m/s | RMSE: %.4f m | %s ' % (v_max, float(tracking_rmse), with_gp) trajectory_tracking_results(reference_timestamps, reference_traj, quad_trajectory, reference_u, u_optimized_seq, title) v_max_abs = np.max(np.sqrt(np.sum(reference_traj[:, 7:10] ** 2, 1))) print("\n:::::::::::::: SIMULATION SETUP ::::::::::::::\n") print("Simulation: Applied disturbances: ") print(json.dumps(simulation_options)) if quad_mpc.gp_ensemble is not None: print("\nModel: Using GP regression model: ", params["version"] + '/' + params["name"]) elif quad_mpc.mlp is not None: print("\nModel: Using MLP regression model: ", params["version"] + '/' + params["name"]) elif quad_mpc.rdrv is not None: print("\nModel: Using RDRv regression model: ", params["version"] + '/' + params["name"]) else: print("\nModel: No regression model loaded") print("\nReference: Executed trajectory", '`' + args.trajectory + '`', "with a peak axial velocity of", args.max_speed, "m/s, and a maximum speed of %2.3f m/s" % v_max_abs) print("\n::::::::::::: SIMULATION RESULTS :::::::::::::\n") print("Mean optimization time: %.3f ms" % mean_opt_time) print("Tracking RMSE: %.4f m\n" % tracking_rmse) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--model_version", type=str, default="", help="Version to load for the regression models. By default it is an 8 digit git hash.") parser.add_argument("--model_name", type=str, default="", help="Name of the regression model within the specified folder.") parser.add_argument("--model_type", type=str, default="gp", help="Type of regression model (GP or RDRv linear)") parser.add_argument("--trajectory", type=str, default="loop", choices=["loop", "lemniscate"], help='path to other necessary data files (eg. vocabularies)') parser.add_argument("--max_speed", type=float, default=8, help="Maximum axial speed executed during the flight in m/s. For the `loop` trajectory, " "velocities are feasible up to 14 m/s, and for the `lemniscate` up to 8 m/s") parser.add_argument("--acceleration", type=float, default=1, help="Acceleration of the reference trajectory. Higher accelerations will shorten the execution" "time of the tracking") parser.add_argument("--trajectory_radius", type=float, default=5, help="Radius of the reference trajectories") input_arguments = parser.parse_args() main(input_arguments) ================================================ FILE: ros_dd_mpc/src/model_fitting/__init__.py ================================================ ================================================ FILE: ros_dd_mpc/src/model_fitting/gp.py ================================================ """ Gaussian Process custom implementation for the data-augmented MPC. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import numpy as np import casadi as cs import joblib from tqdm import tqdm from operator import itemgetter from numpy.linalg import inv, cholesky, lstsq from numpy.random import mtrand from scipy.optimize import minimize from scipy.spatial.distance import pdist, cdist, squareform from src.utils.utils import safe_mknode_recursive, make_bz_matrix class CustomKernelFunctions: def __init__(self, kernel_func, params=None): self.params = params self.kernel_type = kernel_func if self.kernel_type == 'squared_exponential': if params is None: self.params = {'l': [1.0], 'sigma_f': 1.0} self.kernel = self.squared_exponential_kernel else: raise NotImplementedError("only squared_exponential is supported") self.theta = self.get_trainable_parameters() def __call__(self, x_1, x_2): return self.kernel(x_1, x_2) def __str__(self): if self.kernel_type == 'squared_exponential': len_scales = np.reshape(self.params['l'], -1) len_scale_str = '[' for i in range(len(len_scales)): len_scale_str += '%.3f, ' % len_scales[i] if i < len(len_scales) - 1 else '%.3f' % len_scales[i] len_scale_str += ']' summary = '%.3f' % self.params['sigma_f'] summary += '**2*RBF(length_scale=' + len_scale_str + ')' return summary else: raise NotImplementedError("only squared_exponential is supported") def get_trainable_parameters(self): trainable_params = [] if self.kernel_type == 'squared_exponential': trainable_params += \ np.reshape(np.squeeze(self.params['l']), -1).tolist() if hasattr(self.params['l'], "__len__") \ else [self.params['l']] trainable_params += [self.params['sigma_f']] return trainable_params @staticmethod def _check_length_scale(x, length_scale): length_scale = np.squeeze(length_scale).astype(float) if np.ndim(length_scale) > 1: raise ValueError("length_scale cannot be of dimension greater than 1") if np.ndim(length_scale) == 1 and x.shape[1] != length_scale.shape[0]: raise ValueError("Anisotropic kernel must have the same number of dimensions as data (%d!=%d)" % (length_scale.shape[0], x.shape[1])) return length_scale def squared_exponential_kernel(self, x_1, x_2=None): """ Anisotropic (diagonal length-scale) matrix squared exponential kernel. Computes a covariance matrix from points in x_1 and x_2. Args: x_1: Array of m points (m x d). x_2: Array of n points (n x d). Returns: Covariance matrix (m x n). """ if isinstance(x_2, cs.MX): return self._squared_exponential_kernel_cs(x_1, x_2) # Length scale parameter len_scale = self.params['l'] if 'l' in self.params.keys() else 1.0 # Vertical variation parameter sigma_f = self.params['sigma_f'] if 'sigma_f' in self.params.keys() else 1.0 x_1 = np.atleast_2d(x_1) length_scale = self._check_length_scale(x_1, len_scale) if x_2 is None: dists = pdist(x_1 / length_scale, metric='sqeuclidean') k = sigma_f * np.exp(-.5 * dists) # convert from upper-triangular matrix to square matrix k = squareform(k) np.fill_diagonal(k, 1) else: dists = cdist(x_1 / length_scale, x_2 / length_scale, metric='sqeuclidean') k = sigma_f * np.exp(-.5 * dists) return k def _squared_exponential_kernel_cs(self, x_1, x_2): """ Symbolic implementation of the anisotropic squared exponential kernel :param x_1: Array of m points (m x d). :param x_2: Array of n points (m x d). :return: Covariance matrix (m x n). """ # Length scale parameter len_scale = self.params['l'] if 'l' in self.params.keys() else 1.0 # Vertical variation parameter sigma_f = self.params['sigma_f'] if 'sigma_f' in self.params.keys() else 1.0 if x_1.shape != x_2.shape and x_2.shape[0] == 1: tiling_ones = cs.MX.ones(x_1.shape[0], 1) d = x_1 - cs.mtimes(tiling_ones, x_2) dist = cs.sum2(d ** 2 / cs.mtimes(tiling_ones, cs.MX(len_scale ** 2).T)) else: d = x_1 - x_2 dist = cs.sum1(d ** 2 / cs.MX(len_scale ** 2)) return sigma_f * cs.SX.exp(-.5 * dist) def diff(self, z, z_train): """ Computes the symbolic differentiation of the kernel function, evaluated at point z and using the training dataset z_train. This function implements equation (80) from overleaf document, without the c^{v_x} vector, and for all the partial derivatives possible (m), instead of just one. :param z: evaluation point. Symbolic vector of length m :param z_train: training dataset. Symbolic matrix of shape n x m :return: an m x n matrix, which is the derivative of the exponential kernel function evaluated at point z against the training dataset. """ if self.kernel_type != 'squared_exponential': raise NotImplementedError len_scale = self.params['l'] if len(self.params['l']) > 0 else self.params['l'] * cs.MX.ones(z_train.shape[1]) len_scale = np.atleast_2d(len_scale ** 2) # Broadcast z vector to have the shape of z_train (tile z to to the number of training points n) z_tile = cs.mtimes(cs.MX.ones(z_train.shape[0], 1), z.T) # Compute k_zZ. Broadcast it to shape of z_tile and z_train, i.e. by the number of variables in z. k_zZ = cs.mtimes(cs.MX.ones(z_train.shape[1], 1), self.__call__(z_train, z.T).T) return - k_zZ * (z_tile - z_train).T / cs.mtimes(cs.MX.ones(z_train.shape[0], 1), len_scale).T class CustomGPRegression: def __init__(self, x_features, u_features, reg_dim, mean=None, y_mean=None, kernel=None, sigma_n=1e-8, n_restarts=1): """ :param x_features: list of indices for the quadrotor state-derived features :param u_features: list of indices for the input state-derived features :param reg_dim: state dimension that this regressor is meant to be used for. :param mean: prior mean value :param y_mean: average y value for data normalization :param kernel: Kernel Function object :param sigma_n: noise sigma value :param n_restarts: number of optimization re-initializations """ if kernel is None: kernel = CustomKernelFunctions('squared_exponential') # Avoid non-invertible error assert sigma_n != 0.0 # For inference time self.x_features = x_features self.u_features = u_features self.reg_dim = reg_dim self.kernel_ = kernel self.kernel_type = kernel.kernel_type # Noise variance prior self.sigma_n = sigma_n # GP center of local feature space self.mean = mean self.y_mean = y_mean # Pre-computed training data kernel self._K = np.zeros((0, 0)) self._K_inv = np.zeros((0, 0)) self._K_inv_y = np.zeros((0, )) # Training dataset memory self._x_train = np.zeros((0, 0)) self._y_train = np.zeros((0, )) # CasADi symbolic equivalents self._K_cs = None self._K_inv_cs = None self._K_inv_y_cs = None self._x_train_cs = None self._y_train_cs = None self.sym_pred = None self.sym_jacobian_dz = None self.n_restarts = n_restarts @property def kernel(self): return self.kernel_ @kernel.setter def kernel(self, ker): self.kernel_ = ker @property def K(self): return self._K @K.setter def K(self, k): self._K = k self._K_cs = cs.DM(k) @property def K_inv(self): return self._K_inv @K_inv.setter def K_inv(self, k): self._K_inv = k self._K_inv_cs = cs.DM(k) @property def K_inv_y(self): return self._K_inv_y @K_inv_y.setter def K_inv_y(self, k): self._K_inv_y = k self._K_inv_y_cs = cs.DM(k) @property def x_train(self): return self._x_train @x_train.setter def x_train(self, k): self._x_train = k self._x_train_cs = cs.DM(k) @property def y_train(self): return self._y_train @y_train.setter def y_train(self, k): self._y_train = k self._y_train_cs = cs.DM(k) def log_marginal_likelihood(self, theta): l_params = np.exp(theta[:-1]) sigma_f = np.exp(theta[-1]) sigma_n = self.sigma_n kernel = CustomKernelFunctions(self.kernel_type, params={'l': l_params, 'sigma_f': sigma_f}) k_train = kernel(self.x_train, self.x_train) + sigma_n ** 2 * np.eye(len(self.x_train)) l_mat = cholesky(k_train) nll = np.sum(np.log(np.diagonal(l_mat))) + \ 0.5 * self.y_train.T.dot(lstsq(l_mat.T, lstsq(l_mat, self.y_train, rcond=None)[0], rcond=None)[0]) + \ 0.5 * len(self.x_train) * np.log(2 * np.pi) return nll def _nll(self, x_train, y_train): """ Returns a numerically stable function implementation of the negative log likelihood using the cholesky decomposition of the kernel matrix. http://www.gaussianprocess.org/gpml/chapters/RW2.pdf, Section 2.2, Algorithm 2.1. :param x_train: Array of m points (m x d). :param y_train: Array of m points (m x 1) :return: negative log likelihood (scalar) computing function """ def nll_func(theta): l_params = np.exp(theta[:-2]) sigma_f = np.exp(theta[-2]) sigma_n = np.exp(theta[-1]) kernel = CustomKernelFunctions(self.kernel_type, params={'l': l_params, 'sigma_f': sigma_f}) k_train = kernel(x_train, x_train) + sigma_n ** 2 * np.eye(len(x_train)) l_mat = cholesky(k_train) nll = np.sum(np.log(np.diagonal(l_mat))) + \ 0.5 * y_train.T.dot(lstsq(l_mat.T, lstsq(l_mat, y_train, rcond=None)[0], rcond=None)[0]) + \ 0.5 * len(x_train) * np.log(2 * np.pi) return nll return nll_func def _constrained_minimization(self, x_train, y_train, x_0, bounds): try: res = minimize(self._nll(x_train, y_train), x0=x_0, bounds=bounds, method='L-BFGS-B') return np.exp(res.x), res.fun except np.linalg.LinAlgError: return x_0, np.inf def fit(self, x_train, y_train): """ Fit a GP regressor to the training dataset by minimizing the negative log likelihood of the training data :param x_train: Array of m points (m x d). :param y_train: Array of m points (m x 1) """ initial_guess = self.kernel.get_trainable_parameters() initial_guess += [self.sigma_n] initial_guess = np.array(initial_guess) bounds = [(1e-5, 1e1) for _ in range(len(initial_guess) - 1)] bounds = bounds + [(1e-8, 1e0)] log_bounds = np.log(tuple(bounds)) y_train -= self.y_mean optima = [self._constrained_minimization(x_train, y_train, initial_guess, log_bounds)] if self.n_restarts > 1: random_state = mtrand._rand for iteration in range(self.n_restarts - 1): theta_initial = random_state.uniform(log_bounds[:, 0], log_bounds[:, 1]) optima.append(self._constrained_minimization(x_train, y_train, theta_initial, log_bounds)) lml_values = list(map(itemgetter(1), optima)) theta_opt = optima[int(np.argmin(lml_values))][0] # Update kernel with new parameters l_new = theta_opt[:-2] sigma_f_new = theta_opt[-2] self.sigma_n = theta_opt[-1] self.kernel = CustomKernelFunctions(self.kernel_type, params={'l': l_new, 'sigma_f': sigma_f_new}) # Pre-compute kernel matrices self.K = self.kernel(x_train, x_train) + self.sigma_n ** 2 * np.eye(len(x_train)) self.K_inv = inv(self.K) self.K_inv_y = self.K_inv.dot(y_train) # Update training dataset points self.x_train = x_train self.y_train = y_train self.compute_gp_jac() def compute_gp_jac(self): self.sym_jacobian_dz = self._linearized_state_estimate() def eval_gp_jac(self, z): if self.sym_jacobian_dz is None: self.compute_gp_jac() return self.sym_jacobian_dz(z) def _linearized_state_estimate(self): """ Computes the symbolic linearization of the GP prediction expected state with respect to the inputs of the GP itself. :return: a CasADi function that computes the linearized GP prediction estimate wrt the input features of the GP regressor itself. The output of the function is a vector of shape m, where m is the number of regression features. """ if self.kernel_type != 'squared_exponential': raise NotImplementedError # Symbolic variable for input state z = cs.MX.sym('z', self.x_train.shape[1]) # Compute the kernel derivative: dgpdz = cs.mtimes(self.kernel.diff(z, self._x_train_cs), self._K_inv_y_cs) return cs.Function('f', [z], [dgpdz], ['z'], ['dgpdz']) def predict(self, x_test, return_std=False, return_cov=False): """ Computes the sufficient statistics of the GP posterior predictive distribution from m training data X_train and Y_train and n new inputs X_s. Args: x_test: test input locations (n x d). return_std: boolean - return a standard deviation vector of size n return_cov: boolean - return a covariance vector of size n (sqrt of standard deviation) Returns: Posterior mean vector (n) and covariance diagonal or standard deviation vectors (n) if also requested. """ # Ensure at least n=1 x_test = np.atleast_2d(x_test) if isinstance(x_test, np.ndarray) else x_test if isinstance(x_test, cs.MX): return self._predict_sym(x_test=x_test, return_std=return_std, return_cov=return_cov) if isinstance(x_test, cs.DM): x_test = np.array(x_test).T k_s = self.kernel(x_test, self.x_train) k_ss = self.kernel(x_test, x_test) + 1e-8 * np.eye(len(x_test)) # Posterior mean value mu_s = k_s.dot(self.K_inv_y) + self.y_mean # Posterior covariance cov_s = k_ss - k_s.dot(self.K_inv).dot(k_s.T) std_s = np.sqrt(np.diag(cov_s)) if not return_std and not return_cov: return mu_s # Return covariance if return_cov: return mu_s, std_s ** 2 # Return standard deviation return mu_s, std_s def _predict_sym(self, x_test, return_std=False, return_cov=False): """ Computes the GP posterior mean and covariance at a given a test sample using CasADi symbolics. :param x_test: vector of size mx1, where m is the number of features used for the GP prediction :return: the posterior mean (scalar) and covariance (scalar). """ k_s = self.kernel(self._x_train_cs, x_test.T) # Posterior mean value mu_s = cs.mtimes(k_s.T, self._K_inv_y_cs) + self.y_mean if not return_std and not return_cov: return {'mu': mu_s} k_ss = self.kernel(x_test, x_test) + 1e-8 * cs.MX.eye(x_test.shape[1]) # Posterior covariance cov_s = k_ss - cs.mtimes(cs.mtimes(k_s.T, self._K_inv_cs), k_s) cov_s = cs.diag(cov_s) if return_std: return {'mu': mu_s, 'std': np.sqrt(cov_s)} return {'mu': mu_s, 'cov': cov_s} def sample_y(self, x_test, n_samples=3): """ Sample a number of functions from the process at the given test points :param x_test: test input locations (n x d). :param n_samples: integer, number of samples to draw :return: the drawn samples from the gaussian process. An array of shape n x n_samples """ mu, cov = self.predict(x_test, return_cov=True) # Draw three samples from the prior samples = np.random.multivariate_normal(mu.ravel(), np.diag(cov), n_samples) return samples.T def save(self, path): """ Saves the current GP regressor to the specified path as a pickle file. Must be re-loaded with the function load :param path: absolute path to save the regressor to """ saved_vars = { "kernel_params": self.kernel.params, "kernel_type": self.kernel.kernel_type, "x_train": self.x_train, "y_train": self.y_train, "k_inv_y": self.K_inv_y, "k_inv": self.K_inv, "sigma_n": self.sigma_n, "reg_dim": self.reg_dim, "x_features": self.x_features, "u_features": self.u_features, "mean": self.mean, "y_mean": self.y_mean } split_path = path.split('/') directory = '/'.join(split_path[:-1]) file = split_path[-1] safe_mknode_recursive(directory, file, overwrite=True) with open(path, 'wb') as f: joblib.dump(saved_vars, f) def load(self, data_dict): """ Load a pre-trained GP regressor :param data_dict: a dictionary with all the pre-trained matrices of the GP regressor """ self.K_inv = data_dict['k_inv'] self.K_inv_y = data_dict['k_inv_y'] self.x_train = data_dict['x_train'] self.y_train = data_dict['y_train'] self.kernel_type = data_dict['kernel_type'] self.kernel = CustomKernelFunctions(self.kernel_type, data_dict['kernel_params']) self.sigma_n = data_dict['sigma_n'] self.mean = data_dict['mean'] if 'mean' in data_dict.keys() else np.array([0, 0, 0]) self.y_mean = data_dict['y_mean'] if 'y_mean' in data_dict.keys() else np.array(0) self.compute_gp_jac() class GPEnsemble: def __init__(self): """ Sets up a GP regression ensemble. This essentially divides the prediction domain into different GP's, so that less training samples need to be used per GP. """ self.out_dim = 0 self.n_models_dict = {} # Make index to dim to make dimensions iterable self.dim_idx = np.zeros(0, dtype=int) # Dictionary of lists. Each element of the dictionary is indexed by the index of the GP output in the state # space, and contains a with all the GP's (one per cluster) used in that dimension. self.gp = {} # Store the centroids of all GP's self.gp_centroids = {} # Store the B_z matrices self.B_z_dict = {} # Whether the same clustering is used for all dimensions or not self.homogeneous = True # Whether the GP model has no ensembles in it (i.e. no GP has more than 1 cluster) self.no_ensemble = True @property def n_models(self): if self.homogeneous or self.no_ensemble: return self.n_models_dict[next(iter(self.n_models_dict))] return self.n_models_dict @property def B_z(self): return self.B_z_dict[next(iter(self.B_z_dict))] if self.homogeneous else self.B_z_dict def add_model(self, gp): """" :param gp: A list of n CustomGPRegression objects, where n is the number of GP's used to divide the feature space domain of the dimension in particular. :type gp: list """ # Check if dimension is already "occupied" by another GP gp_dim = gp[0].reg_dim if gp_dim in self.gp.keys(): raise ValueError("This dimension is already taken by another GP") self.out_dim += 1 self.dim_idx = np.append(self.dim_idx, gp_dim) self.gp[gp_dim] = np.array(gp) # Store centroids and sort along first dimension for easier comparison self.gp_centroids[gp_dim] = np.array([gp_cluster.mean for gp_cluster in gp]) sorted_cluster_idx = np.argsort(self.gp_centroids[gp_dim][:, 0]) self.gp_centroids[gp_dim] = self.gp_centroids[gp_dim][sorted_cluster_idx] self.gp[gp_dim] = self.gp[gp_dim][sorted_cluster_idx] # Calculate if Ensemble is still homogeneous self.homogeneous = self.homogeneous_feature_space() # Check if current gp is an actual ensemble self.n_models_dict[gp_dim] = len(gp) if len(gp) > 1: self.no_ensemble = False # Pre-compute B_z matrix self.B_z_dict[gp_dim] = make_bz_matrix(x_dims=13, u_dims=4, x_feats=gp[0].x_features, u_feats=gp[0].u_features) def get_z(self, x, u, dim): """ Computes the z features from the x and u vectors, and the target output dimension. :param x: state vector. Shape 13x1. Can be np.array or cs.MX. :param u: control input vector. Shape 4x1. Can be np.array or cs.MX. :param dim: output dimension target. :return: A vector of shape mx1 of the same format as inputs. m is determined by the B_z matrix for dim. """ # Get input feature indices x_feats = self.gp[dim][0].x_features u_feats = self.gp[dim][0].u_features # Stack into a single matrix if isinstance(x, np.ndarray): z = np.concatenate((x[x_feats], u[u_feats]), axis=0) elif isinstance(x, cs.MX): z = cs.mtimes(self.B_z_dict[dim], cs.vertcat(x, u)) else: raise TypeError return z def predict(self, x_test, u_test, return_std=False, return_cov=False, return_gp_id=False, return_z=False, progress_bar=False, gp_idx=None): """ Runs GP inference. First, select the GP optimally for the test samples. Then, run inference on that GP. :param x_test: array of shape d x n. n is the number of test samples and d their dimension. :param u_test: array of shape d' x n. n is the number of test samples and d' their dimension. :param return_std: True if also return the standard deviation of the GP inference. :param return_cov: True if also return the covariance of the GP inference. :param return_gp_id: True if also return the id of the GP used for inference. :param return_z: True if also return the z features computed for inference. :param progress_bar: If True, a progress bar will be shown when evaluating the test data. :param gp_idx: Dictionary of indices with the same length as the GP output dimension indicating which GP to use for each one. If None, select best based on x_test. :type gp_idx: dict :return: m x n arrays, where m is the output dimension and n is the number of samples. """ if return_std: assert not return_cov, "Can only return the std or the cov" if return_cov: assert not return_std, "Can only return the std or the cov" # Dictionary for function return outputs = {} # Build regression features and evaluation cluster indices for each GP output dimension z = {} gp_idx = {} if gp_idx is None else gp_idx if not self.homogeneous: for dim in self.gp.keys(): z[dim] = self.get_z(x_test, u_test, dim) if dim not in gp_idx.keys(): # Calculate optimal GP clusters to use for each test point gp_idx[dim] = self.select_gp(z=z[dim], dim=dim) gp_idx[dim] = np.atleast_1d(gp_idx[dim]) else: z_ = self.get_z(x_test, u_test, self.dim_idx[0]) z = {k: v for k, v in zip(self.dim_idx, [z_] * self.out_dim)} if not bool(gp_idx): gp_idx_ = self.select_gp(z=z_, dim=self.dim_idx[0]) gp_idx = {k: v for k, v in zip(self.dim_idx, [gp_idx_] * self.out_dim)} # Add stuff to output dictionary if return_z: outputs["z"] = z if return_gp_id: outputs["gp_id"] = gp_idx pred = [] cov_or_std = [] noise_prior = [] # Test data loop range_data = tqdm(range(x_test.shape[1])) if progress_bar else range(x_test.shape[1]) for j in range_data: pred_j = [] cov_or_std_j = [] noise_prior_j = [] # Output dim loop for dim in self.gp.keys(): out = self.gp[dim][gp_idx[dim][j]].predict(z[dim][:, j], return_std, return_cov) if not return_std and not return_cov: if isinstance(out, dict): pred_j += [out['mu']] else: pred_j += [out] else: if isinstance(out, dict): pred_j += [out['mu']] cov_or_std_j += [out['cov'] if 'cov' in out.keys() else out['std']] else: pred_j += [out[0]] cov_or_std_j += [out[1]] noise_prior_j += [np.array([self.gp[dim][gp_idx[dim][j]].sigma_n])] pred += [pred_j] cov_or_std += [cov_or_std_j] noise_prior += [noise_prior_j] # Convert to CasADi symbolic or numpy matrix depending on the input type pred = cs.horzcat(*[cs.vertcat(*pred[i]) for i in range(x_test.shape[1])]) \ if isinstance(x_test, cs.MX) else np.squeeze(np.array(pred)).T outputs["pred"] = pred if not return_std and not return_cov: return outputs # Convert to CasADi symbolic or numpy matrix depending on the input type cov_or_std = cs.horzcat(*[cs.vertcat(*cov_or_std[i]) for i in range(x_test.shape[1])]) \ if isinstance(x_test, cs.MX) else np.squeeze(np.array(cov_or_std)).T noise_prior = cs.horzcat(*[cs.vertcat(*noise_prior[i]) for i in range(x_test.shape[1])]) \ if isinstance(x_test, cs.MX) else np.squeeze(np.array(noise_prior)).T outputs["cov_or_std"] = cov_or_std outputs["noise_cov"] = noise_prior return outputs def select_gp(self, dim, x=None, u=None, z=None): """ Selects the best GP's for computing inference at the given test points x for a given regression output dimension. This calculation is done by computing the distance of all n test points to all available GP's centroids and selecting the one minimizing the Euclidean distance. :param z: np array of shape d x n corresponding to the processed feature vector. If unknown one may call this method with x and u instead. :param x: np array of shape 13 x n corresponding to the query quadrotor states. Only necessary if z=None. :param u: np.array of shape 4 x n corresponding to the query quadrotor control vectors. Only necessary if z=None. :param dim: index of GP output dimension. If None, evaluate on all dimensions. :return: a numpy vector of length n, indicating which GP to use for every test sample x. """ if dim is None: dim = self.dim_idx if isinstance(dim, np.ndarray): # If the ensemble is homogeneous only one evaluation is necessary if self.homogeneous or self.no_ensemble: return self.select_gp(dim[0], x, u, z)[0] return np.array([self.select_gp(i, x, u, z) for i in dim]) if z is None: z = self.get_z(x, u, dim) z = np.atleast_2d(z) centroids = self.gp_centroids[dim] # Select subset of features for current dimension return np.argmin(np.sqrt(np.sum((z[np.newaxis, :, :] - centroids[:, :, np.newaxis]) ** 2, 1)), 0) def homogeneous_feature_space(self): if self.out_dim == 1: return True equal_clusters = True last_centroids = None for i, key in enumerate(self.gp_centroids.keys()): centroids = self.gp_centroids[key] if i == 0: last_centroids = centroids continue if np.any(last_centroids != centroids): equal_clusters = False break last_centroids = centroids return equal_clusters ================================================ FILE: ros_dd_mpc/src/model_fitting/gp_common.py ================================================ """ Contains a set of utility functions and classes for the GP training and inference. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import copy import os import joblib import numpy as np import pandas as pd from sklearn.mixture import GaussianMixture from src.model_fitting.gp import GPEnsemble, CustomGPRegression as npGPRegression from src.utils.utils import undo_jsonify, prune_dataset, safe_mknode_recursive, get_data_dir_and_file, \ separate_variables, v_dot_q, quaternion_inverse from src.utils.visualization import visualize_data_distribution class GPDataset: def __init__(self, raw_ds=None, x_features=None, u_features=None, y_dim=None, cap=None, n_bins=None, thresh=None, visualize_data=False): self.data_labels = [r'$p_x$', r'$p_y$', r'$p_z$', r'$q_w$', r'$q_x$', r'$q_y$', r'$q_z$', r'$v_x$', r'$v_y$', r'$v_z$', r'$w_x$', r'$w_y$', r'$w_z$'] # Raw dataset data self.x_raw = None self.x_out_raw = None self.u_raw = None self.y_raw = None self.x_pred_raw = None self.dt_raw = None self.x_features = x_features self.u_features = u_features self.y_dim = y_dim # Data pruning parameters self.cap = cap self.n_bins = n_bins self.thresh = thresh self.plot = visualize_data # GMM clustering self.pruned_idx = [] self.cluster_agency = None self.centroids = None # Number of data clusters self.n_clusters = 1 if raw_ds is not None: self.load_data(raw_ds) self.prune() def load_data(self, ds): if isinstance(ds, np.lib.npyio.NpzFile): x_raw = ds['state_in'] x_out = ds['state_out'] x_pred = ds['state_pred'] u_raw = ds['input_in'] dt = ds["dt"] else: x_raw = undo_jsonify(ds['state_in'].to_numpy()) x_out = undo_jsonify(ds['state_out'].to_numpy()) x_pred = undo_jsonify(ds['state_pred'].to_numpy()) u_raw = undo_jsonify(ds['input_in'].to_numpy()) dt = ds["dt"].to_numpy() invalid = np.where(dt == 0) # Remove invalid entries (dt = 0) x_raw = np.delete(x_raw, invalid, axis=0) x_out = np.delete(x_out, invalid, axis=0) x_pred = np.delete(x_pred, invalid, axis=0) u_raw = np.delete(u_raw, invalid, axis=0) dt = np.delete(dt, invalid, axis=0) # Rotate velocities to body frame and recompute errors x_raw = world_to_body_velocity_mapping(x_raw) x_pred = world_to_body_velocity_mapping(x_pred) x_out = world_to_body_velocity_mapping(x_out) y_err = x_out - x_pred # Normalize error by window time (i.e. predict error dynamics instead of error itself) y_err /= np.expand_dims(dt, 1) # Select features self.x_raw = x_raw self.x_out_raw = x_out self.u_raw = u_raw self.y_raw = y_err self.x_pred_raw = x_pred self.dt_raw = dt def prune(self): # Prune noisy data if self.cap is not None and self.n_bins is not None and self.thresh is not None: x_interest = np.array([7, 8, 9]) y_interest = np.array([7, 8, 9]) labels = [self.data_labels[dim] for dim in np.atleast_1d(y_interest)] prune_x_data = self.get_x(pruned=False, raw=True)[:, x_interest] prune_y_data = self.get_y(pruned=False, raw=True)[:, y_interest] self.pruned_idx.append(prune_dataset(prune_x_data, prune_y_data, self.cap, self.n_bins, self.thresh, plot=self.plot, labels=labels)) def get_x(self, cluster=None, pruned=True, raw=False): if cluster is not None: assert pruned if raw: return self.x_raw[tuple(self.pruned_idx)] if pruned else self.x_raw if self.x_features is not None and self.u_features is not None: x_f = self.x_features u_f = self.u_features data = np.concatenate((self.x_raw[:, x_f], self.u_raw[:, u_f]), axis=1) if u_f else self.x_raw[:, x_f] else: data = np.concatenate((self.x_raw, self.u_raw), axis=1) data = data[:, np.newaxis] if len(data.shape) == 1 else data if pruned or cluster is not None: data = data[tuple(self.pruned_idx)] data = data[self.cluster_agency[cluster]] if cluster is not None else data return data @property def x(self): return self.get_x() def get_x_out(self, cluster=None, pruned=True): if cluster is not None: assert pruned if pruned or cluster is not None: data = self.x_out_raw[tuple(self.pruned_idx)] data = data[self.cluster_agency[cluster]] if cluster is not None else data return data return self.x_out_raw[tuple(self.pruned_idx)] if pruned else self.x_out_raw @property def x_out(self): return self.get_x_out() def get_u(self, cluster=None, pruned=True, raw=False): if cluster is not None: assert pruned if raw: return self.u_raw[tuple(self.pruned_idx)] if pruned else self.u_raw data = self.u_raw[:, self.u_features] if self.u_features is not None else self.u_raw data = data[:, np.newaxis] if len(data.shape) == 1 else data if pruned or cluster is not None: data = data[tuple(self.pruned_idx)] data = data[self.cluster_agency[cluster]] if cluster is not None else data return data @property def u(self): return self.get_u() def get_y(self, cluster=None, pruned=True, raw=False): if cluster is not None: assert pruned if raw: return self.y_raw[self.pruned_idx] if pruned else self.y_raw data = self.y_raw[:, self.y_dim] if self.y_dim is not None else self.y_raw data = data[:, np.newaxis] if len(data.shape) == 1 else data if pruned or cluster is not None: data = data[tuple(self.pruned_idx)] data = data[self.cluster_agency[cluster]] if cluster is not None else data return data @property def y(self): return self.get_y() def get_x_pred(self, pruned=True, raw=False): if raw: return self.x_pred_raw[tuple(self.pruned_idx)] if pruned else self.x_pred_raw data = self.x_pred_raw[:, self.y_dim] if self.y_dim is not None else self.x_pred_raw data = data[:, np.newaxis] if len(data.shape) == 1 else data if pruned: data = data[tuple(self.pruned_idx)] return data @property def x_pred(self): return self.get_x_pred() def get_dt(self, pruned=True): return self.dt_raw[tuple(self.pruned_idx)] if pruned else self.dt_raw @property def dt(self): return self.get_dt() def cluster(self, n_clusters, load_clusters=False, save_dir="", visualize_data=False): self.n_clusters = n_clusters x_pruned = self.x y_pruned = self.y file_name = os.path.join(save_dir, 'gmm.pkl') loaded = False gmm = None if os.path.exists(file_name) and load_clusters: print("Loaded GP clusters from last GP training session. File: {}".format(file_name)) gmm = joblib.load(file_name) if gmm.n_components != n_clusters: print("The loaded GP does not coincide with the number of set clusters: Found {} but requested" "is {}".format(gmm.n_components, n_clusters)) else: loaded = True if not loaded: gmm = GaussianMixture(n_clusters).fit(x_pruned) clusters = gmm.predict_proba(x_pruned) centroids = gmm.means_ if not load_clusters and n_clusters > 1: safe_mknode_recursive(save_dir, 'gmm.pkl', overwrite=True) joblib.dump(gmm, file_name) index_aux = np.arange(0, clusters.shape[0]) cluster_agency = {} # Add to the points corresponding to each cluster the points that correspond to it with top 2 probability, # if this probability is high for cluster in range(n_clusters): top_1 = np.argmax(clusters, 1) clusters_ = copy.deepcopy(clusters) clusters_[index_aux, top_1] *= 0 top_2 = np.argmax(clusters_, 1) idx = np.where(top_1 == cluster)[0] idx = np.append(idx, np.where((top_2 == cluster) * (clusters_[index_aux, top_2] > 0.2))[0]) cluster_agency[cluster] = idx # Only works in len(x_features) >= 3 if visualize_data: x_aux = self.get_x(pruned=False) y_aux = self.get_y(pruned=False) visualize_data_distribution(x_aux, y_aux, cluster_agency, x_pruned, y_pruned) self.cluster_agency = cluster_agency self.centroids = centroids def restore_gp_regressors(pre_trained_models): """ :param pre_trained_models: A dictionary with all the GP models to be restored in 'models'. :return: The GP ensemble restored from the models. :rtype: GPEnsemble """ gp_reg_ensemble = GPEnsemble() # TODO: Deprecate compatibility mode with old models. if all(item in list(pre_trained_models.keys()) for item in ["x_features", "u_features"]): x_features = pre_trained_models["x_features"] u_features = pre_trained_models["u_features"] else: x_features = u_features = None if isinstance(pre_trained_models['models'][0], dict): pre_trained_gp_reg = {} for _, model_dict in enumerate(pre_trained_models['models']): if x_features is not None: gp_reg = npGPRegression(x_features, u_features, model_dict["reg_dim"]) else: gp_reg = npGPRegression(model_dict["x_features"], model_dict["u_features"], model_dict["reg_dim"]) gp_reg.load(model_dict) if model_dict["reg_dim"] not in pre_trained_gp_reg.keys(): pre_trained_gp_reg[model_dict["reg_dim"]] = [gp_reg] else: pre_trained_gp_reg[model_dict["reg_dim"]] += [gp_reg] # Add the GP's in a per-output-dim basis into the Ensemble for key in np.sort(list(pre_trained_gp_reg.keys())): gp_reg_ensemble.add_model(pre_trained_gp_reg[key]) else: raise NotImplementedError("Cannot load this format of GP model.") return gp_reg_ensemble def read_dataset(name, train_split, sim_options): """ Attempts to read a dataset given its name and its metadata. :param name: Name of the dataset :param train_split: Whether to load the training split (True) or the test split (False) :param sim_options: Dictionary of metadata of the dataset to be read. :return: """ data_file = get_data_dir_and_file(name, training_split=train_split, params=sim_options, read_only=True) if data_file is None: raise FileNotFoundError rec_file_dir, rec_file_name = data_file if rec_file_name.startswith('npz'): rec_file_name = rec_file_name[:-4] + '.npz' rec_file = os.path.join(rec_file_dir, rec_file_name) ds = np.load(rec_file) else: rec_file = os.path.join(rec_file_dir, rec_file_name) ds = pd.read_csv(rec_file) return ds def world_to_body_velocity_mapping(state_sequence): """ :param state_sequence: N x 13 state array, where N is the number of states in the sequence. :return: An N x 13 sequence of states, but where the velocities (assumed to be in positions 7, 8, 9) have been rotated from world to body frame. The rotation is made using the quaternion in positions 3, 4, 5, 6. """ p, q, v_w, w = separate_variables(state_sequence) v_b = [] for i in range(len(q)): v_b.append(v_dot_q(v_w[i], quaternion_inverse(q[i]))) v_b = np.stack(v_b) return np.concatenate((p, q, v_b, w), 1) ================================================ FILE: ros_dd_mpc/src/model_fitting/gp_fitting.py ================================================ """ Executable script to train a custom Gaussian Process on recorded flight data. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import os import time import argparse import subprocess import numpy as np from tqdm import tqdm import matplotlib.pyplot as plt from src.utils.utils import safe_mkdir_recursive, load_pickled_models from src.utils.utils import distance_maximizing_points, get_model_dir_and_file from src.utils.utils import sample_random_points from src.model_fitting.gp import CustomKernelFunctions as npKernelFunctions from src.model_fitting.gp import CustomGPRegression as npGPRegression from src.model_fitting.gp import GPEnsemble from src.model_fitting.gp_common import GPDataset, restore_gp_regressors, read_dataset from src.model_fitting.gp_visualization import gp_visualization_experiment from config.configuration_parameters import ModelFitConfig as Conf def plot_gp_regression(x_test, y_test, x_train, y_train, gp_mean, gp_std, gp_regressor, labels, title='', n_samples=3): # Assert the number of provided labels is coherent with the feature dimension [1] of the x vectors if len(x_test.shape) == 1: x_test = np.expand_dims(x_test, 1) n_subplots = 1 else: n_subplots = x_test.shape[1] assert len(labels) == x_test.shape[1] # Generate samples if a gp_regressor is provided if gp_regressor is not None: # Sample from GP & plot samples y_samples = gp_regressor.sample_y(x_test, n_samples) y_samples = np.squeeze(y_samples) else: y_samples = None for i in range(n_subplots): plt.subplot(n_subplots, 1, i + 1) # Sort x axis values x_sort_ind_test = np.argsort(x_test[:, i]) # Plot gp mean line plt.plot(x_test[x_sort_ind_test, i], gp_mean[x_sort_ind_test], 'k', lw=3, zorder=9) # Plot gp std area plt.fill_between(x_test[x_sort_ind_test, i], gp_mean[x_sort_ind_test] - 3 * gp_std[x_sort_ind_test], gp_mean[x_sort_ind_test] + 3 * gp_std[x_sort_ind_test], alpha=0.2, color='k') if y_samples is not None: plt.plot(x_test[x_sort_ind_test, i], y_samples[x_sort_ind_test], '-o', lw=1) plt.xlim(min(x_test[:, i]), max(x_test[:, i])) plt.ylim(min(np.min(y_samples), np.min(gp_mean - 3 * gp_std)), max(np.max(y_samples), np.max(gp_mean + 3 * gp_std))) if x_train is not None and y_train is not None: plt.scatter(x_train[:, i], y_train, c='r', s=50, zorder=10, edgecolors=(0, 0, 0)) if y_test is not None: plt.plot(x_test[x_sort_ind_test, i], y_test[x_sort_ind_test], lw=1, marker='o') if i == 0 and title != '': plt.title(title, fontsize=12) plt.ylabel(labels[i]) plt.tight_layout() def gp_train_and_save(x, y, gp_regressors, save_model, save_file, save_path, y_dims, cluster_n, progress_bar=True): """ Trains and saves the 'm' GP's in the gp_regressors list. Each of these regressors will predict on one of the output variables only. :param x: Feature variables for the regression training. A list of length m where each entry is a Nxn dataset, N is the number of training samples and n is the feature space dimensionality. Each entry of this list will be used to train the respective GP. :param y: Output variables for the regression training. A list of length m where each entry is a N array, N is the number of training samples. Each entry of this list will be used as ground truth output for the respective GP. :param gp_regressors: List of m GPRegression objects (npGPRegression or skGPRegression) :param save_model: Bool. Whether to save the trained models or not. :param save_file: Name of file where to save the model. The 'pkl' extension will be added automatically. :param save_path: Path where to store the trained model pickle file. :param y_dims: List of length m, where each entry represents the state index that corresponds to each GP. :param cluster_n: Number (int) of the cluster currently being trained. :param progress_bar: Bool. Whether to visualize a progress bar or not. :return: the same list as te input gp_regressors variable, but each GP has been trained and saved if requested. """ # If save model, generate saving directory if save_model: safe_mkdir_recursive(save_path) if progress_bar: print("Training {} gp regression models".format(len(y_dims))) prog_range = tqdm(y_dims) if progress_bar else y_dims for y_dim_reg, dim in enumerate(prog_range): # Fit one regressor for each output dimension gp_regressors[y_dim_reg].fit(x[y_dim_reg], y[y_dim_reg]) if save_model: full_path = os.path.join(save_path, save_file + '_' + str(dim) + '_' + str(cluster_n) + '.pkl') gp_regressors[y_dim_reg].save(full_path) return gp_regressors def main(x_features, u_features, reg_y_dim, quad_sim_options, dataset_name, x_cap, hist_bins, hist_thresh, n_train_points=50, n_restarts=10, n_clusters=1, load_clusters=False, model_name="model", dense_gp_name="model", dense_gp_points=100, dense_gp_version="", use_dense=False, visualize_data=False, visualize_model=False): """ Reads the dataset specified and trains a GP model or ensemble on it. The regressed variables is the time-derivative of one of the states. The feature and regressed variables are identified through the index number in the state vector. It is defined as: [0: p_x, 1: p_y, 2:, p_z, 3: q_w, 4: q_x, 5: q_y, 6: q_z, 7: v_x, 8: v_y, 9: v_z, 10: w_x, 11: w_y, 12: w_z] The input vector is also indexed: [0: u_0, 1: u_1, 2: u_2, 3: u_3]. :param x_features: List of n regression feature indices from the quadrotor state indexing. :type x_features: list :param u_features: List of n' regression feature indices from the input state. :type u_features: list :param reg_y_dim: Index of output dimension being regressed as the time-derivative. :type reg_y_dim: float :param dataset_name: Name of the dataset :param quad_sim_options: Dictionary of metadata of the dataset to be read. :param x_cap: cap value (in absolute value) for dataset pruning. Any input feature that exceeds this number in any dimension will be removed. :param hist_bins: Number of bins used for histogram pruning :param hist_thresh: Any bin with less data percentage than this number will be removed. :param n_train_points: Number of training points used for the current GP training. :param dense_gp_points: Number of training points used for the dense GP training. The dense GP will be loaded if possible. :param n_restarts: Number of restarts of nonlinear optimizer. :param n_clusters: Number of clusters used in GP ensemble. If 1, a normal GP is trained. :param load_clusters: True if attempt to load clusters from last GP training. :param model_name: Given name to the currently trained GP. :param dense_gp_name: Given name to the dense GP in case it needs to be used. :param dense_gp_version: Git hash of the folder where to find the dense GP. :param use_dense: Whether to sample the training set from a dense GP or sample from the training data. :param visualize_data: True if display some plots about the data loading, pruning and training process. :param visualize_model: True if display some performance plots about the trained model. """ # #### Prepare saving directory for GP's #### # # Get git commit hash for saving the model git_version = '' try: git_version = subprocess.check_output(['git', 'describe', '--always']).strip().decode("utf-8") except subprocess.CalledProcessError as e: print(e.returncode, e.output) print("The model will be saved using hash: %s" % git_version) gp_name_dict = {"git": git_version, "model_name": model_name, "params": quad_sim_options} save_file_path, save_file_name = get_model_dir_and_file(gp_name_dict) # #### DATASET LOADING #### # if isinstance(dataset_name, str): df_train = read_dataset(dataset_name, True, quad_sim_options) gp_dataset = GPDataset(df_train, x_features, u_features, reg_y_dim, cap=x_cap, n_bins=hist_bins, thresh=hist_thresh, visualize_data=visualize_data) elif isinstance(dataset_name, GPDataset): gp_dataset = dataset_name else: raise TypeError("dataset_name must be a string or a GPDataset instance.") # Make clusters for multi-gp prediction gp_dataset.cluster(n_clusters, load_clusters=load_clusters, save_dir=save_file_path, visualize_data=visualize_data) # #### LOAD DENSE GP IF USING GP ENSEMBLE #### # if use_dense: load_options = {"git": dense_gp_version, "model_name": dense_gp_name, "params": quad_sim_options} loaded_models = load_pickled_models(model_options=load_options) if loaded_models is None: print("Model not found. Training a new dense GP with ") # Train model as accurate as possible. If dense_gp_points is specified, train a gp with that amount of # straining samples and use it to generate a dataset for the GP ensemble. dense_gp = main(x_features, u_features, reg_y_dim, quad_sim_options, dataset_name, x_cap, hist_bins, hist_thresh, n_train_points=dense_gp_points, n_restarts=n_restarts, model_name=dense_gp_name) else: dense_gp = restore_gp_regressors(loaded_models) print("Loaded dense GP model from: %s/%s" % (dense_gp_version, dense_gp_name)) else: dense_gp = None # #### DECLARE SOME PARAMETERS AND VARIABLES #### # # List of trained GP regressors. One for each cluster gp_regressors = [] # Prior parameters sigma_f = 0.5 length_scale = .1 sigma_n = 0.01 gp_params = {"x_features": x_features, "u_features": u_features, "reg_dim": reg_y_dim, "sigma_n": sigma_n, "n_restarts": n_restarts} # Get all cluster centroids for the current output dimension centroids = gp_dataset.centroids print("Training {} cluster model(s)".format(n_clusters)) range_vec = tqdm(range(n_clusters)) for cluster in range_vec: # #### TRAINING POINT SELECTION #### # cluster_mean = centroids[cluster] cluster_x_points = gp_dataset.get_x(cluster=cluster) cluster_y_points = gp_dataset.get_y(cluster=cluster) cluster_u_points = gp_dataset.get_u(cluster=cluster) # Select a base set of training points for the current cluster using PCA that are as separate from each # other as possible selected_points = distance_maximizing_points( cluster_x_points, cluster_mean, n_train_points=n_train_points, dense_gp=dense_gp, plot=False) cluster_y_mean = np.mean(cluster_y_points, 0) # If no dense_gp was provided to the previous function, training_points will be the indices of the training # points to choose from the training set if dense_gp is None: x_train = cluster_x_points[selected_points] y_train = np.squeeze(cluster_y_points[selected_points]) training_points = selected_points else: # Generate a new dataset of synthetic data composed of x and y values x_mock = np.zeros((13, selected_points.shape[1])) if x_features: x_mock[np.array(x_features), :] = selected_points[:len(x_features)] u_mock = np.zeros((4, selected_points.shape[1])) if u_features: u_mock[np.array(u_features), :] = selected_points[len(x_features):] out = dense_gp.predict(x_mock, u_mock) out["pred"] = np.atleast_2d(out["pred"]) y_train = np.squeeze(out["pred"][np.where(dense_gp.dim_idx == reg_y_dim)]) x_train = selected_points.T training_points = [] # Check if we still haven't used the entirety of the available points n_used_points = x_train.shape[0] if n_used_points < n_train_points and n_used_points < cluster_x_points.shape[0]: missing_pts = n_train_points - n_used_points training_points = sample_random_points(cluster_x_points, training_points, missing_pts, dense_gp) if dense_gp is None: # Transform from cluster data index to full dataset index x_train = cluster_x_points[training_points] y_train = np.squeeze(cluster_y_points[training_points]) else: # Generate a new dataset of synthetic data composed of x and y values training_points = training_points.astype(int) x_mock = np.zeros((13, len(training_points))) if x_features: x_mock[np.array(x_features), :] = cluster_x_points[training_points, :len(x_features)].T u_mock = np.zeros((4, len(training_points))) if u_features: u_mock[np.array(u_features), :] = cluster_u_points[len(x_features):] out = dense_gp.predict(x_mock, u_mock) y_additional = np.squeeze(out["pred"][np.where(dense_gp.dim_idx == reg_y_dim)]) y_train = np.append(y_train, y_additional) x_train = np.concatenate((x_train, cluster_x_points[training_points, :len(x_features)]), axis=0) # #### GP TRAINING #### # # Multidimensional input GP regressors l_scale = length_scale * np.ones((x_train.shape[1], 1)) cluster_mean = centroids[cluster] gp_params["mean"] = cluster_mean gp_params["y_mean"] = cluster_y_mean # Train one independent GP for each output dimension exponential_kernel = npKernelFunctions('squared_exponential', params={'l': l_scale, 'sigma_f': sigma_f}) gp_regressors.append(npGPRegression(kernel=exponential_kernel, **gp_params)) gp_regressors[cluster] = gp_train_and_save([x_train], [y_train], [gp_regressors[cluster]], True, save_file_name, save_file_path, [reg_y_dim], cluster, progress_bar=False)[0] if visualize_model: gp_ensemble = GPEnsemble() gp_ensemble.add_model(gp_regressors) x_features = x_features gp_visualization_experiment(quad_sim_options, gp_dataset, x_cap, hist_bins, hist_thresh, x_features, u_features, reg_y_dim, grid_sampling_viz=True, pre_set_gp=gp_ensemble) def gp_evaluate_test_set(gp_regressors, gp_dataset, pruned=False, timed=False, progress_bar=False): """ Runs GP prediction on a specified dataset. :param gp_regressors: GPEnsemble object :type gp_regressors: GPEnsemble :param gp_dataset: Dataset for evaluation :type gp_dataset: GPDataset :param pruned: Whether to use the pruned data or the raw data from the GPDataset. :param timed: whether to return the elapsed time of the data evaluation or not :param progress_bar: If True, a progress bar will be shown when evaluating the test data. :return: Given n number of samples of dimension d, return: - n x d vector of mean predictions, - n x d vector of std predictions, - n x d vector of ground truth values to compare with the predictions """ x_test = gp_dataset.get_x(pruned=pruned, raw=True) u_test = gp_dataset.get_u(pruned=pruned, raw=True) dt_test = gp_dataset.get_dt(pruned=pruned) tic = time.time() out = gp_regressors.predict(x_test.T, u_test.T, return_std=True, progress_bar=progress_bar) mean_post = out["pred"] std_post = out["cov_or_std"] elapsed = time.time() - tic mean_post *= dt_test std_post *= dt_test mean_post = mean_post.T std_post = std_post.T if not timed: return mean_post, std_post else: return mean_post, std_post, elapsed if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--n_points", type=int, default="20", help="Number of training points used to fit the current GP model.") parser.add_argument("--model_name", type=str, default="", help="Name assigned to the trained model.") parser.add_argument('--x', nargs='+', type=int, default=[7], help='Regression X variables. Must be a list of integers between 0 and 12. Velocities xyz ' 'correspond to indices 7, 8, 9.') parser.add_argument('--u', action="store_true", help='Use the control as input to the model.') parser.add_argument("--y", type=int, default=7, help="Regression Y variable. Must be an integer between 0 and 12. Velocities xyz correspond to" "indices 7, 8, 9.") input_arguments = parser.parse_args() # Use vx, vy, vz as input features x_feats = input_arguments.x if input_arguments.u: u_feats = [0, 1, 2, 3] else: u_feats = [] # Regression dimension y_regressed_dim = input_arguments.y n_train_pts = input_arguments.n_points gp_name = input_arguments.model_name ds_name = Conf.ds_name simulation_options = Conf.ds_metadata histogram_pruning_bins = Conf.histogram_bins histogram_pruning_threshold = Conf.histogram_threshold x_value_cap = Conf.velocity_cap gp_dense_name = Conf.dense_model_name gp_id_custom = Conf.dense_model_version dense_n_points = Conf.dense_training_points with_dense = Conf.use_dense_model main(x_feats, u_feats, y_regressed_dim, simulation_options, ds_name, x_value_cap, histogram_pruning_bins, histogram_pruning_threshold, model_name=gp_name, n_train_points=n_train_pts, n_clusters=Conf.clusters, load_clusters=Conf.load_clusters, use_dense=with_dense, dense_gp_points=dense_n_points, dense_gp_name=gp_dense_name, dense_gp_version=gp_id_custom, visualize_data=Conf.visualize_data, visualize_model=Conf.visualize_training_result) ================================================ FILE: ros_dd_mpc/src/model_fitting/gp_visualization.py ================================================ """ Executable script for visual evaluation of the trained GP quality. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import argparse import numpy as np import matplotlib.pyplot as plt from src.model_fitting.gp_common import GPDataset, restore_gp_regressors, read_dataset from config.configuration_parameters import ModelFitConfig as Conf from src.utils.utils import load_pickled_models from src.utils.visualization import visualize_gp_inference def gp_visualization_experiment(quad_sim_options, dataset_name, x_cap, hist_bins, hist_thresh, x_vis_feats, u_vis_feats, y_vis_feats, grid_sampling_viz=False, load_model_version="", load_model_name="", pre_set_gp=None): # #### GP ENSEMBLE LOADING #### # if pre_set_gp is None: load_options = {"git": load_model_version, "model_name": load_model_name, "params": quad_sim_options} loaded_models = load_pickled_models(model_options=load_options) if loaded_models is None: raise FileNotFoundError("Model not found") gp_ensemble = restore_gp_regressors(loaded_models) else: gp_ensemble = pre_set_gp # #### DATASET LOADING #### # # Pre-set labels of the data: labels_x = [ r'$p_x\:\left[m\right]$', r'$p_y\:\left[m\right]$', r'$p_z\:\left[m\right]$', r'$q_w\:\left[rad\right]$', r'$q_x\:\left[rad\right]$', r'$q_y\:\left[rad\right]$', r'$q_z\:\left[rad\right]$', r'$v_x\:\left[\frac{m}{s}\right]$', r'$v_y\:\left[\frac{m}{s}\right]$', r'$v_z\:\left[\frac{m}{s}\right]$', r'$w_x\:\left[\frac{rad}{s}\right]$', r'$w_y\:\left[\frac{rad}{s}\right]$', r'$w_z\:\left[\frac{rad}{s}\right]$' ] labels_u = [r'$u_1$', r'$u_2$', r'$u_3$', r'$u_4$'] labels = [labels_x[feat] for feat in x_vis_feats] labels_ = [labels_u[feat] for feat in u_vis_feats] labels = labels + labels_ if isinstance(dataset_name, str): test_ds = read_dataset(dataset_name, True, quad_sim_options) test_gp_ds = GPDataset(test_ds, x_features=x_vis_feats, u_features=u_vis_feats, y_dim=y_vis_feats, cap=x_cap, n_bins=hist_bins, thresh=hist_thresh, visualize_data=False) else: test_gp_ds = dataset_name x_test = test_gp_ds.get_x(pruned=True, raw=True) u_test = test_gp_ds.get_u(pruned=True, raw=True) y_test = test_gp_ds.get_y(pruned=True, raw=False) dt_test = test_gp_ds.get_dt(pruned=True) x_pred = test_gp_ds.get_x_pred(pruned=True, raw=False) if isinstance(y_vis_feats, list): y_dims = [np.where(gp_ensemble.dim_idx == y_dim)[0][0] for y_dim in y_vis_feats] else: y_dims = np.where(gp_ensemble.dim_idx == y_vis_feats)[0] # #### VISUALIZE GRID SAMPLING RESULTS IN TRAINING SET RANGE #### # if grid_sampling_viz: visualize_gp_inference(x_test, u_test, y_test, gp_ensemble, x_vis_feats, y_dims, labels) # #### EVALUATE GP ON TEST SET #### # print("Test set prediction...") outs = gp_ensemble.predict(x_test.T, u_test.T, return_std=True, progress_bar=True) mean_estimate = np.atleast_2d(np.atleast_2d(outs["pred"])[y_dims]) std_estimate = np.atleast_2d(np.atleast_2d(outs["cov_or_std"])[y_dims]) mean_estimate = mean_estimate.T * dt_test[:, np.newaxis] std_estimate = std_estimate.T * dt_test[:, np.newaxis] # Undo dt normalization y_test *= dt_test[:, np.newaxis] # Error of nominal model nominal_diff = y_test # GP regresses model error, correct the predictions of the nominal model augmented_diff = nominal_diff - mean_estimate mean_estimate += x_pred nominal_rmse = np.mean(np.abs(nominal_diff), 0) augmented_rmse = np.mean(np.abs(augmented_diff), 0) labels = [r'$v_x$ error', r'$v_y$ error', r'$v_z$ error'] t_vec = np.cumsum(dt_test) plt.figure() # mng = plt.get_current_fig_manager() # mng.resize(*mng.window.maxsize()) for i in range(std_estimate.shape[1]): plt.subplot(std_estimate.shape[1], 1, i+1) plt.plot(t_vec, np.zeros(augmented_diff[:, i].shape), 'k') plt.plot(t_vec, augmented_diff[:, i], 'b', label='augmented_err') plt.plot(t_vec, augmented_diff[:, i] - 3 * std_estimate[:, i], ':c') plt.plot(t_vec, augmented_diff[:, i] + 3 * std_estimate[:, i], ':c', label='3 std') if nominal_diff is not None: plt.plot(t_vec, nominal_diff[:, i], 'r', label='nominal_err') plt.title('Mean dt: %.2f s. Nominal RMSE: %.5f [m/s]. Augmented rmse: %.5f [m/s]' % ( float(np.mean(dt_test)), nominal_rmse[i], augmented_rmse[i])) else: plt.title('Mean dt: %.2f s. Augmented RMSE: %.5f [m/s]' % ( float(np.mean(dt_test)), float(augmented_rmse[i]))) plt.ylabel(labels[i]) plt.legend() if i == std_estimate.shape[1] - 1: plt.xlabel('time (s)') plt.show() if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--model_version", type=str, default="", help="Version to load for the regression models. By default it is an 8 digit git hash.") parser.add_argument("--model_name", type=str, default="", help="Name of the regression model within the specified folder.") input_arguments = parser.parse_args() ds_name = Conf.ds_name simulation_options = Conf.ds_metadata histogram_pruning_bins = Conf.histogram_bins histogram_pruning_threshold = Conf.histogram_threshold x_value_cap = Conf.velocity_cap x_viz = Conf.x_viz u_viz = Conf.u_viz y_viz = Conf.y_viz gp_load_model_name = input_arguments.model_name gp_load_model_version = input_arguments.model_version gp_visualization_experiment(simulation_options, ds_name, x_value_cap, histogram_pruning_bins, histogram_pruning_threshold, x_vis_feats=x_viz, u_vis_feats=u_viz, y_vis_feats=y_viz, grid_sampling_viz=True, load_model_version=gp_load_model_version, load_model_name=gp_load_model_name) ================================================ FILE: ros_dd_mpc/src/model_fitting/mlp_common.py ================================================ import numpy as np from torch.utils.data import Dataset import ml_casadi.torch as mc from config.configuration_parameters import GroundEffectMapConfig from src.utils.ground_map import GroundMapWithBox class GPToMLPDataset(Dataset): def __init__(self, gp_dataset, ground_effect=False): super().__init__() self.x = gp_dataset.x self.y = gp_dataset.y if ground_effect: self.x_raw = gp_dataset.x_raw[tuple(gp_dataset.pruned_idx)] assert self.x.shape[0] == self.x_raw.shape[0] map_conf = GroundEffectMapConfig ground_map = GroundMapWithBox(np.array(map_conf.box_min), np.array(map_conf.box_max), map_conf.box_height, horizon=map_conf.horizon, resolution=map_conf.resolution) self._map_res = map_conf.resolution self._static_ground_map, self._org_to_map_org = ground_map.at(np.array(map_conf.origin)) else: self._static_ground_map = None def stats(self): if self._static_ground_map is None: return self.x.mean(axis=0), self.x.std(axis=0), self.y.mean(axis=0), self.y.std(axis=0) else: x_mean = np.hstack([self.x.mean(axis=0), np.zeros(9+4)]) x_std = np.hstack([self.x.std(axis=0), np.ones(9+4)]) return x_mean, x_std, self.y.mean(axis=0), self.y.std(axis=0) def __len__(self): return len(self.x) def __getitem__(self, item): if self._static_ground_map is None: return self.x[item], self.y[item] else: x, y, z = self.x_raw[item][:3] orientation = self.x_raw[item][3:7] x_idxs = np.floor((x - self._org_to_map_org[0]) / self._map_res).astype(int) - 1 y_idxs = np.floor((y - self._org_to_map_org[1]) / self._map_res).astype(int) - 1 ground_patch = self._static_ground_map[x_idxs:x_idxs+3, y_idxs:y_idxs+3] relative_ground_patch = 4 * (np.clip(z - ground_patch, 0, 0.5) - 0.25) flatten_relative_ground_patch = relative_ground_patch.flatten(order='F') ground_effect_in = np.hstack([flatten_relative_ground_patch, orientation*0]) return np.hstack([self.x[item], ground_effect_in]), self.y[item] class NormalizedMLP(mc.TorchMLCasadiModule): def __init__(self, model, x_mean, x_std, y_mean, y_std): super().__init__() self.model = model self.input_size = self.model.input_size self.output_size = self.model.output_size self.register_buffer('x_mean', x_mean) self.register_buffer('x_std', x_std) self.register_buffer('y_mean', y_mean) self.register_buffer('y_std', y_std) def forward(self, x): return (self.model((x - self.x_mean) / self.x_std) * self.y_std) + self.y_mean def cs_forward(self, x): return (self.model((x - self.x_mean.cpu().numpy()) / self.x_std.cpu().numpy()) * self.y_std.cpu().numpy()) + self.y_mean.cpu().numpy() class QuadResidualModel(mc.TorchMLCasadiModule): def __init__(self, hidden_size, hidden_layers): super().__init__() self._input_size = 10 self._output_size = 6 self.force_model = mc.nn.MultiLayerPerceptron(7, hidden_size, 3, hidden_layers, 'Tanh') self.torque_model = mc.nn.MultiLayerPerceptron(10, hidden_size, 3, hidden_layers, 'Tanh') def forward(self, x): v = x[:, :3] w = x[:, 3:6] u = x[:, 6:] force_in = mc.hcat([v, u]) torque_in = mc.hcat([v, w, u]) force_out = self.force_model(force_in) torque_out = self.torque_model(torque_in) out = mc.hcat([force_out, torque_out]) return out def cs_forward(self, x): v = x[:3] w = x[3:6] u = x[6:] force_in = mc.vcat([v, u]) torque_in = mc.vcat([v, w, u]) force_out = self.force_model(force_in) torque_out = self.torque_model(torque_in) out = mc.vcat([force_out, torque_out]) return out ================================================ FILE: ros_dd_mpc/src/model_fitting/mlp_fitting.py ================================================ """ Executable script to train a custom Gaussian Process on recorded flight data. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import os import time import argparse import subprocess import numpy as np import torch from torch.utils.data import DataLoader import ml_casadi.torch as mc from tqdm import tqdm from src.model_fitting.mlp_common import GPToMLPDataset, NormalizedMLP from src.utils.utils import safe_mkdir_recursive, load_pickled_models from src.utils.utils import get_model_dir_and_file from src.model_fitting.gp_common import GPDataset, read_dataset from config.configuration_parameters import ModelFitConfig as Conf def main(x_features, u_features, reg_y_dims, model_ground_effect, quad_sim_options, dataset_name, x_cap, hist_bins, hist_thresh, model_name="model", epochs=100, batch_size=64, hidden_layers=1, hidden_size=32, plot=False): """ Reads the dataset specified and trains a GP model or ensemble on it. The regressed variables is the time-derivative of one of the states. The feature and regressed variables are identified through the index number in the state vector. It is defined as: [0: p_x, 1: p_y, 2:, p_z, 3: q_w, 4: q_x, 5: q_y, 6: q_z, 7: v_x, 8: v_y, 9: v_z, 10: w_x, 11: w_y, 12: w_z] The input vector is also indexed: [0: u_0, 1: u_1, 2: u_2, 3: u_3]. :param x_features: List of n regression feature indices from the quadrotor state indexing. :type x_features: list :param u_features: List of n' regression feature indices from the input state. :type u_features: list :param reg_y_dims: Indices of output dimension being regressed as the time-derivative. :type reg_y_dims: List :param dataset_name: Name of the dataset :param quad_sim_options: Dictionary of metadata of the dataset to be read. :param x_cap: cap value (in absolute value) for dataset pruning. Any input feature that exceeds this number in any dimension will be removed. :param hist_bins: Number of bins used for histogram pruning :param hist_thresh: Any bin with less data percentage than this number will be removed. :param model_name: Given name to the currently trained GP. """ # Get git commit hash for saving the model git_version = '' try: git_version = subprocess.check_output(['git', 'describe', '--always']).strip().decode("utf-8") except subprocess.CalledProcessError as e: print(e.returncode, e.output) print("The model will be saved using hash: %s" % git_version) gp_name_dict = {"git": git_version, "model_name": model_name, "params": quad_sim_options} save_file_path, save_file_name = get_model_dir_and_file(gp_name_dict) safe_mkdir_recursive(save_file_path) # #### DATASET LOADING #### # if isinstance(dataset_name, str): df_train = read_dataset(dataset_name, True, quad_sim_options) gp_dataset_train = GPDataset(df_train, x_features, u_features, reg_y_dims, cap=x_cap, n_bins=hist_bins, thresh=hist_thresh) gp_dataset_val = None try: df_val = read_dataset(dataset_name, False, quad_sim_options) gp_dataset_val = GPDataset(df_val, x_features, u_features, reg_y_dims, cap=x_cap, n_bins=hist_bins, thresh=hist_thresh) except: print('Could not find test dataset.') else: raise TypeError("dataset_name must be a string.") dataset_train = GPToMLPDataset(gp_dataset_train, ground_effect=model_ground_effect) x_mean, x_std, y_mean, y_std = dataset_train.stats() data_loader = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=0) input_dims = len(x_features) + len(u_features) + (9 + 4 if model_ground_effect else 0) mlp_model = mc.nn.MultiLayerPerceptron(input_dims, hidden_size, len(reg_y_dims), hidden_layers, 'Tanh') model = NormalizedMLP(mlp_model, torch.tensor(x_mean).float(), torch.tensor(x_std).float(), torch.tensor(y_mean).float(), torch.tensor(y_std).float()) print(f"Train Dataset has {len(dataset_train)} points.") if gp_dataset_val: dataset_val = GPToMLPDataset(gp_dataset_val, ground_effect=model_ground_effect) data_loader_val = DataLoader(dataset_val, batch_size=10*batch_size, shuffle=False, num_workers=0) print(f"Val Dataset has {len(dataset_val)} points.") print(f"Model has {sum(p.numel() for p in model.parameters() if p.requires_grad)} parameters.") if torch.cuda.is_available(): model = model.to('cuda:0') optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) loss_infos = [] p_bar = tqdm(range(epochs)) for i in p_bar: model.train() losses = [] for x, y in data_loader: if torch.cuda.is_available(): x = x.to('cuda:0') y = y.to('cuda:0') x = x.float() y = y.float() optimizer.zero_grad() y_pred = model(x) loss = torch.square(y - y_pred).mean() loss.backward() optimizer.step() losses.append(loss.item()) losses_val = [] if gp_dataset_val: with torch.no_grad(): for x, y in data_loader_val: if torch.cuda.is_available(): x = x.to('cuda:0') y = y.to('cuda:0') x = x.float() y = y.float() y_pred = model(x) loss = torch.square(y - y_pred) losses_val.append(loss.numpy()) train_loss_info = np.mean(losses) loss_info = np.mean(np.vstack(losses_val)) if losses_val else np.mean(losses) loss_infos.append(loss_info) p_bar.set_description(f'Train Loss: {train_loss_info}, Val Loss {loss_info:.6f}') p_bar.refresh() save_dict = { 'state_dict': model.state_dict(), 'input_size': input_dims, 'hidden_size': hidden_size, 'output_size': len(reg_y_dims), 'hidden_layers': hidden_layers } torch.save(save_dict, os.path.join(save_file_path, f'{save_file_name}.pt')) if plot: import matplotlib.pyplot as plt plt.plot(loss_infos) plt.show() if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--epochs", type=int, default="1000", help="Number of epochs to train the model.") parser.add_argument("--batch_size", type=int, default="64", help="Batch Size.") parser.add_argument("--hidden_layers", type=int, default="1", help="Number of hidden layers.") parser.add_argument("--hidden_size", type=int, default="32", help="Size of hidden layers.") parser.add_argument("--model_name", type=str, default="", help="Name assigned to the trained model.") parser.add_argument('--x', nargs='+', type=int, default=[7], help='Regression X variables. Must be a list of integers between 0 and 12. Velocities xyz ' 'correspond to indices 7, 8, 9.') parser.add_argument('--u', action="store_true", help='Use the control as input to the model.') parser.add_argument("--y", nargs='+', type=int, default=[7], help="Regression Y variable. Must be an integer between 0 and 12. Velocities xyz correspond to" "indices 7, 8, 9.") parser.add_argument('--ge', action="store_true", help='Use the ground map as input to the model.') parser.add_argument("--plot", dest="plot", action="store_true", help="Plot the loss after training.") parser.set_defaults(plot=False) input_arguments = parser.parse_args() # Use vx, vy, vz as input features x_feats = input_arguments.x if input_arguments.u: u_feats = [0, 1, 2, 3] else: u_feats = [] model_ground_effect = input_arguments.ge # Regression dimension y_regressed_dims = input_arguments.y model_name = input_arguments.model_name epochs = input_arguments.epochs batch_size = input_arguments.batch_size hidden_layers = input_arguments.hidden_layers hidden_size = input_arguments.hidden_size plot = input_arguments.plot ds_name = Conf.ds_name simulation_options = Conf.ds_metadata histogram_pruning_bins = Conf.histogram_bins histogram_pruning_threshold = Conf.histogram_threshold x_value_cap = Conf.velocity_cap main(x_feats, u_feats, y_regressed_dims, model_ground_effect, simulation_options, ds_name, x_value_cap, histogram_pruning_bins, histogram_pruning_threshold, model_name=model_name, epochs=epochs, batch_size=batch_size, hidden_layers=hidden_layers, hidden_size=hidden_size, plot=plot) ================================================ FILE: ros_dd_mpc/src/model_fitting/mlp_quad_res_fitting.py ================================================ """ Executable script to train a custom Gaussian Process on recorded flight data. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import os import time import argparse import subprocess import numpy as np import torch from torch.utils.data import DataLoader import ml_casadi.torch as mc from tqdm import tqdm from src.model_fitting.mlp_common import GPToMLPDataset, NormalizedMLP, QuadResidualModel from src.utils.utils import safe_mkdir_recursive, load_pickled_models from src.utils.utils import get_model_dir_and_file from src.model_fitting.gp_common import GPDataset, read_dataset from config.configuration_parameters import ModelFitConfig as Conf def main(x_features, u_features, reg_y_dims, quad_sim_options, dataset_name, x_cap, hist_bins, hist_thresh, model_name="model", epochs=100, batch_size=64, hidden_layers=1, hidden_size=32, plot=False): """ Reads the dataset specified and trains a GP model or ensemble on it. The regressed variables is the time-derivative of one of the states. The feature and regressed variables are identified through the index number in the state vector. It is defined as: [0: p_x, 1: p_y, 2:, p_z, 3: q_w, 4: q_x, 5: q_y, 6: q_z, 7: v_x, 8: v_y, 9: v_z, 10: w_x, 11: w_y, 12: w_z] The input vector is also indexed: [0: u_0, 1: u_1, 2: u_2, 3: u_3]. :param x_features: List of n regression feature indices from the quadrotor state indexing. :type x_features: list :param u_features: List of n' regression feature indices from the input state. :type u_features: list :param reg_y_dims: Indices of output dimension being regressed as the time-derivative. :type reg_y_dims: List :param dataset_name: Name of the dataset :param quad_sim_options: Dictionary of metadata of the dataset to be read. :param x_cap: cap value (in absolute value) for dataset pruning. Any input feature that exceeds this number in any dimension will be removed. :param hist_bins: Number of bins used for histogram pruning :param hist_thresh: Any bin with less data percentage than this number will be removed. :param model_name: Given name to the currently trained GP. """ # Get git commit hash for saving the model git_version = '' try: git_version = subprocess.check_output(['git', 'describe', '--always']).strip().decode("utf-8") except subprocess.CalledProcessError as e: print(e.returncode, e.output) print("The model will be saved using hash: %s" % git_version) gp_name_dict = {"git": git_version, "model_name": model_name, "params": quad_sim_options} save_file_path, save_file_name = get_model_dir_and_file(gp_name_dict) safe_mkdir_recursive(save_file_path) # #### DATASET LOADING #### # if isinstance(dataset_name, str): df_train = read_dataset(dataset_name, True, quad_sim_options) gp_dataset_train = GPDataset(df_train, x_features, u_features, reg_y_dims, cap=x_cap, n_bins=hist_bins, thresh=hist_thresh) gp_dataset_val = None try: df_val = read_dataset(dataset_name, False, quad_sim_options) gp_dataset_val = GPDataset(df_val, x_features, u_features, reg_y_dims, cap=x_cap, n_bins=hist_bins, thresh=hist_thresh) except: print('Could not find test dataset.') else: raise TypeError("dataset_name must be a string.") dataset_train = GPToMLPDataset(gp_dataset_train) x_mean, x_std, y_mean, y_std = dataset_train.stats() data_loader = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=0) input_dims = len(x_features) + len(u_features) mlp_model = QuadResidualModel(hidden_size, hidden_layers) model = NormalizedMLP(mlp_model, torch.tensor(x_mean).float(), torch.tensor(x_std).float(), torch.tensor(y_mean).float(), torch.tensor(y_std).float()) print(f"Train Dataset has {len(dataset_train)} points.") if gp_dataset_val: dataset_val = GPToMLPDataset(gp_dataset_val) data_loader_val = DataLoader(dataset_val, batch_size=10*batch_size, shuffle=False, num_workers=0) print(f"Val Dataset has {len(dataset_val)} points.") print(f"Model has {sum(p.numel() for p in model.parameters() if p.requires_grad)} parameters.") if torch.cuda.is_available(): print('Pushing model to GPU.') model = model.to('cuda:0') optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) loss_infos = [] p_bar = tqdm(range(epochs)) for i in p_bar: model.train() losses = [] for x, y in data_loader: if torch.cuda.is_available(): x = x.to('cuda:0') y = y.to('cuda:0') x = x.float() y = y.float() optimizer.zero_grad() y_pred = model(x) loss = torch.square(y - y_pred).mean() loss.backward() optimizer.step() losses.append(loss.item()) losses_val = [] if gp_dataset_val: with torch.no_grad(): for x, y in data_loader_val: if torch.cuda.is_available(): x = x.to('cuda:0') y = y.to('cuda:0') x = x.float() y = y.float() y_pred = model(x) loss = torch.square(y - y_pred) losses_val.append(loss.cpu().numpy()) train_loss_info = np.mean(losses) loss_info = np.mean(np.vstack(losses_val)) if losses_val else np.mean(losses) loss_infos.append(loss_info) p_bar.set_description(f'Train Loss: {train_loss_info}, Val Loss {loss_info:.6f}') p_bar.refresh() save_dict = { 'state_dict': model.state_dict(), 'input_size': input_dims, 'hidden_size': hidden_size, 'output_size': len(reg_y_dims), 'hidden_layers': hidden_layers } torch.save(save_dict, os.path.join(save_file_path, f'{save_file_name}_{i}.pt')) print(f'{i}: Val Loss {loss_info:.6f}') if plot: import matplotlib.pyplot as plt plt.plot(loss_infos) plt.show() if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--epochs", type=int, default="1000", help="Number of epochs to train the model.") parser.add_argument("--batch_size", type=int, default="64", help="Batch Size.") parser.add_argument("--hidden_layers", type=int, default="1", help="Number of hidden layers.") parser.add_argument("--hidden_size", type=int, default="32", help="Size of hidden layers.") parser.add_argument("--model_name", type=str, default="", help="Name assigned to the trained model.") parser.add_argument("--plot", dest="plot", action="store_true", help="Plot the loss after training.") parser.set_defaults(plot=False) input_arguments = parser.parse_args() # Use vx, vy, vz as input features x_feats = [7, 8, 9, 10, 11, 12] u_feats = [0, 1, 2, 3] # Regression dimension y_regressed_dims = [7, 8, 9, 10, 11, 12] model_name = input_arguments.model_name epochs = input_arguments.epochs batch_size = input_arguments.batch_size hidden_layers = input_arguments.hidden_layers hidden_size = input_arguments.hidden_size plot = input_arguments.plot ds_name = Conf.ds_name simulation_options = Conf.ds_metadata histogram_pruning_bins = Conf.histogram_bins histogram_pruning_threshold = Conf.histogram_threshold x_value_cap = Conf.velocity_cap main(x_feats, u_feats, y_regressed_dims, simulation_options, ds_name, x_value_cap, histogram_pruning_bins, histogram_pruning_threshold, model_name=model_name, epochs=epochs, batch_size=batch_size, hidden_layers=hidden_layers, hidden_size=hidden_size, plot=plot) ================================================ FILE: ros_dd_mpc/src/model_fitting/process_neurobem_dataset.py ================================================ import argparse import os import random import numpy as np import pandas as pd from scipy.signal import savgol_filter from tqdm import tqdm from config.configuration_parameters import DirectoryConfig from config.configuration_parameters import ModelFitConfig as Conf from src.experiments.point_tracking_and_record import make_record_dict, jsonify from src.quad_mpc.create_ros_dd_mpc import custom_quad_param_loader from src.quad_mpc.quad_3d_mpc import Quad3DMPC from src.utils.utils import safe_mkdir_recursive, v_dot_q val_files = ['merged_2021-02-03-16-58-13_seg_2.csv'] def main(quad): full_path = os.path.join(DirectoryConfig.DATA_DIR, 'neurobem_processed_data_cmd') files = os.listdir(full_path) random.shuffle(files) rec_dicts = [] for file in tqdm(files): if file in val_files: print('Skipping Validation File') continue try: rec_dict = make_record_dict(state_dim=13) process_file(os.path.join(full_path, file), quad, rec_dict) rec_dicts.append(rec_dict) except Exception as e: print(e) rec_dict = {} rec_dict['state_in'] = np.concatenate([data_dict['state_in'] for data_dict in rec_dicts], axis=0) rec_dict['input_in'] = np.concatenate([data_dict['input_in'] for data_dict in rec_dicts], axis=0) rec_dict['state_out'] = np.concatenate([data_dict['state_out'] for data_dict in rec_dicts], axis=0) rec_dict['state_ref'] = np.concatenate([data_dict['state_ref'] for data_dict in rec_dicts], axis=0) rec_dict['timestamp'] = np.concatenate([data_dict['timestamp'] for data_dict in rec_dicts], axis=0) rec_dict['dt'] = np.concatenate([data_dict['dt'] for data_dict in rec_dicts], axis=0) rec_dict['state_pred'] = np.concatenate([data_dict['state_pred'] for data_dict in rec_dicts], axis=0) rec_dict['error'] = np.concatenate([data_dict['error'] for data_dict in rec_dicts], axis=0) del rec_dicts # Save datasets save_file_folder = os.path.join(DirectoryConfig.DATA_DIR, 'neurobem_dataset', 'train') safe_mkdir_recursive(save_file_folder) save_file = os.path.join(save_file_folder, f'npz_dataset_001.npz') np.savez(save_file, **rec_dict) # # Validation # rec_dicts = [] # for file in tqdm(files[-20:]): # try: # rec_dict = make_record_dict(state_dim=13) # process_file(os.path.join(full_path, file), quad, rec_dict) # rec_dicts.append(rec_dict) # except Exception as e: # print(e) # # rec_dict = {} # rec_dict['state_in'] = np.concatenate([data_dict['state_in'] for data_dict in rec_dicts], axis=0) # rec_dict['input_in'] = np.concatenate([data_dict['input_in'] for data_dict in rec_dicts], axis=0) # rec_dict['state_out'] = np.concatenate([data_dict['state_out'] for data_dict in rec_dicts], axis=0) # rec_dict['state_ref'] = np.concatenate([data_dict['state_ref'] for data_dict in rec_dicts], axis=0) # rec_dict['timestamp'] = np.concatenate([data_dict['timestamp'] for data_dict in rec_dicts], axis=0) # rec_dict['dt'] = np.concatenate([data_dict['dt'] for data_dict in rec_dicts], axis=0) # rec_dict['state_pred'] = np.concatenate([data_dict['state_pred'] for data_dict in rec_dicts], axis=0) # rec_dict['error'] = np.concatenate([data_dict['error'] for data_dict in rec_dicts], axis=0) # # del rec_dicts # # # Save datasets # save_file_folder = os.path.join(DirectoryConfig.DATA_DIR, 'neurobem_dataset', 'test') # safe_mkdir_recursive(save_file_folder) # save_file = os.path.join(save_file_folder, f'npz_dataset_001.npz') # np.savez(save_file, **rec_dict) def quaternion_to_euler(q): yaw = np.arctan2(2 * (q[:, 0] * q[:, 3] - q[:, 1] * q[:, 2]), 1 - 2 * (q[:, 2] ** 2 + q[:, 3] ** 2)) pitch = np.arcsin(2 * (q[:, 0] * q[:, 2] + q[:, 3] * q[:, 1])) roll = np.arctan2(2 * (q[:, 0] * q[:, 1] - q[:, 2] * q[:, 3]), 1 - 2 * (q[:, 1] ** 2 + q[:, 2] ** 2)) return np.stack([yaw, pitch, roll], axis=-1) def val(quad): full_path = os.path.join(DirectoryConfig.DATA_DIR, 'neurobem_processed_data_cmd') #val_files = ['merged_2021-02-23-22-54-17_seg_1.csv', 'merged_2021-02-23-17-35-26_seg_1.csv'] # 'merged_2021-02-18-18-09-47_seg_1.csv', rec_dicts = [] for file in tqdm(val_files): try: rec_dict = make_record_dict(state_dim=13) process_file(os.path.join(full_path, file), quad, rec_dict) rec_dicts.append(rec_dict) except Exception as e: print(e) rec_dict = {} rec_dict['state_in'] = np.concatenate([data_dict['state_in'] for data_dict in rec_dicts], axis=0) rec_dict['input_in'] = np.concatenate([data_dict['input_in'] for data_dict in rec_dicts], axis=0) rec_dict['state_out'] = np.concatenate([data_dict['state_out'] for data_dict in rec_dicts], axis=0) rec_dict['state_ref'] = np.concatenate([data_dict['state_ref'] for data_dict in rec_dicts], axis=0) rec_dict['timestamp'] = np.concatenate([data_dict['timestamp'] for data_dict in rec_dicts], axis=0) rec_dict['dt'] = np.concatenate([data_dict['dt'] for data_dict in rec_dicts], axis=0) rec_dict['state_pred'] = np.concatenate([data_dict['state_pred'] for data_dict in rec_dicts], axis=0) rec_dict['error'] = np.concatenate([data_dict['error'] for data_dict in rec_dicts], axis=0) del rec_dicts # Save datasets save_file_folder = os.path.join(DirectoryConfig.DATA_DIR, 'neurobem_dataset', 'test') safe_mkdir_recursive(save_file_folder) save_file = os.path.join(save_file_folder, f'npz_dataset_001.npz') np.savez(save_file, **rec_dict) def process_file(file_path, quad, rec_dict): data = pd.read_csv(file_path, encoding='latin-1') for x_0, x_f, u, dt in consecutive_data_points(data): resimulate(x_0, x_f, u, dt, quad, rec_dict) def consecutive_data_points(data): for i in range(25, len(data)-1, 1): data_0 = data.iloc[i] data_1 = data.iloc[i+1] data_c = data.iloc[i] # Communication Delay t0 = data_0['t'] t1 = data_1['t'] dt = t1 - t0 x_0 = np.hstack([ data_0['pos x'], data_0['pos y'], data_0['pos z'], data_0['quat w'], data_0['quat x'], data_0['quat y'], data_0['quat z'], data_0['vel x'], data_0['vel y'], data_0['vel z'], data_0['ang vel x'], data_0['ang vel y'], data_0['ang vel z'], ]) x_1 = np.hstack([ data_1['pos x'], data_1['pos y'], data_1['pos z'], data_1['quat w'], data_1['quat x'], data_1['quat y'], data_1['quat z'], data_1['vel x'], data_1['vel y'], data_1['vel z'], data_1['ang vel x'], data_1['ang vel y'], data_1['ang vel z'], ]) # Velocity to world coordinate frame x_0[7:10] = v_dot_q(x_0[7:10], x_0[3:7]) x_1[7:10] = v_dot_q(x_1[7:10], x_1[3:7]) # Solve single rotor thrusts Jm1 = 1 / quad.J a1 = Jm1[0] * quad.y_f b1 = data_0['ang acc x'] - Jm1[0] * (quad.J[1] - quad.J[2]) * x_0[11] * x_0[12] a2 = -Jm1[1] * quad.x_f b2 = data_0['ang acc y'] - Jm1[1] * (quad.J[2] - quad.J[0]) * x_0[12] * x_0[10] a3 = Jm1[2] * quad.z_l_tau b3 = data_0['ang acc z'] - Jm1[2] * (quad.J[0] - quad.J[1]) * x_0[10] * x_0[11] a4 = np.ones(4,) b4 = data_c['cmd_thrust'] * quad_mpc.quad.mass u = np.linalg.solve([a1, a2, a3, a4], [b1, b2, b3, b4]) / quad.max_thrust # ypr_0 = quaternion_to_euler(x_0[np.newaxis, 3:7])[0][::-1] # x_0[12] = 1.0 * x_0[10] + np.sin(ypr_0[0]) * np.tan(ypr_0[1]) * x_0[11] + np.cos(ypr_0[0]) * np.tan(ypr_0[1]) * x_0[12] # x_0[11] = np.cos(ypr_0[0]) * x_0[11] - np.sin(ypr_0[0]) * x_0[12] # x_0[10] = np.sin(ypr_0[0]) / np.cos(ypr_0[1]) * x_0[11] + np.cos(ypr_0[0]) / np.cos(ypr_0[1]) * x_0[12] # # ypr_1 = quaternion_to_euler(x_1[np.newaxis, 3:7])[0][::-1] # x_1[10] = 1.0 * x_1[10] + np.sin(ypr_1[0]) * np.tan(ypr_1[1]) * x_1[11] + np.cos(ypr_1[0]) * np.tan(ypr_1[1]) * \ # x_1[12] # x_1[11] = np.cos(ypr_1[0]) * x_1[11] - np.sin(ypr_1[0]) * x_1[12] # x_1[12] = np.sin(ypr_1[0]) / np.cos(ypr_1[1]) * x_1[11] + np.cos(ypr_1[0]) / np.cos(ypr_1[1]) * x_1[12] # ypr_0 = quaternion_to_euler(x_0[np.newaxis, 3:7])[0][::-1] # x_0[10] = 1.0 * x_0[10] - np.sin(ypr_0[1]) * x_0[12] # x_0[11] = np.cos(ypr_0[0]) * x_0[11] + np.sin(ypr_0[0]) * np.cos(ypr_0[1]) * x_0[12] # x_0[12] = -np.sin(ypr_0[0]) * x_0[11] + np.cos(ypr_0[0]) * np.cos(ypr_0[1]) * x_0[12] # # ypr_1 = quaternion_to_euler(x_1[np.newaxis, 3:7])[0][::-1] # x_1[10] = 1.0 * x_1[10] - np.sin(ypr_1[1]) * x_1[12] # x_1[11] = np.cos(ypr_1[0]) * x_1[11] + np.sin(ypr_1[0]) * np.cos(ypr_1[1]) * x_1[12] # x_1[12] = -np.sin(ypr_1[0]) * x_1[11] + np.cos(ypr_1[0]) * np.cos(ypr_1[1]) * x_1[12] # u = np.hstack([ # data_0['mot 2'], # data_0['mot 3'], # data_0['mot 1'], # data_0['mot 4'], # ]) # # u = u ** 2 * quad.thrust_map[0] / quad.max_thrust yield x_0, x_1, u, dt def f_rate(x, u, quad): """ Time-derivative of the angular rate :param x: 4-length array of input state with components: 3D pos, quaternion angle, 3D vel, 3D rate :param u: control input vector (4-dimensional): [trust_motor_1, ..., thrust_motor_4] :param t_d: disturbance torque (3D) :return: angular rate differential increment (scalar): dr/dt """ rate = x[10:] return np.array([ 1 / quad.J[0] * (u.dot(quad.y_f) + (quad.J[1] - quad.J[2]) * rate[1] * rate[2]), 1 / quad.J[1] * (-u.dot(quad.x_f) + (quad.J[2] - quad.J[0]) * rate[2] * rate[0]), 1 / quad.J[2] * (u.dot(quad.z_l_tau) + (quad.J[0] - quad.J[1]) * rate[0] * rate[1]) ]).squeeze() def resimulate(x_0, x_f, u, dt, quad_mpc, rec_dict): x_pred, _ = quad_mpc.forward_prop(x_0, u, t_horizon=dt, use_gp=False) x_pred = x_pred[-1, np.newaxis, :] rec_dict['state_in'] = np.append(rec_dict['state_in'], x_0[np.newaxis, :], axis=0) rec_dict['input_in'] = np.append(rec_dict['input_in'], u[np.newaxis, :], axis=0) rec_dict['state_out'] = np.append(rec_dict['state_out'], x_f[np.newaxis, :], axis=0) rec_dict['state_ref'] = np.append(rec_dict['state_ref'], np.zeros_like(x_f[np.newaxis, :]), axis=0) rec_dict['timestamp'] = np.append(rec_dict['timestamp'], np.zeros_like(dt)) rec_dict['dt'] = np.append(rec_dict['dt'], dt) rec_dict['state_pred'] = np.append(rec_dict['state_pred'], x_pred, axis=0) rec_dict['error'] = np.append(rec_dict['error'], x_f - x_pred, axis=0) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--quad", type=str, default="kingfisher", help="Name of the quad.") input_arguments = parser.parse_args() ds_name = Conf.ds_name simulation_options = Conf.ds_metadata quad = custom_quad_param_loader(input_arguments.quad) quad_mpc = Quad3DMPC(quad) main(quad_mpc) val(quad_mpc) ================================================ FILE: ros_dd_mpc/src/model_fitting/rdrv_fitting.py ================================================ """ Implementation and fitting of the RDRv linear regression model on flight data. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import os import joblib import argparse import subprocess import numpy as np import matplotlib.pyplot as plt from sklearn.linear_model import LinearRegression from src.model_fitting.gp_common import GPDataset, read_dataset from src.utils.utils import v_dot_q, get_model_dir_and_file, safe_mknode_recursive from config.configuration_parameters import ModelFitConfig as Conf def linear_rdrv_fitting(x, y, feat_idx): drag_coeffs = np.zeros((3, 3)) for i in range(x.shape[1]): reg = LinearRegression(fit_intercept=False).fit(x[:, i, np.newaxis], y[:, i]) drag_coeffs[feat_idx[i], feat_idx[i]] = reg.coef_ return drag_coeffs def load_rdrv(model_options): directory, file_name = get_model_dir_and_file(model_options) rdrv_d = joblib.load(os.path.join(directory, file_name + '.pkl')) return rdrv_d def main(model_name, features, quad_sim_options, dataset_name, x_cap, hist_bins, hist_thresh, plot=False): df_train = read_dataset(dataset_name, True, quad_sim_options) gp_dataset = GPDataset(df_train, features, [], features, cap=x_cap, n_bins=hist_bins, thresh=hist_thresh, visualize_data=False) # Get X,Y datasets for the specified regression dimensions (subset of [7, 8, 9]) a_err_b = gp_dataset.y v_b = gp_dataset.x drag_coeffs = linear_rdrv_fitting(v_b, a_err_b, np.array(features) - 7) # Get git commit hash for saving the model git_version = '' try: git_version = subprocess.check_output(['git', 'describe', '--always']).strip().decode("utf-8") except subprocess.CalledProcessError as e: print(e.returncode, e.output) gp_name_dict = {"git": git_version, "model_name": model_name, "params": quad_sim_options} save_file_path, save_file_name = get_model_dir_and_file(gp_name_dict) safe_mknode_recursive(save_file_path, save_file_name + '.pkl', overwrite=True) file_name = os.path.join(save_file_path, save_file_name + '.pkl') joblib.dump(drag_coeffs, file_name) if not plot: return drag_coeffs # Get X,Y datasets for velocity dimensions [7, 8, 9] a_err_b = gp_dataset.get_y(pruned=True, raw=True)[:, 7:10] v_b = gp_dataset.get_x(pruned=True, raw=True)[:, 7:10] # Compute predictions with RDRv model preds = [] for i in range(len(a_err_b)): preds.append(np.matmul(drag_coeffs, v_b[i])) preds = np.stack(preds) ax_names = [r'$v_B^x$', r'$v_B^y$', r'$v_B^z$'] y_dim_name = [r'drag $a_B^x$', 'drag $a_B^y$', 'drag $a_B^z$'] fig, ax = plt.subplots(len(features), 1, sharex='all') for i in range(len(features)): ax[i].scatter(v_b[:, i], a_err_b[:, i], label='data') ax[i].scatter(v_b[:, i], preds[:, i], label='RDRv') ax[i].legend() ax[i].set_ylabel(y_dim_name[features[i] - 7]) ax[i].set_xlabel(ax_names[features[i] - 7]) ax[0].set_title('Body reference frame') # Remap error to world coordinates and predictions too q = gp_dataset.get_x(pruned=True, raw=True)[:, 3:7] for i in range(len(a_err_b)): a_err_b[i] = v_dot_q(a_err_b[i], q[i]) preds[i] = v_dot_q(preds[i], q[i]) ax_names = [r'$v_W^x$', r'$v_W^y$', r'$v_W^z$'] y_dim_name = [r'drag $a_W^x$', 'drag $a_W^y$', 'drag $a_W^z$'] fig, ax = plt.subplots(len(features), 1, sharex='all') for i in range(len(features)): ax[i].scatter(v_b[:, i], a_err_b[:, i], label='data') ax[i].scatter(v_b[:, i], preds[:, i], label='RDRv') ax[i].legend() ax[i].set_ylabel(y_dim_name[features[i] - 7]) ax[i].set_xlabel(ax_names[features[i] - 7]) ax[0].set_title('World reference frame') plt.show() return drag_coeffs if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--model_name", type=str, default="", help="Name assigned to the trained model.") parser.add_argument('--x', nargs='+', type=int, default=[7, 8, 9], help='Regression X and Y variables. Must be a list of integers between 0 and 12.' 'Velocities xyz correspond to indices 7, 8, 9.') input_arguments = parser.parse_args() reg_dimensions = input_arguments.x rdrv_name = input_arguments.model_name histogram_bins = Conf.histogram_bins histogram_threshold = Conf.histogram_threshold histogram_cap = Conf.velocity_cap ds_name = Conf.ds_name simulation_options = Conf.ds_metadata main(model_name=rdrv_name, features=reg_dimensions, quad_sim_options=simulation_options, dataset_name=ds_name, x_cap=histogram_cap, hist_bins=histogram_bins, hist_thresh=histogram_threshold, plot=Conf.visualize_training_result) ================================================ FILE: ros_dd_mpc/src/model_fitting/system_identification.py ================================================ import argparse import os import numpy as np import pandas as pd import rosbag import rospy from config.configuration_parameters import ModelFitConfig as Conf from src.experiments.point_tracking_and_record import make_record_dict from src.quad_mpc.create_ros_dd_mpc import custom_quad_param_loader from src.quad_mpc.quad_3d_mpc import Quad3DMPC from src.utils.utils import jsonify, get_data_dir_and_file def odometry_parse(odom_msg): p = [odom_msg.pose.pose.position.x, odom_msg.pose.pose.position.y, odom_msg.pose.pose.position.z] q = [odom_msg.pose.pose.orientation.w, odom_msg.pose.pose.orientation.x, odom_msg.pose.pose.orientation.y, odom_msg.pose.pose.orientation.z] v = [odom_msg.twist.twist.linear.x, odom_msg.twist.twist.linear.y, odom_msg.twist.twist.linear.z] w = [odom_msg.twist.twist.angular.x, odom_msg.twist.twist.angular.y, odom_msg.twist.twist.angular.z] return np.array(p + q + v + w) def thrust_motor_model(motor_tau, thrust, thrust_des, dt): if motor_tau < 1e-12: return thrust_des tau_inv = 1 / motor_tau thrust_out = ( tau_inv ** 2 * (thrust_des - 2 * (thrust * thrust_des) ** 0.5 + thrust) * dt ** 2 + 2 * tau_inv * ((thrust * thrust_des) ** 0.5 - thrust) * dt + thrust ) return thrust_out def system_identification(quad_mpc, ds_name, sim_options): """ Processes a rosbag with recorded odometry and send commands to do a system identification of the residual error. :param quad_mpc: Quad model used for forward simulation :param ds_name: Dataset name :param sim_options: Simulation options. :return: """ # Open bag data_file = get_data_dir_and_file(ds_name, training_split=True, params=sim_options, read_only=True) if data_file is None: raise FileNotFoundError rec_file_dir, rec_file_name = data_file rec_file_name_bag = rec_file_name.replace('.csv', '.bag') rec_file = os.path.join(rec_file_dir, rec_file_name_bag) save_file = os.path.join(rec_file_dir, rec_file_name) bag = rosbag.Bag(rec_file) # Re-simulate thrust with motor model and communication delay thrust = None control = {'t': [], 'thrust': []} for topic, msg, t in bag.read_messages(topics=['control']): t = msg.header.stamp desired_thrust = np.array(msg.thrusts) if thrust is None: thrust = desired_thrust t = t + rospy.Duration.from_sec(quad_mpc.quad.comm_delay) + rospy.Duration.from_sec(0.001) for dt in np.arange(0.001, 0.021, step=0.001): thrust = thrust_motor_model(quad_mpc.quad.motor_tau, thrust, desired_thrust, 0.001) t_at = t + rospy.Duration.from_sec(dt) control['t'].append(t_at.to_sec()) control['thrust'].append(thrust) control = pd.DataFrame(control) control = control.set_index('t') control = control[~control.index.duplicated(keep='last')] control = control.sort_index() state = {'t': [], 'state': []} states_list = [] recording = False for topic, msg, t in bag.read_messages(topics=['recording_ctrl', 'state']): if topic == 'recording_ctrl': recording = msg.data if topic == 'state' and recording: t = msg.header.stamp state['t'].append(t.to_sec()) x = odometry_parse(msg) state['state'].append(x) states_list.append(np.array(x)) states_np = np.array(states_list) filtered_states_np = states_np.copy() from scipy.signal import savgol_filter window_size_xy = 31 window_size_z = 31 window_size_q = 31 window_size_v = 31 window_size_w = 121 poly_order = 4 filtered_states_np[:, :2] = savgol_filter(filtered_states_np[:, :2], window_size_xy, poly_order, axis=0) filtered_states_np[:, 2] = savgol_filter(filtered_states_np[:, 2], window_size_z, poly_order, axis=0) filtered_states_np[:, 3:7] = savgol_filter(filtered_states_np[:, 3:7], window_size_q, 2, axis=0) filtered_states_np[:, 3:7] = filtered_states_np[:, 3:7] / np.linalg.norm(filtered_states_np[:, 3:7], axis=1, keepdims=True) filtered_states_np[:, 7:10] = savgol_filter(filtered_states_np[:, 7:10], window_size_v, 2, axis=0) filtered_states_np[:, 10:] = savgol_filter(filtered_states_np[:, 10:], window_size_w, poly_order, axis=0) # def extract_windows(array, sub_window_size): # examples = [] # # for i in range(0, array.shape[0]-sub_window_size, sub_window_size): # example = array[i:sub_window_size + i] # examples.append(np.expand_dims(example, 0)) # # return np.vstack(examples) # # from scipy.spatial.transform import Rotation as R # from scipy.spatial.transform import RotationSpline # window_rot = extract_windows(filtered_states_np[:, 3:7], 20) # mean_rot = [] # for i in range(window_rot.shape[0]): # mean_rot.append(R.from_quat(window_rot[i]).mean().as_quat()) # mean_rot = np.vstack(mean_rot) # rotations = R.from_quat(mean_rot) # spline = RotationSpline(state['t'][9::20][:mean_rot.shape[0]], rotations) # filtered_states_np[:, 3:7] = spline(state['t']).as_quat() import matplotlib.pyplot as plt # plt.plot(states_np[:, 9]) #plt.plot(filtered_states_np[3975:4025, 0]) #plt.plot(states_np[3975:4025, 0]) #plt.plot(np.gradient(filtered_states_np[3975:4025, 1])/1e-2) #plt.plot(filtered_states_np[3975:4025, 7]) #plt.plot(filtered_states_np[1800:2000, 3]) #plt.plot(states_np[1800:2000, 3]) #plt.plot(filtered_states_np[3775:4025, 3:7] * np.sign(filtered_states_np[3775:4025, 3])) #plt.plot(filtered_states_np[3775:4025, 10]) #plt.plot(states_np[3775:4025, 10]) #plt.show() state = {'t': state['t'], 'state': []} for i in range(filtered_states_np.shape[0]): state['state'].append(filtered_states_np[i]) rec_dict = make_record_dict(state_dim=13) for i in range(1, len(state['t'])): last_state_idx = control.index.get_loc(state['t'][i-1], method='nearest') curr_state_idx = control.index.get_loc(state['t'][i], method='nearest') u0 = control['thrust'].iloc[last_state_idx] / quad_mpc.quad.max_thrust u1 = control['thrust'].iloc[curr_state_idx] / quad_mpc.quad.max_thrust # Only use subsequent states with similar thrust state for system identification (within 1%) if np.all(np.abs(u0 - u1) < 0.01): x_0 = state['state'][i-1] x_f = state['state'][i] u = np.vstack([u0, u1]).mean(0) dt = state['t'][i] - state['t'][i-1] if dt > 0.015: continue x_pred, _ = quad_mpc.forward_prop(x_0, u, t_horizon=dt, use_gp=False) x_pred = x_pred[-1, np.newaxis, :] rec_dict['state_in'] = np.append(rec_dict['state_in'], x_0[np.newaxis, :], axis=0) rec_dict['input_in'] = np.append(rec_dict['input_in'], u[np.newaxis, :], axis=0) rec_dict['state_out'] = np.append(rec_dict['state_out'], x_f[np.newaxis, :], axis=0) rec_dict['state_ref'] = np.append(rec_dict['state_ref'], np.zeros_like(x_f[np.newaxis, :]), axis=0) rec_dict['timestamp'] = np.append(rec_dict['timestamp'], state['t'][i]) rec_dict['dt'] = np.append(rec_dict['dt'], dt) rec_dict['state_pred'] = np.append(rec_dict['state_pred'], x_pred, axis=0) rec_dict['error'] = np.append(rec_dict['error'], x_f - x_pred, axis=0) # Save datasets for key in rec_dict.keys(): print(key, " ", rec_dict[key].shape) rec_dict[key] = jsonify(rec_dict[key]) df = pd.DataFrame(rec_dict) df.to_csv(save_file, index=True, header=True) return import matplotlib.pyplot as plt #plt.plot(states_np[:, 9]) plt.plot(np.gradient(out[:, 8])) plt.show() return # Match each state to the corresponding thrust state of the quad rec_dict = make_record_dict(state_dim=13) recording = False last_state_msg = None for topic, msg, t in bag.read_messages(topics=['recording_ctrl', 'state']): if topic == 'recording_ctrl': recording = msg.data if not recording: last_state_msg = None if topic == 'state' and recording: if last_state_msg is not None: last_state_idx = control.index.get_loc(last_state_msg.header.stamp.to_sec(), method='nearest') curr_state_idx = control.index.get_loc(msg.header.stamp.to_sec(), method='nearest') u0 = control['thrust'].iloc[last_state_idx] / quad_mpc.quad.max_thrust u1 = control['thrust'].iloc[curr_state_idx] / quad_mpc.quad.max_thrust # Only use subsequent states with similar thrust state for system identification (within 1%) if np.all(np.abs(u0 - u1) < 0.01): x_0 = odometry_parse(last_state_msg) x_f = odometry_parse(msg) u = np.vstack([u0, u1]).mean(0) dt = msg.header.stamp.to_sec() - last_state_msg.header.stamp.to_sec() x_pred, _ = quad_mpc.forward_prop(x_0, u, t_horizon=dt, use_gp=False) x_pred = x_pred[-1, np.newaxis, :] rec_dict['state_in'] = np.append(rec_dict['state_in'], x_0[np.newaxis, :], axis=0) rec_dict['input_in'] = np.append(rec_dict['input_in'], u[np.newaxis, :], axis=0) rec_dict['state_out'] = np.append(rec_dict['state_out'], x_f[np.newaxis, :], axis=0) rec_dict['state_ref'] = np.append(rec_dict['state_ref'], np.zeros_like(x_f[np.newaxis, :]), axis=0) rec_dict['timestamp'] = np.append(rec_dict['timestamp'], msg.header.stamp.to_sec()) rec_dict['dt'] = np.append(rec_dict['dt'], dt) rec_dict['state_pred'] = np.append(rec_dict['state_pred'], x_pred, axis=0) rec_dict['error'] = np.append(rec_dict['error'], x_f - x_pred, axis=0) last_state_msg = msg # Save datasets for key in rec_dict.keys(): print(key, " ", rec_dict[key].shape) rec_dict[key] = jsonify(rec_dict[key]) df = pd.DataFrame(rec_dict) df.to_csv(save_file, index=True, header=True) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--quad", type=str, default="kingfisher", help="Name of the quad.") input_arguments = parser.parse_args() ds_name = Conf.ds_name simulation_options = Conf.ds_metadata quad = custom_quad_param_loader(input_arguments.quad) quad_mpc = Quad3DMPC(quad) system_identification(quad_mpc, ds_name, simulation_options) ================================================ FILE: ros_dd_mpc/src/quad_mpc/__init__.py ================================================ ================================================ FILE: ros_dd_mpc/src/quad_mpc/create_ros_dd_mpc.py ================================================ """ Class for interfacing the data-augmented MPC with ROS. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ from src.quad_mpc.quad_3d import Quadrotor3D from src.quad_mpc.quad_3d_mpc import Quad3DMPC import numpy as np import os import yaml def custom_quad_param_loader(quad_name): this_path = os.path.dirname(os.path.realpath(__file__)) params_file = os.path.join(this_path, '..', '..', 'config', quad_name + '.yaml') # Get parameters for drone with open(params_file, "r") as stream: attrib = yaml.safe_load(stream) quad = Quadrotor3D(noisy=False, drag=False, payload=False, motor_noise=False) quad.mass = float(attrib['mass']) + (float(attrib['mass_rotor']) if 'mass_rotor' in attrib else 0) * 4 quad.J = np.array(attrib['inertia']) quad.thrust_map = attrib["thrust_map"] if 'thrust_max' in attrib: quad.max_thrust = attrib['thrust_max'] else: float(attrib["motor_omega_max"]) ** 2 * quad.thrust_map[0] quad.c = float(attrib['kappa']) if 'arm_length' in attrib and attrib['rotors_config'] == 'cross': quad.length = float(attrib['arm_length']) h = np.cos(np.pi / 4) * quad.length quad.x_f = np.array([h, -h, -h, h]) quad.y_f = np.array([-h, h, -h, h]) elif 'arm_length' in attrib and attrib['rotors_config'] == 'plus': quad.length = float(attrib['arm_length']) quad.x_f = np.array([0, 0, -quad.length, quad.length]) quad.y_f = np.array([-quad.length, quad.length, 0, 0]) else: tbm_fr = np.array(attrib['tbm_fr']) tbm_bl = np.array(attrib['tbm_bl']) tbm_br = np.array(attrib['tbm_br']) tbm_fl = np.array(attrib['tbm_fl']) quad.length = np.linalg.norm(tbm_fr) quad.x_f = np.array([tbm_fr[0], tbm_bl[0], tbm_br[0], tbm_fl[0]]) quad.y_f = np.array([tbm_fr[1], tbm_bl[1], tbm_br[1], tbm_fl[1]]) quad.z_l_tau = np.array([-quad.c, -quad.c, quad.c, quad.c]) if 'motor_tau' in attrib: quad.motor_tau = attrib['motor_tau'] if 'comm_delay' in attrib: quad.comm_delay = attrib['comm_delay'] return quad class ROSDDMPC: def __init__(self, t_horizon, n_mpc_nodes, opt_dt, quad_name, point_reference=False, models=None, model_conf=None, rdrv=None): quad = custom_quad_param_loader(quad_name) # Initialize quad MPC if point_reference: acados_config = { "solver_type": "SQP", "terminal_cost": True } else: acados_config = { "solver_type": "SQP_RTI", "terminal_cost": False } q_diagonal = np.array([0.5, 0.5, 0.5, 0.1, 0.1, 0.1, 0.05, 0.05, 0.05, 0.01, 0.01, 0.01]) r_diagonal = np.array([1.0, 1.0, 1.0, 1.0]) q_mask = np.array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]).T quad_mpc = Quad3DMPC(quad, t_horizon=t_horizon, optimization_dt=opt_dt, n_nodes=n_mpc_nodes, pre_trained_models=models, model_conf=model_conf, model_name=quad_name, solver_options=acados_config, q_mask=q_mask, q_cost=q_diagonal, r_cost=r_diagonal, rdrv_d_mat=rdrv) self.quad_name = quad_name self.quad = quad self.quad_mpc = quad_mpc # Last optimization self.last_w = None def set_state(self, x): """ Set quadrotor state estimate from an odometry measurement :param x: measured state from odometry. List with 13 components with the format: [p_xyz, q_wxyz, v_xyz, w_xyz] """ self.quad.set_state(x) def set_gp_state(self, x): """ Set a quadrotor state estimate from an odometry measurement. While the input state in the function set_state() will be used as initial state for the MPC optimization, the input state of this function will be used to evaluate the GP's. If this function is never called, then the GP's will be evaluated with the state from set_state() too. :param x: measured state from odometry. List with 13 components with the format: [p_xyz, q_wxyz, v_xyz, w_xyz] """ self.quad.set_gp_state(x) def set_reference(self, x_ref, u_ref): """ Set a reference state for the optimizer. :param x_ref: list with 4 sub-components (position, angle quaternion, velocity, body rate). If these four are lists, then this means a single target point is used. If they are Nx3 and Nx4 (for quaternion) numpy arrays, then they are interpreted as a sequence of N tracking points. :param u_ref: Optional target for the optimized control inputs """ return self.quad_mpc.set_reference(x_reference=x_ref, u_reference=u_ref) def optimize(self, model_data): w_opt, x_opt = self.quad_mpc.optimize(use_model=model_data, return_x=True) return w_opt, x_opt ================================================ FILE: ros_dd_mpc/src/quad_mpc/quad_3d.py ================================================ """ Implementation of the Simplified Simulator and its quadrotor dynamics. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ from math import sqrt import numpy as np from src.utils.utils import quaternion_to_euler, skew_symmetric, v_dot_q, unit_quat, quaternion_inverse class Quadrotor3D: def __init__(self, noisy=False, drag=False, payload=False, motor_noise=False): """ Initialization of the 3D quadrotor class :param noisy: Whether noise is used in the simulation :type noisy: bool :param drag: Whether to simulate drag or not. :type drag: bool :param payload: Whether to simulate a payload force in the simulation :type payload: bool :param motor_noise: Whether non-gaussian noise is considered in the motor inputs :type motor_noise: bool """ # Either 'x' or '+'. The xy configuration of the thrusters in the body plane. configuration = 'x' # Maximum thrust in Newtons of a thruster when rotating at maximum speed. self.max_thrust = 20 self.thrust_map = 1e-6 # Motor time constant self.motor_tau = 0. # Communication delay self.comm_delay = 0. # System state space self.pos = np.zeros((3,)) self.vel = np.zeros((3,)) self.angle = np.array([1., 0., 0., 0.]) # Quaternion format: qw, qx, qy, qz self.a_rate = np.zeros((3,)) # System state space for GP evaluation. self.gp_pos = None self.gp_vel = None self.gp_angle = None self.gp_a_rate = None # Input constraints self.max_input_value = 1 # Motors at full thrust self.min_input_value = 0 # Motors turned off # Quadrotor intrinsic parameters self.J = np.array([.03, .03, .06]) # N m s^2 = kg m^2 self.mass = 1.0 # kg # Length of motor to CoG segment self.length = 0.47 / 2 # m # Positions of thrusters if configuration == '+': self.x_f = np.array([self.length, 0, -self.length, 0]) self.y_f = np.array([0, self.length, 0, -self.length]) elif configuration == 'x': h = np.cos(np.pi / 4) * self.length self.x_f = np.array([h, -h, -h, h]) self.y_f = np.array([-h, -h, h, h]) # For z thrust torque calculation self.c = 0.013 # m (z torque generated by each motor) self.z_l_tau = np.array([-self.c, self.c, -self.c, self.c]) # Gravity vector self.g = np.array([[0], [0], [9.8066]]) # m s^-2 # Actuation thrusts self.u_noiseless = np.array([0.0, 0.0, 0.0, 0.0]) self.u = np.array([0.0, 0.0, 0.0, 0.0]) # N # Drag coefficients [kg / m] self.rotor_drag_xy = 0.3 self.rotor_drag_z = 0.0 # No rotor drag in the z dimension self.rotor_drag = np.array([self.rotor_drag_xy, self.rotor_drag_xy, self.rotor_drag_z])[:, np.newaxis] self.aero_drag = 0.08 self.drag = drag self.noisy = noisy self.motor_noise = motor_noise self.payload_mass = 0.3 # kg self.payload_mass = self.payload_mass * payload def set_state(self, *args, **kwargs): if len(args) != 0: assert len(args) == 1 and len(args[0]) == 13 self.pos[0], self.pos[1], self.pos[2], \ self.angle[0], self.angle[1], self.angle[2], self.angle[3], \ self.vel[0], self.vel[1], self.vel[2], \ self.a_rate[0], self.a_rate[1], self.a_rate[2] \ = args[0] else: self.pos = kwargs["pos"] self.angle = kwargs["angle"] self.vel = kwargs["vel"] self.a_rate = kwargs["rate"] def set_gp_state(self, *args, **kwargs): if self.gp_pos is None: # Instantiate these variables for the first time self.gp_pos = np.zeros((3,)) self.gp_vel = np.zeros((3,)) self.gp_angle = np.array([1., 0., 0., 0.]) # Quaternion format: qw, qx, qy, qz self.gp_a_rate = np.zeros((3,)) if len(args) != 0: assert len(args) == 1 and len(args[0]) == 13 self.gp_pos[0], self.gp_pos[1], self.gp_pos[2], \ self.gp_angle[0], self.gp_angle[1], self.gp_angle[2], self.gp_angle[3], \ self.gp_vel[0], self.gp_vel[1], self.gp_vel[2], \ self.gp_a_rate[0], self.gp_a_rate[1], self.gp_a_rate[2] \ = args[0] else: self.gp_pos = kwargs["pos"] self.gp_angle = kwargs["angle"] self.gp_vel = kwargs["vel"] self.gp_a_rate = kwargs["rate"] def get_state(self, quaternion=False, stacked=False): if quaternion and not stacked: return [self.pos, self.angle, self.vel, self.a_rate] if quaternion and stacked: return [self.pos[0], self.pos[1], self.pos[2], self.angle[0], self.angle[1], self.angle[2], self.angle[3], self.vel[0], self.vel[1], self.vel[2], self.a_rate[0], self.a_rate[1], self.a_rate[2]] angle = quaternion_to_euler(self.angle) if not quaternion and stacked: return [self.pos[0], self.pos[1], self.pos[2], angle[0], angle[1], angle[2], self.vel[0], self.vel[1], self.vel[2], self.a_rate[0], self.a_rate[1], self.a_rate[2]] return [self.pos, angle, self.vel, self.a_rate] def get_gp_state(self, quaternion=False, stacked=False): if self.gp_pos is None: return None if quaternion and not stacked: return [self.gp_pos, self.gp_angle, self.gp_vel, self.gp_a_rate] if quaternion and stacked: return [self.gp_pos[0], self.gp_pos[1], self.gp_pos[2], self.gp_angle[0], self.gp_angle[1], self.gp_angle[2], self.gp_angle[3], self.gp_vel[0], self.gp_vel[1], self.gp_vel[2], self.gp_a_rate[0], self.gp_a_rate[1], self.gp_a_rate[2]] gp_angle = quaternion_to_euler(self.gp_angle) if not quaternion and stacked: return [self.gp_pos[0], self.gp_pos[1], self.gp_pos[2], gp_angle[0], gp_angle[1], gp_angle[2], self.gp_vel[0], self.gp_vel[1], self.gp_vel[2], self.gp_a_rate[0], self.gp_a_rate[1], self.gp_a_rate[2]] return [self.gp_pos, gp_angle, self.gp_vel, self.gp_a_rate] def get_control(self, noisy=False): if not noisy: return self.u_noiseless else: return self.u def update(self, u, dt): """ Runge-Kutta 4th order dynamics integration :param u: 4-dimensional vector with components between [0.0, 1.0] that represent the activation of each motor. :param dt: time differential """ # Clip inputs for i, u_i in enumerate(u): self.u_noiseless[i] = max(min(u_i, self.max_input_value), self.min_input_value) # Apply noise to inputs (uniformly distributed noise with standard deviation proportional to input magnitude) if self.motor_noise: for i, u_i in enumerate(self.u_noiseless): std = 0.02 * sqrt(u_i) noise_u = np.random.normal(loc=0.1 * (u_i / 1.3) ** 2, scale=std) self.u[i] = max(min(u_i - noise_u, self.max_input_value), self.min_input_value) * self.max_thrust else: self.u = self.u_noiseless * self.max_thrust # Generate disturbance forces / torques if self.noisy: f_d = np.random.normal(size=(3, 1), scale=10 * dt) t_d = np.random.normal(size=(3, 1), scale=10 * dt) else: f_d = np.zeros((3, 1)) t_d = np.zeros((3, 1)) x = self.get_state(quaternion=True, stacked=False) # RK4 integration k1 = [self.f_pos(x), self.f_att(x), self.f_vel(x, self.u, f_d), self.f_rate(x, self.u, t_d)] x_aux = [x[i] + dt / 2 * k1[i] for i in range(4)] k2 = [self.f_pos(x_aux), self.f_att(x_aux), self.f_vel(x_aux, self.u, f_d), self.f_rate(x_aux, self.u, t_d)] x_aux = [x[i] + dt / 2 * k2[i] for i in range(4)] k3 = [self.f_pos(x_aux), self.f_att(x_aux), self.f_vel(x_aux, self.u, f_d), self.f_rate(x_aux, self.u, t_d)] x_aux = [x[i] + dt * k3[i] for i in range(4)] k4 = [self.f_pos(x_aux), self.f_att(x_aux), self.f_vel(x_aux, self.u, f_d), self.f_rate(x_aux, self.u, t_d)] x = [x[i] + dt * (1.0 / 6.0 * k1[i] + 2.0 / 6.0 * k2[i] + 2.0 / 6.0 * k3[i] + 1.0 / 6.0 * k4[i]) for i in range(4)] # Ensure unit quaternion x[1] = unit_quat(x[1]) self.pos, self.angle, self.vel, self.a_rate = x def f_pos(self, x): """ Time-derivative of the position vector :param x: 4-length array of input state with components: 3D pos, quaternion angle, 3D vel, 3D rate :return: position differential increment (vector): d[pos_x; pos_y]/dt """ vel = x[2] return vel def f_att(self, x): """ Time-derivative of the attitude in quaternion form :param x: 4-length array of input state with components: 3D pos, quaternion angle, 3D vel, 3D rate :return: attitude differential increment (quaternion qw, qx, qy, qz): da/dt """ rate = x[3] angle_quaternion = x[1] return 1 / 2 * skew_symmetric(rate).dot(angle_quaternion) def f_vel(self, x, u, f_d): """ Time-derivative of the velocity vector :param x: 4-length array of input state with components: 3D pos, quaternion angle, 3D vel, 3D rate :param u: control input vector (4-dimensional): [trust_motor_1, ..., thrust_motor_4] :param f_d: disturbance force vector (3-dimensional) :return: 3D velocity differential increment (vector): d[vel_x; vel_y; vel_z]/dt """ a_thrust = np.array([[0], [0], [np.sum(u)]]) / self.mass if self.drag: # Transform velocity to body frame v_b = v_dot_q(x[2], quaternion_inverse(x[1]))[:, np.newaxis] # Compute aerodynamic drag acceleration in world frame a_drag = -self.aero_drag * v_b ** 2 * np.sign(v_b) / self.mass # Add rotor drag a_drag -= self.rotor_drag * v_b / self.mass # Transform drag acceleration to world frame a_drag = v_dot_q(a_drag, x[1]) else: a_drag = np.zeros((3, 1)) angle_quaternion = x[1] a_payload = -self.payload_mass * self.g / self.mass return np.squeeze(-self.g + a_payload + a_drag + v_dot_q(a_thrust + f_d / self.mass, angle_quaternion)) def f_rate(self, x, u, t_d): """ Time-derivative of the angular rate :param x: 4-length array of input state with components: 3D pos, quaternion angle, 3D vel, 3D rate :param u: control input vector (4-dimensional): [trust_motor_1, ..., thrust_motor_4] :param t_d: disturbance torque (3D) :return: angular rate differential increment (scalar): dr/dt """ rate = x[3] return np.array([ 1 / self.J[0] * (u.dot(self.y_f) + t_d[0] + (self.J[1] - self.J[2]) * rate[1] * rate[2]), 1 / self.J[1] * (-u.dot(self.x_f) + t_d[1] + (self.J[2] - self.J[0]) * rate[2] * rate[0]), 1 / self.J[2] * (u.dot(self.z_l_tau) + t_d[2] + (self.J[0] - self.J[1]) * rate[0] * rate[1]) ]).squeeze() ================================================ FILE: ros_dd_mpc/src/quad_mpc/quad_3d_mpc.py ================================================ """ Implementation of the data-augmented MPC for quadrotors. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import numpy as np import torch from src.quad_mpc.quad_3d_optimizer import Quad3DOptimizer from src.model_fitting.gp_common import restore_gp_regressors from src.utils.quad_3d_opt_utils import simulate_plant, uncertainty_forward_propagation from src.utils.utils import make_bx_matrix class Quad3DMPC: def __init__(self, my_quad, t_horizon=1.0, n_nodes=5, q_cost=None, r_cost=None, optimization_dt=5e-2, simulation_dt=5e-4, pre_trained_models=None, model_name="my_quad", q_mask=None, solver_options=None, rdrv_d_mat=None, model_conf=None): """ :param my_quad: Quadrotor3D simulator object :type my_quad: Quadrotor3D :param t_horizon: time horizon for optimization loop MPC controller :param n_nodes: number of MPC nodes :param optimization_dt: time step between two successive optimizations intended to be used. :param simulation_dt: discretized time-step for the quadrotor simulation :param pre_trained_models: additional pre-trained GP regressors to be combined with nominal model in the MPC :param q_cost: diagonal of Q matrix for LQR cost of MPC cost function. Must be a numpy array of length 13. :param r_cost: diagonal of R matrix for LQR cost of MPC cost function. Must be a numpy array of length 4. :param q_mask: Optional boolean mask that determines which variables from the state compute towards the cost function. In case no argument is passed, all variables are weighted. :param solver_options: Optional set of extra options dictionary for acados solver. :param rdrv_d_mat: 3x3 matrix that corrects the drag with a linear model according to Faessler et al. 2018. None if not used """ if rdrv_d_mat is not None: # rdrv is currently not compatible with covariance mode or with GP-MPC. print("RDRv mode") self.rdrv = rdrv_d_mat assert pre_trained_models is None else: self.rdrv = None self.quad = my_quad self.simulation_dt = simulation_dt self.optimization_dt = optimization_dt # motor commands from last step self.motor_u = np.array([0., 0., 0., 0.]) self.n_nodes = n_nodes self.t_horizon = t_horizon self.mlp = None self.mlp_approx = None # Load augmented dynamics model with GP regressor if isinstance(pre_trained_models, torch.nn.Module): self.gp_ensemble = None self.B_x = {} x_dims = len(my_quad.get_state(quaternion=True, stacked=True)) pred_dims = [7, 8, 9] + ([10, 11, 12] if model_conf['torque_output'] else []) for y_dim in pred_dims: self.B_x[y_dim] = make_bx_matrix(x_dims, [y_dim]) self.mlp = pre_trained_models elif pre_trained_models is not None: self.gp_ensemble = restore_gp_regressors(pre_trained_models) x_dims = len(my_quad.get_state(quaternion=True, stacked=True)) self.B_x = {} for y_dim in self.gp_ensemble.gp.keys(): self.B_x[y_dim] = make_bx_matrix(x_dims, [y_dim]) else: self.gp_ensemble = None self.B_x = {} # Selection matrix of the GP regressor-modified system states # For MPC optimization use self.quad_opt = Quad3DOptimizer(my_quad, t_horizon=t_horizon, n_nodes=n_nodes, q_cost=q_cost, r_cost=r_cost, B_x=self.B_x, gp_regressors=self.gp_ensemble, mlp_regressor=self.mlp, mlp_conf=model_conf, model_name=model_name, q_mask=q_mask, solver_options=solver_options, rdrv_d_mat=rdrv_d_mat) def clear(self): self.quad_opt.clear_acados_model() def get_state(self): """ Returns the state of the drone, with the angle described as a wxyz quaternion :return: 13x1 array with the drone state: [p_xyz, a_wxyz, v_xyz, r_xyz] """ x = np.expand_dims(self.quad.get_state(quaternion=True, stacked=True), 1) return x def set_reference(self, x_reference, u_reference=None): """ Sets a target state for the MPC optimizer :param x_reference: list with 4 sub-components (position, angle quaternion, velocity, body rate). If these four are lists, then this means a single target point is used. If they are Nx3 and Nx4 (for quaternion) numpy arrays, then they are interpreted as a sequence of N tracking points. :param u_reference: Optional target for the optimized control inputs """ if isinstance(x_reference[0], list): # Target state is just a point return self.quad_opt.set_reference_state(x_reference, u_reference) else: # Target state is a sequence return self.quad_opt.set_reference_trajectory(x_reference, u_reference) def optimize(self, use_model=0, return_x=False): """ Runs MPC optimization to reach the pre-set target. :param use_model: Integer. Select which dynamics model to use from the available options. :param return_x: bool, whether to also return the optimized sequence of states alongside with the controls. :return: 4*m vector of optimized control inputs with the format: [u_1(0), u_2(0), u_3(0), u_4(0), u_1(1), ..., u_3(m-1), u_4(m-1)]. If return_x is True, will also return a vector of shape N+1 x 13 containing the optimized state prediction. """ quad_current_state = self.quad.get_state(quaternion=True, stacked=True) quad_gp_state = self.quad.get_gp_state(quaternion=True, stacked=True) # Remove rate state for simplified model NLP out_out = self.quad_opt.run_optimization(quad_current_state, use_model=use_model, return_x=return_x, gp_regression_state=quad_gp_state) return out_out def simulate(self, ref_u): """ Runs the simulation step for the dynamics model of the quadrotor 3D. :param ref_u: 4-length reference vector of control inputs """ # Simulate step self.quad.update(ref_u, self.simulation_dt) def simulate_plant(self, w_opt, t_horizon=None, dt_vec=None, progress_bar=False): """ Given a sequence of n inputs, evaluates the simulated discrete-time plant model n steps into the future. The current drone state will not be changed by calling this method. :param w_opt: sequence of control n x m control inputs, where n is the number of steps and m is the dimensionality of a control input. :param t_horizon: time corresponding to the duration of the n control inputs. In the case that the w_opt comes from an MPC optimization, this parameter should be the MPC time horizon. :param dt_vec: a vector of timestamps, the same length as w_opt, corresponding to the total time each input is applied. :param progress_bar: boolean - whether to show a progress bar on the console or not. :return: the sequence of simulated quadrotor states. """ if t_horizon is None and dt_vec is None: t_horizon = self.t_horizon return simulate_plant(self.quad, w_opt, simulation_dt=self.simulation_dt, simulate_func=self.simulate, t_horizon=t_horizon, dt_vec=dt_vec, progress_bar=progress_bar) def forward_prop(self, x_0, w_opt, cov_0=None, t_horizon=None, dt=None, use_gp=False, use_model=0): """ Computes the forward propagation of the state using an MPC-optimized control input sequence. :param x_0: Initial n-length state vector :param w_opt: Optimized m*4-length sequence of control inputs from MPC, with the vector format: [u_1(1), u_2(1), u_3(1), u_4(1), ..., u_3(m), u_4(m)] :param cov_0: Initial covariance estimate (default 0). Can be either a positive semi-definite matrix or a 1D vector, which will be the diagonal of the covariance matrix. In both cases, the resulting covariance matrix must be nxn shape, where n is the length of x_0. :param t_horizon: time span of the control inputs (default is the time horizon of the MPC) :param dt: Optional. Vector of length m, with the corresponding integration time for every control input in w_opt. If none is provided, the default integration step is used. :param use_gp: Boolean, whether to use GP regressors when performing the integration or not. :param use_model: Integer. Select which dynamics model to use from the available options. :return: An m x n array of propagated (expected) state estimates, and an m x n x n array with the m propagated covariance matrices. """ # Default parameters if dt is not None: assert len(dt) == len(w_opt) / 4 t_horizon = np.sum(dt) if t_horizon is None: t_horizon = self.t_horizon if cov_0 is None: cov_0 = np.diag(np.zeros_like(np.squeeze(x_0))) elif len(cov_0.shape) == 1: cov_0 = np.diag(cov_0) elif len(cov_0.shape) > 2: TypeError("The initial covariance value must be either a 1D vector of a 2D matrix") gp_ensemble = self.gp_ensemble if use_gp else None # Compute forward propagation of state pdf return uncertainty_forward_propagation(x_0, w_opt, t_horizon=t_horizon, covar=cov_0, dt=dt, discrete_dynamics_f=self.quad_opt.discretize_f_and_q, dynamics_jac_f=self.quad_opt.quad_xdot_jac, B_x=self.B_x, gp_regressors=gp_ensemble, use_model=use_model, m_integration_steps=1) @staticmethod def reshape_input_sequence(u_seq): """ Reshapes the an output trajectory from the 1D format: [u_0(0), u_1(0), ..., u_0(n-1), u_1(n-1), ..., u_m-1(n-1)] to a 2D n x m array. :param u_seq: 1D input sequence :return: 2D input sequence, were n is the number of control inputs and m is the dimension of a single input. """ k = np.arange(u_seq.shape[0] / 4, dtype=int) u_seq = np.atleast_2d(u_seq).T if len(u_seq.shape) == 1 else u_seq u_seq = np.concatenate((u_seq[4 * k], u_seq[4 * k + 1], u_seq[4 * k + 2], u_seq[4 * k + 3]), 1) return u_seq def reset(self): return ================================================ FILE: ros_dd_mpc/src/quad_mpc/quad_3d_optimizer.py ================================================ """ Implementation of the nonlinear optimizer for the data-augmented MPC. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import os import sys import shutil import casadi as cs import numpy as np from copy import copy import torch from acados_template import AcadosOcp, AcadosOcpSolver, AcadosModel from config.configuration_parameters import GroundEffectMapConfig from src.quad_mpc.quad_3d import Quadrotor3D from src.model_fitting.gp import GPEnsemble from src.utils.ground_map import GroundMapWithBox from src.utils.utils import skew_symmetric, v_dot_q, safe_mkdir_recursive, quaternion_inverse from src.utils.quad_3d_opt_utils import discretize_dynamics_and_cost class Quad3DOptimizer: def __init__(self, quad, t_horizon=1, n_nodes=20, q_cost=None, r_cost=None, q_mask=None, B_x=None, gp_regressors=None, mlp_conf=None, mlp_regressor=None, rdrv_d_mat=None, model_name="quad_3d_acados_mpc", solver_options=None): """ :param quad: quadrotor object :type quad: Quadrotor3D :param t_horizon: time horizon for MPC optimization :param n_nodes: number of optimization nodes until time horizon :param q_cost: diagonal of Q matrix for LQR cost of MPC cost function. Must be a numpy array of length 12. :param r_cost: diagonal of R matrix for LQR cost of MPC cost function. Must be a numpy array of length 4. :param B_x: dictionary of matrices that maps the outputs of the gp regressors to the state space. :param gp_regressors: Gaussian Process ensemble for correcting the nominal model :type gp_regressors: GPEnsemble :param q_mask: Optional boolean mask that determines which variables from the state compute towards the cost function. In case no argument is passed, all variables are weighted. :param solver_options: Optional set of extra options dictionary for solvers. :param rdrv_d_mat: 3x3 matrix that corrects the drag with a linear model according to Faessler et al. 2018. None if not used """ # Weighted squared error loss function q = (p_xyz, a_xyz, v_xyz, r_xyz), r = (u1, u2, u3, u4) if q_cost is None: q_cost = np.array([10, 10, 10, 0.1, 0.1, 0.1, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05]) if r_cost is None: r_cost = np.array([0.1, 0.1, 0.1, 0.1]) self.T = t_horizon # Time horizon self.N = n_nodes # number of control nodes within horizon self.quad = quad self.max_u = quad.max_input_value self.min_u = quad.min_input_value # Declare model variables self.p = cs.MX.sym('p', 3) # position self.q = cs.MX.sym('a', 4) # angle quaternion (wxyz) self.v = cs.MX.sym('v', 3) # velocity self.r = cs.MX.sym('r', 3) # angle rate # Full state vector (13-dimensional) self.x = cs.vertcat(self.p, self.q, self.v, self.r) self.state_dim = 13 # Control input vector u1 = cs.MX.sym('u1') u2 = cs.MX.sym('u2') u3 = cs.MX.sym('u3') u4 = cs.MX.sym('u4') self.u = cs.vertcat(u1, u2, u3, u4) # Nominal model equations symbolic function (no GP) self.quad_xdot_nominal = self.quad_dynamics(rdrv_d_mat) # Linearized model dynamics symbolic function self.quad_xdot_jac = self.linearized_quad_dynamics() # Initialize objective function, 0 target state and integration equations self.L = None self.target = None self.u_target = None self.mlp_regressor = None self.mlp_conf = None self.x_opt_acados = None self.w_opt_acados = None self.mlp_params = None # Check if GP ensemble has an homogeneous feature space (if actual Ensemble) if gp_regressors is not None and gp_regressors.homogeneous: self.gp_reg_ensemble = gp_regressors self.B_x = [B_x[dim] for dim in gp_regressors.dim_idx] self.B_x = np.squeeze(np.stack(self.B_x, axis=1)) self.B_x = self.B_x[:, np.newaxis] if len(self.B_x.shape) == 1 else self.B_x self.B_z = gp_regressors.B_z elif gp_regressors and gp_regressors.no_ensemble: # If not homogeneous, we have to treat each z feature vector independently self.gp_reg_ensemble = gp_regressors self.B_x = [B_x[dim] for dim in gp_regressors.dim_idx] self.B_x = np.squeeze(np.stack(self.B_x, axis=1)) self.B_x = self.B_x[:, np.newaxis] if len(self.B_x.shape) == 1 else self.B_x self.B_z = gp_regressors.B_z elif mlp_regressor is not None: self.gp_reg_ensemble = None self.mlp_regressor = mlp_regressor self.B_x = np.hstack(list(B_x.values())) self.mlp_conf = mlp_conf else: self.gp_reg_ensemble = None # Declare model variables for GP prediction (only used in real quadrotor flight with EKF estimator). # Will be used as initial state for GP prediction as Acados parameters. self.gp_p = cs.MX.sym('gp_p', 3) self.gp_q = cs.MX.sym('gp_a', 4) self.gp_v = cs.MX.sym('gp_v', 3) self.gp_r = cs.MX.sym('gp_r', 3) self.gp_x = cs.vertcat(self.gp_p, self.gp_q, self.gp_v, self.gp_r) # The trigger variable is used to tell ACADOS to use the additional GP state estimate in the first optimization # node and the regular integrated state in the rest self.trigger_var = cs.MX.sym('trigger', 1) # Build full model. Will have 13 variables. self.dyn_x contains the symbolic variable that # should be used to evaluate the dynamics function. It corresponds to self.x if there are no GP's, or # self.x_with_gp otherwise acados_models, nominal_with_gp = self.acados_setup_model( self.quad_xdot_nominal(x=self.x, u=self.u)['x_dot'], model_name) # Convert dynamics variables to functions of the state and input vectors self.quad_xdot = {} for dyn_model_idx in nominal_with_gp.keys(): dyn = nominal_with_gp[dyn_model_idx] self.quad_xdot[dyn_model_idx] = cs.Function( 'x_dot', [self.x, self.u], [dyn], ['x', 'u'], ['x_dot'], {'allow_free': True}) # ### Setup and compile Acados OCP solvers ### # self.acados_ocp_solver = {} # Check if GP's have been loaded self.with_gp = self.gp_reg_ensemble is not None self.with_mlp = self.mlp_regressor is not None # Add one more weight to the rotation (use quaternion norm weighting in acados) q_diagonal = np.concatenate((q_cost[:3], np.mean(q_cost[3:6])[np.newaxis], q_cost[3:])) if q_mask is not None: q_mask = np.concatenate((q_mask[:3], np.zeros(1), q_mask[3:])) q_diagonal *= q_mask # Ensure current working directory is current folder os.chdir(os.path.dirname(os.path.realpath(__file__))) self.acados_models_dir = '../../acados_models' safe_mkdir_recursive(os.path.join(os.getcwd(), self.acados_models_dir)) for key, key_model in zip(acados_models.keys(), acados_models.values()): nx = key_model.x.size()[0] nu = key_model.u.size()[0] ny = nx + nu n_param = key_model.p.size()[0] if isinstance(key_model.p, cs.MX) else 0 acados_source_path = os.environ['ACADOS_SOURCE_DIR'] sys.path.insert(0, '../common') # Create OCP object to formulate the optimization ocp = AcadosOcp() ocp.acados_include_path = acados_source_path + '/include' ocp.acados_lib_path = acados_source_path + '/lib' ocp.model = key_model ocp.dims.N = self.N ocp.solver_options.tf = t_horizon # Initialize parameters ocp.dims.np = n_param ocp.parameter_values = np.zeros(n_param) ocp.cost.cost_type = 'LINEAR_LS' ocp.cost.cost_type_e = 'LINEAR_LS' ocp.cost.W = np.diag(np.concatenate((q_diagonal, r_cost))) ocp.cost.W_e = np.diag(q_diagonal) terminal_cost = 0 if solver_options is None or not solver_options["terminal_cost"] else 1 ocp.cost.W_e *= terminal_cost ocp.cost.Vx = np.zeros((ny, nx)) ocp.cost.Vx[:nx, :nx] = np.eye(nx) ocp.cost.Vu = np.zeros((ny, nu)) ocp.cost.Vu[-4:, -4:] = np.eye(nu) ocp.cost.Vx_e = np.eye(nx) # Initial reference trajectory (will be overwritten) x_ref = np.zeros(nx) ocp.cost.yref = np.concatenate((x_ref, np.array([0.0, 0.0, 0.0, 0.0]))) ocp.cost.yref_e = x_ref # Initial state (will be overwritten) ocp.constraints.x0 = x_ref # Set constraints ocp.constraints.lbu = np.array([self.min_u] * 4) ocp.constraints.ubu = np.array([self.max_u] * 4) ocp.constraints.idxbu = np.array([0, 1, 2, 3]) # Solver options ocp.solver_options.qp_solver = 'FULL_CONDENSING_HPIPM' ocp.solver_options.hessian_approx = 'GAUSS_NEWTON' ocp.solver_options.integrator_type = 'ERK' ocp.solver_options.print_level = 0 ocp.solver_options.nlp_solver_type = 'SQP_RTI' if solver_options is None else solver_options["solver_type"] # Compile acados OCP solver if necessary json_file = os.path.join(self.acados_models_dir, key_model.name + '_acados_ocp.json') self.acados_ocp_solver[key] = AcadosOcpSolver(ocp, json_file=json_file) def clear_acados_model(self): """ Removes previous stored acados models to avoid name conflicts. """ json_file = os.path.join(self.acados_models_dir, 'acados_ocp.json') if os.path.exists(json_file): os.remove(os.path.join(os.getcwd(), json_file)) compiled_model_dir = os.path.join(os.getcwd(), 'c_generated_code') if os.path.exists(compiled_model_dir): shutil.rmtree(compiled_model_dir) def add_missing_states(self, gp_outs): """ Adds 0's to the GP predictions in case the GP's only regressed a subset of the three velocity states. :param gp_outs: dictionary of GP predictions :type gp_outs: dict :return: The same dictionary as the input but all entries will be 3-dimensional, with 0's added to the right dimensions if they were missing. """ # Dims predicted by the GP. A subset of real_pred_dims pred_dims = self.gp_reg_ensemble.dim_idx # All the dims that should be filled in real_pred_dims = [7, 8, 9] for i, dim in enumerate(real_pred_dims): if dim not in pred_dims: for k, v in zip(gp_outs.keys(), gp_outs.values()): if (isinstance(v, cs.MX) or isinstance(v, cs.DM)) and k == 'pred': gp_outs[k] = cs.vertcat(v[:i], cs.MX.zeros(1), v[i:]) return gp_outs def remove_extra_states(self, vec): # Dims predicted by the GP. A subset of real_pred_dims pred_dims = self.gp_reg_ensemble.dim_idx # All the dims that should be filled in real_pred_dims = [7, 8, 9] for i, dim in enumerate(real_pred_dims): if dim not in pred_dims: vec = cs.vertcat(vec[:i], vec[i + 1:]) return vec def acados_setup_model(self, nominal, model_name): """ Builds an Acados symbolic models using CasADi expressions. :param model_name: name for the acados model. Must be different from previously used names or there may be problems loading the right model. :param nominal: CasADi symbolic nominal model of the quadrotor: f(self.x, self.u) = x_dot, dimensions 13x1. :return: Returns a total of three outputs, where m is the number of GP's in the GP ensemble, or 1 if no GP: - A dictionary of m AcadosModel of the GP-augmented quadrotor - A dictionary of m CasADi symbolic nominal dynamics equations with GP mean value augmentations (if with GP) :rtype: dict, dict, cs.MX """ def fill_in_acados_model(x, u, p, dynamics, name): x_dot = cs.MX.sym('x_dot', dynamics.shape) f_impl = x_dot - dynamics # Dynamics model model = AcadosModel() model.f_expl_expr = dynamics model.f_impl_expr = f_impl model.x = x model.xdot = x_dot model.u = u model.p = p model.name = name return model acados_models = {} dynamics_equations = {} # Run GP inference if GP's available if self.gp_reg_ensemble is not None: # Feature vector are the elements of x and u determined by the selection matrix B_z. The trigger var is used # to select the gp-specific state estimate in the first optimization node, and the regular integrated state # in the rest. The computing of the features z is done within the GPEnsemble. gp_x = self.gp_x * self.trigger_var + self.x * (1 - self.trigger_var) # Transform velocity to body frame v_b = v_dot_q(gp_x[7:10], quaternion_inverse(gp_x[3:7])) gp_x = cs.vertcat(gp_x[:7], v_b, gp_x[10:]) gp_u = self.u gp_dims = self.gp_reg_ensemble.dim_idx # Get number of models in GP for i in range(self.gp_reg_ensemble.n_models): # Evaluate cluster of the GP ensemble cluster_id = {k: [v] for k, v in zip(gp_dims, i * np.ones_like(gp_dims, dtype=int))} outs = self.gp_reg_ensemble.predict(gp_x, gp_u, return_cov=False, gp_idx=cluster_id, return_z=False) # Unpack prediction outputs. Transform back to world reference frame outs = self.add_missing_states(outs) gp_means = v_dot_q(outs["pred"], gp_x[3:7]) gp_means = self.remove_extra_states(gp_means) # Add GP mean prediction dynamics_equations[i] = nominal + cs.mtimes(self.B_x, gp_means) x_ = self.x dynamics_ = dynamics_equations[i] # Add again the gp augmented dynamics for the GP state dynamics_ = cs.vertcat(dynamics_) dynamics_equations[i] = cs.vertcat(dynamics_equations[i]) i_name = model_name + "_domain_" + str(i) params = cs.vertcat(self.gp_x, self.trigger_var) acados_models[i] = fill_in_acados_model(x=x_, u=self.u, p=params, dynamics=dynamics_, name=i_name) elif self.mlp_regressor is not None: state = self.gp_x * self.trigger_var + self.x * (1 - self.trigger_var) # Transform velocity to body frame v_b = v_dot_q(state[7:10], quaternion_inverse(state[3:7])) state = cs.vertcat(state[:7], v_b, state[10:]) mlp_in = v_b if self.mlp_conf['torque_output']: mlp_in = cs.vertcat(mlp_in, state[10:]) if self.mlp_conf['u_inp']: mlp_in = cs.vertcat(mlp_in, self.u) if self.mlp_conf['ground_map_input']: map_conf = GroundEffectMapConfig map = GroundMapWithBox(np.array(map_conf.box_min), np.array(map_conf.box_max), map_conf.box_height, horizon=map_conf.horizon, resolution=map_conf.resolution) self._map_res = map_conf.resolution self._static_ground_map, self._org_to_map_org = map.at(np.array(map_conf.origin)) ground_map_dx = cs.MX(self._static_ground_map) idx = cs.DM(np.arange(0, 3, 1)) x, y, z = state[0], state[1], state[2] orientation = state[3:7] x_idxs = cs.floor((x - self._org_to_map_org[0]) / map_conf.resolution) + idx - 1 y_idxs = cs.floor((y - self._org_to_map_org[1]) / map_conf.resolution) + idx - 1 ground_patch = ground_map_dx[x_idxs, y_idxs] relative_ground_patch = z - ground_patch relative_ground_patch = 4 * (cs.fmax(cs.fmin(relative_ground_patch, 0.5), 0.0) - 0.25) ground_effect_in = cs.vertcat(cs.reshape(relative_ground_patch, 9, 1), orientation*0) mlp_in = cs.vertcat(mlp_in, ground_effect_in) if not self.mlp_conf['approximated']: outs = self.mlp_regressor(mlp_in) else: outs = self.mlp_regressor.approx(mlp_in, order=self.mlp_conf['approx_order'], parallel=False) if self.mlp_conf['torque_output']: outs_force = outs[:3] outs_torque = outs[3:] mlp_means = cs.vertcat(v_dot_q(outs_force, state[3:7]), outs_torque) else: # Unpack prediction outputs. Transform back to world reference frame mlp_means = v_dot_q(outs, state[3:7]) # Add GP mean prediction dynamics_equations[0] = nominal + cs.mtimes(self.B_x, mlp_means) x_ = self.x dynamics_ = dynamics_equations[0] # Add again the gp augmented dynamics for the GP state dynamics_ = cs.vertcat(dynamics_) dynamics_equations[0] = cs.vertcat(dynamics_equations[0]) i_name = model_name + "_domain_" + str(0) if not self.mlp_conf['approximated']: params = cs.vertcat(self.gp_x, self.trigger_var) else: params = cs.vertcat(self.gp_x, self.trigger_var, self.mlp_regressor.sym_approx_params(order=self.mlp_conf['approx_order'], flat=True)) acados_models[0] = fill_in_acados_model(x=x_, u=self.u, p=params, dynamics=dynamics_, name=i_name) else: # No available GP so return nominal dynamics dynamics_equations[0] = nominal x_ = self.x dynamics_ = nominal acados_models[0] = fill_in_acados_model(x=x_, u=self.u, p=[], dynamics=dynamics_, name=model_name) return acados_models, dynamics_equations def quad_dynamics(self, rdrv_d): """ Symbolic dynamics of the 3D quadrotor model. The state consists on: [p_xyz, a_wxyz, v_xyz, r_xyz]^T, where p stands for position, a for angle (in quaternion form), v for velocity and r for body rate. The input of the system is: [u_1, u_2, u_3, u_4], i.e. the activation of the four thrusters. :param rdrv_d: a 3x3 diagonal matrix containing the D matrix coefficients for a linear drag model as proposed by Faessler et al. :return: CasADi function that computes the analytical differential state dynamics of the quadrotor model. Inputs: 'x' state of quadrotor (6x1) and 'u' control input (2x1). Output: differential state vector 'x_dot' (6x1) """ x_dot = cs.vertcat(self.p_dynamics(), self.q_dynamics(), self.v_dynamics(rdrv_d), self.w_dynamics()) return cs.Function('x_dot', [self.x[:13], self.u], [x_dot], ['x', 'u'], ['x_dot']) def p_dynamics(self): return self.v def q_dynamics(self): return 1 / 2 * cs.mtimes(skew_symmetric(self.r), self.q) def v_dynamics(self, rdrv_d): """ :param rdrv_d: a 3x3 diagonal matrix containing the D matrix coefficients for a linear drag model as proposed by Faessler et al. None, if no linear compensation is to be used. """ f_thrust = self.u * self.quad.max_thrust g = cs.vertcat(0.0, 0.0, 9.81) a_thrust = cs.vertcat(0.0, 0.0, f_thrust[0] + f_thrust[1] + f_thrust[2] + f_thrust[3]) / self.quad.mass v_dynamics = v_dot_q(a_thrust, self.q) - g if rdrv_d is not None: # Velocity in body frame: v_b = v_dot_q(self.v, quaternion_inverse(self.q)) rdrv_drag = v_dot_q(cs.mtimes(rdrv_d, v_b), self.q) v_dynamics += rdrv_drag return v_dynamics def w_dynamics(self): f_thrust = self.u * self.quad.max_thrust y_f = cs.MX(self.quad.y_f) x_f = cs.MX(self.quad.x_f) c_f = cs.MX(self.quad.z_l_tau) return cs.vertcat( (cs.mtimes(f_thrust.T, y_f) + (self.quad.J[1] - self.quad.J[2]) * self.r[1] * self.r[2]) / self.quad.J[0], (-cs.mtimes(f_thrust.T, x_f) + (self.quad.J[2] - self.quad.J[0]) * self.r[2] * self.r[0]) / self.quad.J[1], (cs.mtimes(f_thrust.T, c_f) + (self.quad.J[0] - self.quad.J[1]) * self.r[0] * self.r[1]) / self.quad.J[2]) def linearized_quad_dynamics(self): """ Jacobian J matrix of the linearized dynamics specified in the function quad_dynamics. J[i, j] corresponds to the partial derivative of f_i(x) wrt x(j). :return: a CasADi symbolic function that calculates the 13 x 13 Jacobian matrix of the linearized simplified quadrotor dynamics """ jac = cs.MX(self.state_dim, self.state_dim) # Position derivatives jac[0:3, 7:10] = cs.diag(cs.MX.ones(3)) # Angle derivatives jac[3:7, 3:7] = skew_symmetric(self.r) / 2 jac[3, 10:] = 1 / 2 * cs.horzcat(-self.q[1], -self.q[2], -self.q[3]) jac[4, 10:] = 1 / 2 * cs.horzcat(self.q[0], -self.q[3], self.q[2]) jac[5, 10:] = 1 / 2 * cs.horzcat(self.q[3], self.q[0], -self.q[1]) jac[6, 10:] = 1 / 2 * cs.horzcat(-self.q[2], self.q[1], self.q[0]) # Velocity derivatives a_u = (self.u[0] + self.u[1] + self.u[2] + self.u[3]) * self.quad.max_thrust / self.quad.mass jac[7, 3:7] = 2 * cs.horzcat(a_u * self.q[2], a_u * self.q[3], a_u * self.q[0], a_u * self.q[1]) jac[8, 3:7] = 2 * cs.horzcat(-a_u * self.q[1], -a_u * self.q[0], a_u * self.q[3], a_u * self.q[2]) jac[9, 3:7] = 2 * cs.horzcat(0, -2 * a_u * self.q[1], -2 * a_u * self.q[1], 0) # Rate derivatives jac[10, 10:] = (self.quad.J[1] - self.quad.J[2]) / self.quad.J[0] * cs.horzcat(0, self.r[2], self.r[1]) jac[11, 10:] = (self.quad.J[2] - self.quad.J[0]) / self.quad.J[1] * cs.horzcat(self.r[2], 0, self.r[0]) jac[12, 10:] = (self.quad.J[0] - self.quad.J[1]) / self.quad.J[2] * cs.horzcat(self.r[1], self.r[0], 0) return cs.Function('J', [self.x, self.u], [jac]) def set_reference_state(self, x_target=None, u_target=None): """ Sets the target state and pre-computes the integration dynamics with cost equations :param x_target: 13-dimensional target state (p_xyz, a_wxyz, v_xyz, r_xyz) :param u_target: 4-dimensional target control input vector (u_1, u_2, u_3, u_4) """ if x_target is None: x_target = [[0, 0, 0], [1, 0, 0, 0], [0, 0, 0], [0, 0, 0]] if u_target is None: u_target = [0, 0, 0, 0] # Set new target state self.target = copy(x_target) self.u_target = copy(u_target) ref = np.concatenate([x_target[i] for i in range(4)]) # Transform velocity to body frame v_b = v_dot_q(ref[7:10], quaternion_inverse(ref[3:7])) ref = np.concatenate((ref[:7], v_b, ref[10:])) # Determine which dynamics model to use based on the GP optimal input feature region. Should get one for each # output dimension of the GP if self.gp_reg_ensemble is not None: gp_ind = self.gp_reg_ensemble.select_gp(dim=None, x=ref, u=u_target) else: gp_ind = 0 ref = np.concatenate((ref, u_target)) for j in range(self.N): self.acados_ocp_solver[gp_ind].set(j, "yref", ref) self.acados_ocp_solver[gp_ind].set(self.N, "yref", ref[:-4]) return gp_ind def set_reference_trajectory(self, x_target, u_target): """ Sets the reference trajectory and pre-computes the cost equations for each point in the reference sequence. :param x_target: Nx13-dimensional reference trajectory (p_xyz, angle_wxyz, v_xyz, rate_xyz). It is passed in the form of a 4-length list, where the first element is a Nx3 numpy array referring to the position targets, the second is a Nx4 array referring to the quaternion, two more Nx3 arrays for the velocity and body rate targets. :param u_target: Nx4-dimensional target control input vector (u1, u2, u3, u4) """ if u_target is not None: assert x_target[0].shape[0] == (u_target.shape[0] + 1) or x_target[0].shape[0] == u_target.shape[0] # If not enough states in target sequence, append last state until required length is met while x_target[0].shape[0] < self.N + 1: x_target = [np.concatenate((x, np.expand_dims(x[-1, :], 0)), 0) for x in x_target] if u_target is not None: u_target = np.concatenate((u_target, np.expand_dims(u_target[-1, :], 0)), 0) stacked_x_target = np.concatenate([x for x in x_target], 1) # Transform velocity to body frame x_mean = stacked_x_target[int(self.N / 2)] v_b = v_dot_q(x_mean[7:10], quaternion_inverse(x_mean[3:7])) x_target_mean = np.concatenate((x_mean[:7], v_b, x_mean[10:])) # Determine which dynamics model to use based on the GP optimal input feature region if self.gp_reg_ensemble is not None: gp_ind = self.gp_reg_ensemble.select_gp(dim=None, x=x_target_mean, u=u_target[int(self.N / 2)]) else: gp_ind = 0 self.target = copy(x_target) self.u_target = copy(u_target) for j in range(self.N): ref = stacked_x_target[j, :] ref = np.concatenate((ref, u_target[j, :])) self.acados_ocp_solver[gp_ind].set(j, "yref", ref) # the last MPC node has only a state reference but no input reference self.acados_ocp_solver[gp_ind].set(self.N, "yref", stacked_x_target[self.N, :]) return gp_ind def discretize_f_and_q(self, t_horizon, n, m=1, i=0, use_gp=True, use_model=0): """ Discretize the model dynamics and the pre-computed cost function if available. :param t_horizon: time horizon in seconds :param n: number of control steps until time horizon :param m: number of integration steps per control step :param i: Only used for trajectory tracking. Index of cost function to use. :param use_gp: Whether to use the dynamics with the GP correction or not. :param use_model: integer, select which model to use from the available options. :return: the symbolic, discretized dynamics. The inputs of the symbolic function are x0 (the initial state) and p, the control input vector. The outputs are xf (the updated state) and qf. qf is the corresponding cost function of the integration, which is calculated from the pre-computed discrete-time model dynamics (self.L) """ dynamics = self.quad_xdot[use_model] if use_gp else self.quad_xdot_nominal # Call with self.x_with_gp even if use_gp=False return discretize_dynamics_and_cost(t_horizon, n, m, self.x, self.u, dynamics, self.L, i) def run_optimization(self, initial_state=None, use_model=0, return_x=False, gp_regression_state=None): """ Optimizes a trajectory to reach the pre-set target state, starting from the input initial state, that minimizes the quadratic cost function and respects the constraints of the system :param initial_state: 13-element list of the initial state. If None, 0 state will be used :param use_model: integer, select which model to use from the available options. :param return_x: bool, whether to also return the optimized sequence of states alongside with the controls. :param gp_regression_state: 13-element list of state for GP prediction. If None, initial_state will be used. :return: optimized control input sequence (flattened) """ if initial_state is None: initial_state = [0, 0, 0] + [1, 0, 0, 0] + [0, 0, 0] + [0, 0, 0] # Set initial state. Add gp state if needed x_init = initial_state x_init = np.stack(x_init) # Set initial condition, equality constraint self.acados_ocp_solver[use_model].set(0, 'lbx', x_init) self.acados_ocp_solver[use_model].set(0, 'ubx', x_init) # Set parameters if self.with_gp: gp_state = gp_regression_state if gp_regression_state is not None else initial_state self.acados_ocp_solver[use_model].set(0, 'p', np.array(gp_state + [1])) for j in range(1, self.N): self.acados_ocp_solver[use_model].set(j, 'p', np.array([0.0] * (len(gp_state) + 1))) if self.with_mlp: if self.x_opt_acados is None: if isinstance(self.target[0], list): self.x_opt_acados = np.expand_dims( np.concatenate([self.target[i] for i in range(len(self.target))]), 0) self.x_opt_acados = self.x_opt_acados.repeat(self.N, 0) else: self.x_opt_acados = np.hstack(self.target) if self.w_opt_acados is None: if len(self.u_target.shape) == 1: self.w_opt_acados = self.u_target[np.newaxis] self.w_opt_acados = self.w_opt_acados.repeat(self.N, 0) else: self.w_opt_acados = np.hstack(self.u_target) gp_state = gp_regression_state if gp_regression_state is not None else initial_state if not self.mlp_conf['approximated']: self.acados_ocp_solver[use_model].set(0, 'p', np.hstack([np.array(gp_state + [1])])) for j in range(1, self.N): self.acados_ocp_solver[use_model].set(j, 'p', np.hstack([np.array([0.0] * (len(gp_state) + 1))])) else: state = np.vstack([np.array([initial_state]), self.x_opt_acados[1:]]) a_list = [] for i in range(state.shape[0]): a_list.append(v_dot_q(np.array(state[i, 7:10]), quaternion_inverse(np.array(state[i, 3:7])))) a = np.array(a_list)[:self.N] if self.mlp_conf['torque_output']: a = np.concatenate([a, state[:self.N, 10:]], axis=-1) if self.mlp_conf['u_inp']: a = np.concatenate([a, self.w_opt_acados], axis=-1) if self.mlp_conf['ground_map_input']: ground_maps = [] for i in range(state.shape[0]): pos = state[i][:3] x_idxs = np.floor((pos[0] - self._org_to_map_org[0]) / self._map_res).astype(int) - 1 y_idxs = np.floor((pos[1] - self._org_to_map_org[1]) / self._map_res).astype(int) - 1 ground_patch = self._static_ground_map[x_idxs:x_idxs + 3, y_idxs:y_idxs + 3] relative_ground_patch = 4 * (np.clip(pos[2] - ground_patch, 0, 0.5) - 0.25) flatten_relative_ground_patch = relative_ground_patch.flatten(order='F') ground_effect_in = np.hstack([flatten_relative_ground_patch, flatten_relative_ground_patch[..., :4] * 0]) ground_maps.append(ground_effect_in) a = np.concatenate([a, np.array(ground_maps)[:self.N]], axis=-1) mlp_params = self.mlp_regressor.approx_params(a, order=self.mlp_conf['approx_order'], flat=True) mlp_params = np.vstack([mlp_params, mlp_params[[-1]]]) self.acados_ocp_solver[use_model].set(0, 'p', np.hstack([np.array(gp_state + [1]), mlp_params[0]])) for j in range(1, self.N): self.acados_ocp_solver[use_model].set(j, 'p', np.hstack([np.array([0.0] * (len(gp_state) + 1)), mlp_params[j]])) # Solve OCP self.acados_ocp_solver[use_model].solve() # Get u w_opt_acados = np.ndarray((self.N, 4)) x_opt_acados = np.ndarray((self.N + 1, len(x_init))) x_opt_acados[0, :] = self.acados_ocp_solver[use_model].get(0, "x") for i in range(self.N): w_opt_acados[i, :] = self.acados_ocp_solver[use_model].get(i, "u") x_opt_acados[i + 1, :] = self.acados_ocp_solver[use_model].get(i + 1, "x") self.x_opt_acados = x_opt_acados.copy() self.w_opt_acados = w_opt_acados.copy() w_opt_acados = np.reshape(w_opt_acados, (-1)) return w_opt_acados if not return_x else (w_opt_acados, x_opt_acados) ================================================ FILE: ros_dd_mpc/src/utils/__init__.py ================================================ ================================================ FILE: ros_dd_mpc/src/utils/animator.py ================================================ """ Class for generating a comprehensive post-processed visualization of experimental flight results. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import matplotlib.pyplot as plt import numpy as np from mpl_toolkits.mplot3d import axes3d import matplotlib.animation as animation from config.configuration_parameters import DirectoryConfig class Dynamic3DTrajectory: def __init__(self, pos_data, vel_data, pos_ref, vel_ref, t_vec_ref, legends, sparsing_factor=1): # Add reference to data so it also moves pos_data = pos_ref + pos_data self.data = np.array(pos_data) self.reference = pos_ref[0] self.data_len = len(pos_data[0]) self.n_lines = len(pos_data) self.t_vec = t_vec_ref - t_vec_ref[0] self.legends = legends vel_data = vel_ref + vel_data self.vel_data = np.array(vel_data) if sparsing_factor == 0: sparse_data = np.arange(0, self.data_len) else: sparse_data = np.arange(0, self.data_len, sparsing_factor) self.sparsed_data = self.data[:, sparse_data, :] self.vel_data = self.vel_data[:, sparse_data, :] self.t_vec = self.t_vec[sparse_data] self.max_buffer_size = 50 self.colors = ["tab:blue", "tab:orange", "tab:green", "tab:red"] self.data_len = len(sparse_data) x_data = np.concatenate(tuple([dat[sparse_data, 0] for dat in pos_data])) y_data = np.concatenate(tuple([dat[sparse_data, 1] for dat in pos_data])) z_data = np.concatenate(tuple([dat[sparse_data, 2] for dat in pos_data])) self.max_x = np.max(x_data) self.min_x = np.min(x_data) self.max_y = np.max(y_data) self.min_y = np.min(y_data) self.max_z = np.max(z_data) self.min_z = np.min(z_data) # Make dimensions more similar and add a bit of padding range_x = self.max_x - self.min_x range_y = self.max_y - self.min_y range_z = self.max_z - self.min_z max_range = max(range_x, range_y, range_z) self.max_x = self.max_x + 0.25 * (max_range - range_x) + (self.max_x - self.min_x) * 0.2 self.min_x = self.min_x - 0.25 * (max_range - range_x) - (self.max_x - self.min_x) * 0.2 self.max_y = self.max_y + 0.25 * (max_range - range_y) + (self.max_y - self.min_y) * 0.2 self.min_y = self.min_y - 0.25 * (max_range - range_y) - (self.max_y - self.min_y) * 0.2 self.max_z = self.max_z + 0.25 * (max_range - range_z) + (self.max_z - self.min_z) * 0.2 self.min_z = self.min_z - 0.25 * (max_range - range_z) - (self.max_z - self.min_z) * 0.2 self.figure = None self.ax = None self.pos_err_ax = None self.speed_ax = None self.top_down_ax = None self.x_time_ax = None self.y_time_ax = None self.lines = None self.n_3d_lines = None self.vel_bars_0_idx = None self.pos_err_0_idx = None self.top_down_0_idx = None self.x_time_0_idx = None self.y_time_0_idx = None def on_launch(self): self.figure = plt.figure(figsize=(14, 10)) self.ax = axes3d.Axes3D(self.figure, rect=(-0.02, 0.3, 0.65, 0.7)) self.ax.set_zlim3d([self.min_z, self.max_z]) self.ax.set_ylim3d([self.min_y, self.max_y]) self.ax.set_xlim3d([self.min_x, self.max_x]) self.ax.set_xlabel(r'$p_x\: [m]$') self.ax.set_ylabel(r'$p_y\: [m]$') self.ax.set_zlabel(r'$p_z\: [m]$') self.ax.plot(self.reference[:, 0], self.reference[:, 1], self.reference[:, 2], '-', alpha=0.2) self.lines = sum([self.ax.plot([], [], [], '-', c=self.colors[i], label=self.legends[i]) for i in range(self.n_lines)], []) projection_lines = sum([self.ax.plot([], [], [], '-', c=self.colors[i], alpha=0.2) for i in range(self.n_lines)], []) projection_lines += sum([self.ax.plot([], [], [], '-', c=self.colors[i], alpha=0.2) for i in range(self.n_lines)], []) projection_lines += sum([self.ax.plot([], [], [], '-', c=self.colors[i], alpha=0.2) for i in range(self.n_lines)], []) self.lines += projection_lines pos_balls = sum([self.ax.plot([], [], [], 'o', c=self.colors[i]) for i in range(self.n_lines)], []) self.lines += pos_balls self.n_3d_lines = len(self.lines) self.ax.legend() self.ax.set_title('3D Visualization') self.vel_bars_0_idx = len(self.lines) bar_height = 0.1 self.speed_ax = plt.axes((0.65, 0.8, 0.3, 0.05)) vel_bar = [self.speed_ax.barh([0], [0], color=self.colors[0], height=bar_height)[0]] self.speed_ax.set_xlim([0, np.max(np.sqrt(np.sum(self.vel_data[0] ** 2, -1))) * 1.05]) self.speed_ax.set_ylim([-bar_height, bar_height * 1.05]) self.speed_ax.set_xlabel(r'$\|\mathbf{v}\|\:[m/s]$') self.speed_ax.set_yticks([0]) self.speed_ax.set_yticklabels([self.legends[0]]) self.speed_ax.grid() self.speed_ax.set_title('Quadrotor speed') self.lines += vel_bar self.pos_err_ax = plt.axes((0.65, 0.6, 0.3, 0.1)) self.pos_err_ax.grid() pos_err = [] p_errors = np.stack([self.data[0] - self.data[i+1] for i in range(self.n_lines - 1)]) pos_err += [self.pos_err_ax.barh([(i - 1) * bar_height * 1.05], [i], color=self.colors[i], height=bar_height)[0] for i in range(1, self.n_lines)] self.pos_err_ax.set_xlim([np.min(p_errors) * 1.05, np.max(p_errors) * 1.05]) self.pos_err_ax.set_ylim([-bar_height, ((self.n_lines - 1) * 1.05) * bar_height]) self.pos_err_ax.set_title('XY position error') self.pos_err_ax.set_xlabel(r'$\|\mathbf{p}^* - \mathbf{p}\|\: [m]$') self.pos_err_ax.set_yticks([i * bar_height * 1.05 for i in range(self.n_lines - 1)]) self.pos_err_ax.set_yticklabels(self.legends[1:]) self.pos_err_0_idx = len(self.lines) self.lines += pos_err self.top_down_ax = plt.axes((0.65, 0.08, 0.3, 0.4)) top_down_lines = sum([self.top_down_ax.plot([], [], color=self.colors[i], label=self.legends[i]) for i in range(self.n_lines)], []) top_down_lines += sum([self.top_down_ax.plot([], [], 'o', color=self.colors[i]) for i in range(self.n_lines)], []) self.top_down_ax.plot(self.reference[:, 0], self.reference[:, 1], '-', alpha=0.2) self.top_down_0_idx = len(self.lines) self.top_down_ax.grid() self.top_down_ax.set_xlim([self.min_x, self.max_x]) self.top_down_ax.set_ylim([self.min_y, self.max_y]) self.top_down_ax.set_title('Top down view') self.top_down_ax.set_xlabel(r'$p_x\:[m]$') self.top_down_ax.set_ylabel(r'$p_y\:[m]$') self.top_down_ax.legend() self.lines += top_down_lines self.x_time_ax = plt.axes((0.05, 0.08, 0.26, 0.2)) self.x_time_ax.axhline(y=0, linestyle='--', alpha=0.5) x_time_lines = sum([self.x_time_ax.plot([], [], color=self.colors[i], label=self.legends[i]) for i in range(1, self.n_lines)], []) self.x_time_0_idx = len(self.lines) self.lines += x_time_lines self.x_time_ax.set_xlim([self.t_vec[0], self.t_vec[-1]]) self.x_time_ax.set_ylim([0, np.max(np.abs(p_errors[:, :, 0])) * 1.05]) self.x_time_ax.grid() self.x_time_ax.set_title(r'$p_x$ error') self.x_time_ax.set_ylabel('pos [m]') self.x_time_ax.legend(loc='upper right') self.y_time_ax = plt.axes((0.34, 0.08, 0.26, 0.2)) self.y_time_ax.axhline(y=0, linestyle='--', alpha=0.5) y_time_lines = sum([self.y_time_ax.plot([], [], color=self.colors[i], label=self.legends[i]) for i in range(1, self.n_lines)], []) self.y_time_0_idx = len(self.lines) self.lines += y_time_lines self.y_time_ax.set_xlim([self.t_vec[0], self.t_vec[-1]]) self.y_time_ax.set_ylim([0, np.max(np.abs(p_errors[:, :, 1])) * 1.05]) self.y_time_ax.set_xlabel('Time [s]') self.y_time_ax.set_title(r'$p_y$ error') self.y_time_ax.grid() self.y_time_ax.legend(loc='upper right') def on_init(self): for i in range(self.n_3d_lines): self.lines[i].set_data([], []) self.lines[i].set_3d_properties([]) self.lines[self.vel_bars_0_idx].set_width(0) for i in range(self.n_lines - 1): self.lines[self.pos_err_0_idx + i].set_width(0) for i in range(self.n_lines * 2): self.lines[self.top_down_0_idx + i].set_data([], []) for i in range(self.n_lines - 1): self.lines[self.x_time_0_idx + i].set_data([], []) self.lines[self.y_time_0_idx + i].set_data([], []) return self.lines def animate(self, i): #i = (2 * i) % self.data.shape[1] for j, (line, xi) in enumerate(zip(self.lines[:self.n_lines], self.sparsed_data)): x, y, z = xi[:i].T if len(x) > self.max_buffer_size: x = x[len(x) - self.max_buffer_size:] y = y[len(y) - self.max_buffer_size:] z = z[len(z) - self.max_buffer_size:] line.set_data(x, y) line.set_3d_properties(z) self.lines[j + self.n_lines].set_data(x, self.max_y) self.lines[j + self.n_lines].set_3d_properties(z) self.lines[j + 2 * self.n_lines].set_data(np.ones(len(y)) * self.min_x, y) self.lines[j + 2 * self.n_lines].set_3d_properties(z) self.lines[j + 3 * self.n_lines].set_data(x, y) self.lines[j + 3 * self.n_lines].set_3d_properties(self.min_z) if len(x) > 0: self.lines[j + 4 * self.n_lines].set_data([x[-1]], [y[-1]]) self.lines[j + 4 * self.n_lines].set_3d_properties([z[-1]]) self.lines[self.top_down_0_idx + j + self.n_lines].set_data(x[-1], y[-1]) self.lines[self.top_down_0_idx + j].set_data(x, y) self.lines[self.vel_bars_0_idx].set_width(np.sqrt(np.sum(self.vel_data[0, i, :] ** 2))) for j in range(self.n_lines - 1): self.lines[self.pos_err_0_idx + j].set_width( np.sqrt(np.sum((self.data[0][i, :2] - self.data[j+1][i, :2]) ** 2))) for j in reversed(range(self.n_lines - 1)): self.lines[self.x_time_0_idx + j].set_data(self.t_vec[:i], np.abs(self.data[0, :i, 0] - self.data[j + 1, :i, 0])) self.lines[self.y_time_0_idx + j].set_data(self.t_vec[:i], np.abs(self.data[0, :i, 1] - self.data[j + 1, :i, 1])) return self.lines def __call__(self, save=False): self.on_launch() ani = animation.FuncAnimation(self.figure, self.animate, init_func=self.on_init, frames=self.data_len, interval=20, blit=True, repeat=False) if save: results_dir = DirectoryConfig.RESULTS_DIR plt.rcParams['animation.ffmpeg_path'] = '/usr/local/bin/ffmpeg' writer = animation.FFMpegWriter(fps=50) ani.save(results_dir + '/animation.mp4', writer=writer) else: plt.show() ================================================ FILE: ros_dd_mpc/src/utils/ground_map.py ================================================ import numpy as np class GroundMap: def __init__(self, horizon=((-1, 1), (-1, 1)), resolution=0.1): assert ((horizon[0][1] - horizon[0][0]) / resolution).is_integer() assert ((horizon[1][1] - horizon[1][0]) / resolution).is_integer() self._horizon = horizon self._resolution = resolution def at(self, center: np.array): map_center_to_map_origin = np.array([self._horizon[0][0], self._horizon[1][0]]) map_origin = center + map_center_to_map_origin return self.draw(map_origin - 0.5 * self._resolution, self.empty_map), map_origin def draw(self, pos, map): raise NotImplementedError @property def empty_map(self): x_len = int((self._horizon[0][1] - self._horizon[0][0]) / self._resolution) + 1 y_len = int((self._horizon[1][1] - self._horizon[1][0]) / self._resolution) + 1 empty_map = np.zeros((x_len, y_len)) return empty_map class GroundMapWithBox(GroundMap): def __init__(self, box_min, box_max, height, *args, **kwargs): super().__init__(*args, **kwargs) self._box_min = box_min self._box_max = box_max self._height = height def draw(self, pos, map): x_from = max(int((self._box_min[0] - pos[0]) // self._resolution), 0) x_to = min(int((self._box_max[0] - pos[0]) // self._resolution), map.shape[0] - 1) y_from = max(int((self._box_min[1] - pos[1]) // self._resolution), 0) y_to = min(int((self._box_max[1] - pos[1]) // self._resolution), map.shape[1] - 1) map[x_from:x_to, y_from:y_to] = self._height return map if __name__ == '__main__': horizon = ((-7, 7), (-7, 7)) box_min_ = (-4.25, 9.37) box_max_ = (-2.76, 10.13) box_height_ = 1.0 map = GroundMapWithBox(box_min_, box_max_, box_height_, horizon=horizon) map_at, o = map.at(np.array([3, 9])) print(map_at) m_x = int((1.01 + o[0]) / 0.1) print(m_x) ================================================ FILE: ros_dd_mpc/src/utils/keyframe_3d_gen.py ================================================ """ Generates a set of keypoints to generate a piece-wise polynomial trajectory between each pair. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ from sklearn.gaussian_process import GaussianProcessRegressor from sklearn.gaussian_process.kernels import ExpSineSquared import numpy as np import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D def apply_map_limits(x, y, z, limits): # Find out which axis is the most constrained x_max_range = limits["x"][1] - limits["x"][0] y_max_range = limits["y"][1] - limits["y"][0] z_max_range = limits["z"][1] - limits["z"][0] x_actual_range = np.max(x) - np.min(x) y_actual_range = np.max(y) - np.min(y) z_actual_range = np.max(z) - np.min(z) # One or more of the ranges violates the constraints. if x_actual_range > x_max_range or y_actual_range > y_max_range or z_actual_range > z_max_range: shrink_ratio = max(x_actual_range / x_max_range, y_actual_range / y_max_range, z_actual_range / z_max_range) x = (x - np.mean(x)) / shrink_ratio y = (y - np.mean(y)) / shrink_ratio z = (z - np.mean(z)) / shrink_ratio x += (limits["x"][0] - np.min(x)) y += (limits["y"][0] - np.min(y)) z += (limits["z"][0] - np.min(z)) return x, y, z def center_and_scale(x, y, z, x_max, x_min, y_max, y_min, z_max, z_min): x -= (x_min + (x_max - x_min) / 2) y -= (y_min + (y_max - y_min) / 2) z -= (z_min + (z_max - z_min) / 2) scaling = np.mean([x_max - x_min, y_max - y_min, z_max - z_min]) x = x * 6 / scaling y = y * 6 / scaling z = z * 6 / scaling return x, y, z def random_periodical_trajectory(plot=False, random_state=None, map_limits=None): if random_state is None: random_state = np.random.randint(0, 9999) kernel_z = ExpSineSquared(length_scale=5.5, periodicity=60) kernel_y = ExpSineSquared(length_scale=4.5, periodicity=30) + ExpSineSquared(length_scale=4.0, periodicity=15) kernel_x = ExpSineSquared(length_scale=4.5, periodicity=30) + ExpSineSquared(length_scale=4.5, periodicity=60) gp_x = GaussianProcessRegressor(kernel=kernel_x) gp_y = GaussianProcessRegressor(kernel=kernel_y) gp_z = GaussianProcessRegressor(kernel=kernel_z) # High resolution sampling for track boundaries inputs_x = np.linspace(0, 60, 100) inputs_y = np.linspace(0, 30, 100) inputs_z = np.linspace(0, 60, 100) x_sample_hr = gp_x.sample_y(inputs_x[:, np.newaxis], 1, random_state=random_state) y_sample_hr = gp_y.sample_y(inputs_y[:, np.newaxis], 1, random_state=random_state) z_sample_hr = gp_z.sample_y(inputs_z[:, np.newaxis], 1, random_state=random_state) max_x_coords = np.max(x_sample_hr, 0) max_y_coords = np.max(y_sample_hr, 0) max_z_coords = np.max(z_sample_hr, 0) min_x_coords = np.min(x_sample_hr, 0) min_y_coords = np.min(y_sample_hr, 0) min_z_coords = np.min(z_sample_hr, 0) x_sample_hr, y_sample_hr, z_sample_hr = center_and_scale( x_sample_hr, y_sample_hr, z_sample_hr, max_x_coords, min_x_coords, max_y_coords, min_y_coords, max_z_coords, min_z_coords) # Additional constraint on map limits if map_limits is not None: x_sample_hr, y_sample_hr, z_sample_hr = apply_map_limits(x_sample_hr, y_sample_hr, z_sample_hr, map_limits) # Low resolution for control points lr_ind = np.linspace(0, len(x_sample_hr) - 1, 10, dtype=int) lr_ind[-1] = 0 x_sample_lr = x_sample_hr[lr_ind, :] y_sample_lr = y_sample_hr[lr_ind, :] z_sample_lr = z_sample_hr[lr_ind, :] x_sample_diff = x_sample_hr[lr_ind + 1, :] - x_sample_lr y_sample_diff = y_sample_hr[lr_ind + 1, :] - y_sample_lr z_sample_diff = z_sample_hr[lr_ind + 1, :] - z_sample_lr # Get angles at keypoints a_x = np.arctan2(z_sample_diff, y_sample_diff) * 0 a_y = np.arctan2(x_sample_diff, z_sample_diff) * 0 a_z = (np.arctan2(y_sample_diff, x_sample_diff) - np.pi/4) * 0 if plot: # Plot checking # Build rotation matrices rx = [np.array([[1, 0, 0], [0, np.cos(a), -np.sin(a)], [0, np.sin(a), np.cos(a)]]) for a in a_x[:, 0]] ry = [np.array([[np.cos(a), 0, np.sin(a)], [0, 1, 0], [-np.sin(a), 0, np.cos(a)]]) for a in a_y[:, 0]] rz = [np.array([[np.cos(a), -np.sin(a), 0], [np.sin(a), np.cos(a), 0], [0, 0, 1]]) for a in a_z[:, 0]] main_axes = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) quiver_axes = np.zeros((len(lr_ind), 3, 3)) for i in range(len(lr_ind)): r_mat = rz[i].dot(ry[i].dot(rx[i])) rot_body = r_mat.dot(main_axes) quiver_axes[i, :, :] = rot_body shortest_axis = min(np.max(x_sample_hr) - np.min(x_sample_hr), np.max(y_sample_hr) - np.min(y_sample_hr), np.max(z_sample_hr) - np.min(z_sample_hr)) fig = plt.figure() ax = fig.add_subplot(111, projection='3d') ax.scatter(x_sample_lr, y_sample_lr, z_sample_lr) ax.plot(x_sample_hr[:, 0], y_sample_hr[:, 0], z_sample_hr[:, 0], '-', alpha=0.5) ax.quiver(x_sample_lr[:, 0], y_sample_lr[:, 0], z_sample_lr[:, 0], x_sample_diff[:, 0], y_sample_diff[:, 0], z_sample_diff[:, 0], color='g', length=shortest_axis / 6, normalize=True, label='traj. norm') ax.quiver(x_sample_lr, y_sample_lr, z_sample_lr, quiver_axes[:, 0, :], quiver_axes[:, 1, :], quiver_axes[:, 2, :], color='b', length=shortest_axis / 6, normalize=True, label='quad. att.') ax.tick_params(labelsize=16) ax.legend(fontsize=16) ax.set_xlabel('x [m]', size=16, labelpad=10) ax.set_ylabel('y [m]', size=16, labelpad=10) ax.set_zlabel('z [m]', size=16, labelpad=10) ax.set_xlim([np.min(x_sample_hr), np.max(x_sample_hr)]) ax.set_ylim([np.min(y_sample_hr), np.max(y_sample_hr)]) ax.set_zlim([np.min(z_sample_hr), np.max(z_sample_hr)]) ax.set_title('Source keypoints', size=18) plt.show() curve = np.concatenate((x_sample_lr, y_sample_lr, z_sample_lr), 1) attitude = np.concatenate((a_x, a_y, a_z), 1) return curve, attitude if __name__ == "__main__": limits = { "x": [-0.6, 4], "y": [-2, 2], "z": [0.1, 2] } random_periodical_trajectory(plot=True, map_limits=limits) ================================================ FILE: ros_dd_mpc/src/utils/quad_3d_opt_utils.py ================================================ """ Set of utility functions for the quadrotor optimizer and simplified simulator. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import casadi as cs import numpy as np from src.quad_mpc.quad_3d import Quadrotor3D from tqdm import tqdm def discretize_dynamics_and_cost(t_horizon, n_points, m_steps_per_point, x, u, dynamics_f, cost_f, ind): """ Integrates the symbolic dynamics and cost equations until the time horizon using a RK4 method. :param t_horizon: time horizon in seconds :param n_points: number of control input points until time horizon :param m_steps_per_point: number of integrations steps per control input :param x: 4-element list with symbolic vectors for position (3D), angle (4D), velocity (3D) and rate (3D) :param u: 4-element symbolic vector for control input :param dynamics_f: symbolic dynamics function written in CasADi symbolic syntax. :param cost_f: symbolic cost function written in CasADi symbolic syntax. If None, then cost 0 is returned. :param ind: Only used for trajectory tracking. Index of cost function to use. :return: a symbolic function that computes the dynamics integration and the cost function at n_control_inputs points until the time horizon given an initial state and """ if isinstance(cost_f, list): # Select the list of cost functions cost_f = cost_f[ind * m_steps_per_point:(ind + 1) * m_steps_per_point] else: cost_f = [cost_f] * m_steps_per_point # Fixed step Runge-Kutta 4 integrator dt = t_horizon / n_points / m_steps_per_point x0 = x q = 0 for j in range(m_steps_per_point): k1 = dynamics_f(x=x, u=u)['x_dot'] k2 = dynamics_f(x=x + dt / 2 * k1, u=u)['x_dot'] k3 = dynamics_f(x=x + dt / 2 * k2, u=u)['x_dot'] k4 = dynamics_f(x=x + dt * k3, u=u)['x_dot'] x_out = x + dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4) x = x_out if cost_f and cost_f[j] is not None: q = q + cost_f[j](x=x, u=u)['q'] return cs.Function('F', [x0, u], [x, q], ['x0', 'p'], ['xf', 'qf']) def _forward_prop_core(x_0, u_seq, t_horizon, discrete_dynamics_f, dynamics_jac_f, B_x, gp_ensemble, covar, dt, m_int_steps, use_model): """ Propagates forward the state estimate described by the mean vector x_0 and the covariance matrix covar, and a sequence of inputs for the system u_seq. These inputs can either be numerical or symbolic. :param x_0: initial mean state of the state probability density function. Vector of length m :param u_seq: sequence of flattened N control inputs. I.e. vector of size N*4 :param t_horizon: time horizon corresponding to sequence of inputs :param discrete_dynamics_f: symbolic function to compute the discrete dynamics of the system. :param dynamics_jac_f: symbolic function to compute the dynamics jacobian of the system. :param B_x: Matrix to convert map from regressor output to state. :param gp_ensemble: The ensemble of GP's. Can be None if no GP's are used. :param covar: Initial covariance matrix of shape m x m, of the state probability density function. :param dt: Vector of N timestamps, the same length as w_opt / 2, corresponding to the time each input is applied. :param m_int_steps: number of intermediate integration steps per control node. :param use_model: The number (index) of the dynamics model to use from the available options. :return: The sequence of mean and covariance estimates for every corresponding input, as well as the computed cost for each stage. """ # Reshape input sequence to a N x 4 (1D) vector. Assume control input dim = 4 k = np.arange(int(u_seq.shape[0] / 4)) input_sequence = cs.horzcat(u_seq[4 * k], u_seq[4 * k + 1], u_seq[4 * k + 2], u_seq[4 * k + 3]) N = int(u_seq.shape[0] / 4) if dt is None: dt = t_horizon / N * np.ones((N, 1)) if len(dt.shape) == 1: dt = np.expand_dims(dt, 1) # Initialize sequence of propagated states mu_x = [x_0] cov_x = [covar] cost_x = [0] for k in range(N): # Get current control input and current state mean and covariance u_k = input_sequence[k, :] mu_k = mu_x[k] sig_k = cov_x[k] # mu(k+1) vector from propagation equations. Pass state through nominal dynamics with GP mean augmentation if GP # is available. Otherwise use nominal dynamics only. f_func = discrete_dynamics_f(dt[k], 1, m_int_steps, k, use_gp=gp_ensemble is not None, use_model=use_model) fk = f_func(x0=mu_k, p=u_k) mu_next = fk['xf'] stage_cost = fk['qf'] # K(k+1) matrix from propagation equations K_mat = sig_k # Evaluate linearized dynamics at current state l_mat = dynamics_jac_f(mu_k, u_k) * dt[k] + np.eye(mu_k.shape[0]) # Adjust matrices if GP model available. if gp_ensemble is not None: raise NotImplementedError # TODO: re-implement covariance propagation with GP # Get subset of features for GP regressor if GP regressor available z_k = cs.mtimes(B_z, cs.vertcat(mu_k, u_k.T)) # GP prediction. Stack predictions vertically along prediction dimension (e.g. vx, vy, ...) _, gp_covar_preds, gp_noise_prior = gp_ensemble.predict(z_k, return_cov=True, gp_idx=use_model) # Covariance forward propagation. l_mat = cs.horzcat(l_mat, B_x * dt[k]) # left-multiplying matrix sig_ld_comp = cs.mtimes(gp_prediction_jac(z_k, B_x, B_z, gp_ensemble, use_model), sig_k) K_mat = cs.vertcat(K_mat, sig_ld_comp) K_mat = cs.horzcat(K_mat, cs.vertcat(sig_ld_comp.T, cs.diag(gp_covar_preds + gp_noise_prior))) # Add next state estimate to lists mu_x += [mu_next] cov_x += [cs.mtimes(cs.mtimes(l_mat, K_mat), l_mat.T)] cost_x += [stage_cost] cov_x = [cov for cov in cov_x] return mu_x, cov_x, cost_x def uncertainty_forward_propagation(x_0, u_seq, t_horizon, discrete_dynamics_f, dynamics_jac_f, B_x=None, gp_regressors=None, covar=None, dt=None, m_integration_steps=1, use_model=0): if covar is None: covar = np.zeros((len(x_0), len(x_0))) else: assert covar.shape == (len(x_0), len(x_0)) x_0 = np.array(x_0) mu_x, cov_x, _ = _forward_prop_core(x_0, u_seq, t_horizon, discrete_dynamics_f, dynamics_jac_f, B_x, gp_regressors, covar, dt, m_integration_steps, use_model) mu_x = cs.horzcat(*mu_x) cov_x = cs.horzcat(*cov_x) mu_prop = np.array(mu_x).T cov_prop = np.array(cov_x).reshape((mu_prop.shape[1], mu_prop.shape[1], -1), order='F').transpose(2, 0, 1) return mu_prop, cov_prop def gp_prediction_jac(z, Bx, Bz, gp_ensemble, gp_idx): """ Computes the symbolic function for the Jacobian of the expected values of the joint GP regression model, evaluated at point z. :param z: A dictionary of symbolic vector at which the Jacobian must be evaluated. Each entry in the dictionary is indexed by the output dimension index. The dimension m_z on any given value must be the expected dimension of the feature inputs for the regressor. :param Bx: Matrix to convert map from regressor output to state. :param Bz: Matrix to map from (stacked) state vector and input vector to feature vector. If the gp_ensemble is not homogeneous, this parameter must be a dictionary specifying in each dimension which Bz matrix to use. :param gp_ensemble: GPEnsemble object with all the gp regressors :param gp_idx: which gp to compute the jacobian to from the ensemble :return: A Jacobian matrix of size n x m_x, where n is the number of variables regressed by the joint GP regressor and m_x is the dimension of the state vector. """ # Figure out which variables are necessary to compute the Jacobian wrt. out_dims = np.matmul(Bx, np.ones((Bx.shape[1], 1))) if isinstance(Bz, dict): z_dims = {} for dim in Bz.keys(): z_dims[dim] = np.matmul(Bz[dim][:, :len(out_dims)], out_dims) else: bz = np.matmul(Bz[:, :len(out_dims)], out_dims) z_dims = {k: v for k, v in zip(z.keys(), [bz] * len(z.keys()))} Bz = {k: v for k, v in zip(z.keys(), [Bz] * len(z.keys()))} jac = [] for dim in gp_idx.keys(): # Mapping from z to x inv_Bz = Bz[dim][:, :len(out_dims)].T gp = gp_ensemble.gp[dim][gp_idx[dim][0]] jac += [cs.mtimes(inv_Bz, gp.eval_gp_jac(z[dim]) * z_dims[dim])] return cs.horzcat(*jac).T def simulate_plant(quad, w_opt, simulation_dt, simulate_func, t_horizon=None, dt_vec=None, progress_bar=False): """ Given a sequence of n inputs, evaluates the simulated discrete-time plant model n steps into the future. The current drone state will not be changed by calling this method. :param quad: Quadrotor3D simulator object :type quad: Quadrotor3D :param w_opt: sequence of control n x m control inputs, where n is the number of steps and m is the dimensionality of a control input. :param simulation_dt: simulation step :param simulate_func: simulation function (inner loop) from the quadrotor MPC. :param t_horizon: time corresponding to the duration of the n control inputs. In the case that the w_opt comes from an MPC optimization, this parameter should be the MPC time horizon. :param dt_vec: a vector of timestamps, the same length as w_opt, corresponding to the total time each input is applied. :param progress_bar: boolean - whether to show a progress bar on the console or not. :return: the sequence of simulated quadrotor states. """ # Default_parameters if t_horizon is None and dt_vec is None: raise ValueError("At least the t_horizon or dt should be provided") if t_horizon is None: t_horizon = np.sum(dt_vec) state_safe = quad.get_state(quaternion=True, stacked=True) # Compute true simulated trajectory total_sim_time = t_horizon sim_traj = [] if dt_vec is None: change_control_input = np.arange(0, t_horizon, t_horizon / (w_opt.shape[0]))[1:] first_dt = t_horizon / w_opt.shape[0] sim_traj.append(quad.get_state(quaternion=True, stacked=True)) else: change_control_input = np.cumsum(dt_vec) first_dt = dt_vec[0] t_start_ep = 1e-6 int_range = tqdm(np.arange(t_start_ep, total_sim_time, simulation_dt)) if progress_bar else \ np.arange(t_start_ep, total_sim_time, simulation_dt) current_ind = 0 past_ind = 0 for t_elapsed in int_range: ref_u = w_opt[current_ind, :].T simulate_func(ref_u) if t_elapsed + simulation_dt >= first_dt: current_ind = np.argwhere(change_control_input <= t_elapsed + simulation_dt)[-1, 0] + 1 if past_ind != current_ind: sim_traj.append(quad.get_state(quaternion=True, stacked=True)) past_ind = current_ind if dt_vec is None: sim_traj.append(quad.get_state(quaternion=True, stacked=True)) sim_traj = np.squeeze(sim_traj) quad.set_state(state_safe) return sim_traj def get_reference_chunk(reference_traj, reference_u, current_idx, n_mpc_nodes, reference_over_sampling): """ Extracts the reference states and controls for the current MPC optimization given the over-sampled counterparts. :param reference_traj: The reference trajectory, which has been finely over-sampled by a factor of reference_over_sampling. It should be a vector of shape (Nx13), where N is the length of the trajectory in samples. :param reference_u: The reference controls, following the same requirements as reference_traj. Should be a vector of shape (Nx4). :param current_idx: Current index of the trajectory tracking. Should be an integer number between 0 and N-1. :param n_mpc_nodes: Number of MPC nodes considered in the optimization. :param reference_over_sampling: The over-sampling factor of the reference trajectories. Should be a positive integer. :return: Returns the chunks of reference selected for the current MPC iteration. Two numpy arrays will be returned: - An ((N+1)x13) array, corresponding to the reference trajectory. The first row is the state of current_idx. - An (Nx4) array, corresponding to the reference controls. """ # Dense references ref_traj_chunk = reference_traj[current_idx:current_idx + (n_mpc_nodes + 1) * reference_over_sampling, :] ref_u_chunk = reference_u[current_idx:current_idx + n_mpc_nodes * reference_over_sampling, :] # Indices for down-sampling the reference to number of MPC nodes downsample_ref_ind = np.arange(0, min(reference_over_sampling * (n_mpc_nodes + 1), ref_traj_chunk.shape[0]), reference_over_sampling, dtype=int) # Sparser references (same dt as node separation) ref_traj_chunk = ref_traj_chunk[downsample_ref_ind, :] ref_u_chunk = ref_u_chunk[downsample_ref_ind[:max(len(downsample_ref_ind) - 1, 1)], :] return ref_traj_chunk, ref_u_chunk ================================================ FILE: ros_dd_mpc/src/utils/trajectories.py ================================================ """ Trajectory generation functions. For the circle, lemniscate and random trajectories. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import numpy as np from src.utils.utils import undo_quaternion_flip, rotation_matrix_to_quat from src.utils.utils import quaternion_inverse, q_dot_q from src.utils.trajectory_generator import draw_poly, get_full_traj, fit_multi_segment_polynomial_trajectory from src.utils.keyframe_3d_gen import random_periodical_trajectory from config.configuration_parameters import DirectoryConfig from src.quad_mpc.quad_3d import Quadrotor3D import matplotlib.pyplot as plt import os import yaml import json def check_trajectory(trajectory, inputs, tvec, plot=False): """ @param trajectory: @param inputs: @param tvec: @param plot: @return: """ print("Checking trajectory integrity...") dt = np.expand_dims(np.gradient(tvec, axis=0), axis=1) numeric_derivative = np.gradient(trajectory, axis=0) / dt errors = np.zeros((dt.shape[0], 3)) num_bodyrates = [] for i in range(dt.shape[0]): # 1) check if velocity is consistent with position numeric_velocity = numeric_derivative[i, 0:3] analytic_velocity = trajectory[i, 7:10] errors[i, 0] = np.linalg.norm(numeric_velocity - analytic_velocity) if not np.allclose(analytic_velocity, numeric_velocity, atol=1e-2, rtol=1e-2): print("inconsistent linear velocity") print(numeric_velocity) print(analytic_velocity) return False # 2) check if attitude is consistent with acceleration gravity = 9.81 numeric_thrust = numeric_derivative[i, 7:10] + np.array([0.0, 0.0, gravity]) numeric_thrust = numeric_thrust / np.linalg.norm(numeric_thrust) analytic_attitude = trajectory[i, 3:7] if np.abs(np.linalg.norm(analytic_attitude) - 1.0) > 1e-6: print("quaternion does not have unit norm!") print(analytic_attitude) print(np.linalg.norm(analytic_attitude)) return False e_z = np.array([0.0, 0.0, 1.0]) q_w = 1.0 + np.dot(e_z, numeric_thrust) q_xyz = np.cross(e_z, numeric_thrust) numeric_attitude = 0.5 * np.array([q_w] + q_xyz.tolist()) numeric_attitude = numeric_attitude / np.linalg.norm(numeric_attitude) # the two attitudes can only differ in yaw --> check x,y component q_diff = q_dot_q(quaternion_inverse(analytic_attitude), numeric_attitude) errors[i, 1] = np.linalg.norm(q_diff[1:3]) if not np.allclose(q_diff[1:3], np.zeros(2, ), atol=1e-3, rtol=1e-3): print("Attitude and acceleration do not match!") print(analytic_attitude) print(numeric_attitude) print(q_diff) return False # 3) check if bodyrates agree with attitude difference numeric_bodyrates = 2.0 * q_dot_q(quaternion_inverse(trajectory[i, 3:7]), numeric_derivative[i, 3:7])[1:] num_bodyrates.append(numeric_bodyrates) analytic_bodyrates = trajectory[i, 10:13] errors[i, 2] = np.linalg.norm(numeric_bodyrates - analytic_bodyrates) if not np.allclose(numeric_bodyrates, analytic_bodyrates, atol=0.05, rtol=0.05): print("inconsistent angular velocity") print(numeric_bodyrates) print(analytic_bodyrates) return False print("Trajectory check successful") print("Maximum linear velocity error: %.5f" % np.max(errors[:, 0])) print("Maximum attitude error: %.5f" % np.max(errors[:, 1])) print("Maximum angular velocity error: %.5f" % np.max(errors[:, 2])) if plot: num_bodyrates = np.stack(num_bodyrates) plt.figure() for i in range(3): plt.subplot(3, 2, i * 2 + 1) plt.plot(numeric_derivative[:, i], label='numeric') plt.plot(trajectory[:, 7 + i], label='analytic') plt.ylabel('m/s') if i == 0: plt.title("Velocity check") plt.legend() for i in range(3): plt.subplot(3, 2, i * 2 + 2) plt.plot(num_bodyrates[:, i], label='numeric') plt.plot(trajectory[:, 10 + i], label='analytic') plt.ylabel('rad/s') if i == 0: plt.title("Body rate check") plt.legend() plt.suptitle('Integrity check of reference trajectory') plt.show() return True def minimum_snap_trajectory_generator(traj_derivatives, yaw_derivatives, t_ref, quad, map_limits, plot, adjust_pos=True): """ Follows the Minimum Snap Trajectory paper to generate a full trajectory given the position reference and its derivatives, and the yaw trajectory and its derivatives. :param traj_derivatives: np.array of shape 4x3xN. N corresponds to the length in samples of the trajectory, and: - The 4 components of the first dimension correspond to position, velocity, acceleration and jerk. - The 3 components of the second dimension correspond to x, y, z. :param yaw_derivatives: np.array of shape 2xN. N corresponds to the length in samples of the trajectory. The first row is the yaw trajectory, and the second row is the yaw time-derivative trajectory. :param t_ref: vector of length N, containing the reference times (starting from 0) for the trajectory. :param quad: Quadrotor3D object, corresponding to the quadrotor model that will track the generated reference. :type quad: Quadrotor3D :param map_limits: dictionary of map limits if available, None otherwise. :param plot: True if show a plot of the generated trajectory. :return: tuple of 3 arrays: - Nx13 array of generated reference trajectory. The 13 dimension contains the components: position_xyz, attitude_quaternion_wxyz, velocity_xyz, body_rate_xyz. - N array of reference timestamps. The same as in the input - Nx4 array of reference controls, corresponding to the four motors of the quadrotor. """ discretization_dt = t_ref[1] - t_ref[0] len_traj = traj_derivatives.shape[2] # Add gravity to accelerations gravity = 9.81 thrust = traj_derivatives[2, :, :].T + np.tile(np.array([[0, 0, 1]]), (len_traj, 1)) * gravity # Compute body axes z_b = thrust / np.sqrt(np.sum(thrust ** 2, 1))[:, np.newaxis] yawing = np.any(yaw_derivatives[0, :] != 0) rate = np.zeros((len_traj, 3)) f_t = np.zeros((len_traj, 1)) for i in range(len_traj): f_t[i, 0] = quad.mass * z_b[i].dot(thrust[i, :].T) if yawing: # yaw is defined as the projection of the body-x axis on the horizontal plane x_c = np.concatenate((np.cos(yaw_derivatives[0, :])[:, np.newaxis], np.sin(yaw_derivatives[0, :])[:, np.newaxis], np.zeros(len_traj)[:, np.newaxis]), 1) y_b = np.cross(z_b, x_c) y_b = y_b / np.sqrt(np.sum(y_b ** 2, axis=1))[:, np.newaxis] x_b = np.cross(y_b, z_b) # Rotation matrix (from body to world) b_r_w = np.concatenate((x_b[:, :, np.newaxis], y_b[:, :, np.newaxis], z_b[:, :, np.newaxis]), -1) q = [] for i in range(len_traj): # Transform to quaternion q.append(rotation_matrix_to_quat(b_r_w[i])) if i > 1: q[-1] = undo_quaternion_flip(q[-2], q[-1]) q = np.stack(q) # Compute angular rate vector # Total thrust acceleration must be equal to the projection of the quadrotor acceleration into the Z body axis a_proj = np.zeros((len_traj, 1)) for i in range(len_traj): a_proj[i, 0] = z_b[i].dot(traj_derivatives[3, :, i]) h_omega = quad.mass / f_t * (traj_derivatives[3, :, :].T - a_proj * z_b) for i in range(len_traj): rate[i, 0] = -h_omega[i].dot(y_b[i]) rate[i, 1] = h_omega[i].dot(x_b[i]) rate[i, 2] = -yaw_derivatives[1, i] * np.array([0, 0, 1]).dot(z_b[i]) else: # new way to compute attitude: # https://math.stackexchange.com/questions/2251214/calculate-quaternions-from-two-directional-vectors e_z = np.array([[0.0, 0.0, 1.0]]) q_w = 1.0 + np.sum(e_z * z_b, axis=1) q_xyz = np.cross(e_z, z_b) q = 0.5 * np.concatenate([np.expand_dims(q_w, axis=1), q_xyz], axis=1) q = q / np.sqrt(np.sum(q ** 2, 1))[:, np.newaxis] # Use numerical differentiation of quaternions q_dot = np.gradient(q, axis=0) / discretization_dt w_int = np.zeros((len_traj, 3)) for i in range(len_traj): w_int[i, :] = 2.0 * q_dot_q(quaternion_inverse(q[i, :]), q_dot[i])[1:] rate[:, 0] = w_int[:, 0] rate[:, 1] = w_int[:, 1] rate[:, 2] = w_int[:, 2] go_crazy_about_yaw = True if go_crazy_about_yaw: print("Maximum yawrate before adaption: %.3f" % np.max(np.abs(rate[:, 2]))) q_new = q yaw_corr_acc = 0.0 for i in range(1, len_traj): yaw_corr = -rate[i, 2] * discretization_dt yaw_corr_acc += yaw_corr q_corr = np.array([np.cos(yaw_corr_acc / 2.0), 0.0, 0.0, np.sin(yaw_corr_acc / 2.0)]) q_new[i, :] = q_dot_q(q[i, :], q_corr) w_int[i, :] = 2.0 * q_dot_q(quaternion_inverse(q[i, :]), q_dot[i])[1:] q_new_dot = np.gradient(q_new, axis=0) / discretization_dt for i in range(1, len_traj): w_int[i, :] = 2.0 * q_dot_q(quaternion_inverse(q_new[i, :]), q_new_dot[i])[1:] q = q_new rate[:, 0] = w_int[:, 0] rate[:, 1] = w_int[:, 1] rate[:, 2] = w_int[:, 2] print("Maximum yawrate after adaption: %.3f" % np.max(np.abs(rate[:, 2]))) # Compute inputs rate_dot = np.gradient(rate, axis=0) / discretization_dt rate_x_Jrate = np.array([(quad.J[2] - quad.J[1]) * rate[:, 2] * rate[:, 1], (quad.J[0] - quad.J[2]) * rate[:, 0] * rate[:, 2], (quad.J[1] - quad.J[0]) * rate[:, 1] * rate[:, 0]]).T tau = rate_dot * quad.J[np.newaxis, :] + rate_x_Jrate b = np.concatenate((tau, f_t), axis=-1) a_mat = np.concatenate((quad.y_f[np.newaxis, :], -quad.x_f[np.newaxis, :], quad.z_l_tau[np.newaxis, :], np.ones_like(quad.z_l_tau)[np.newaxis, :]), 0) reference_u = np.zeros((len_traj, 4)) for i in range(len_traj): reference_u[i, :] = np.linalg.solve(a_mat, b[i, :]) full_pos = traj_derivatives[0, :, :].T full_vel = traj_derivatives[1, :, :].T reference_traj = np.concatenate((full_pos, q, full_vel, rate), 1) if adjust_pos: if map_limits is None: # Locate starting point right at x=0 and y=0. reference_traj[:, 0] -= reference_traj[0, 0] reference_traj[:, 1] -= reference_traj[0, 1] else: x_max_range = map_limits["x"][1] - map_limits["x"][0] y_max_range = map_limits["y"][1] - map_limits["y"][0] z_max_range = map_limits["z"][1] - map_limits["z"][0] x_center = x_max_range / 2 + map_limits["x"][0] y_center = y_max_range / 2 + map_limits["y"][0] z_center = z_max_range / 2 + map_limits["z"][0] # Center circle to center of map XY plane reference_traj[:, :3] += np.array([x_center, y_center, 0]) reference_traj[:, 2] = z_center if plot: draw_poly(reference_traj, reference_u, t_ref) # Change format of reference input to motor activation, in interval [0, 1] reference_u = reference_u / quad.max_thrust return reference_traj, t_ref, reference_u def load_map_limits_from_file(map_limits): import rospy if map_limits is not None and map_limits != "None": config_path = DirectoryConfig.CONFIG_DIR params_file = os.path.join(config_path, map_limits + '.yaml') try: with open(params_file) as file: limits = yaml.full_load(file) map_limits = {"x": [limits["x_min"], limits["x_max"]], "y": [limits["y_min"], limits["y_max"]], "z": [limits["z_min"], limits["z_max"]]} rospy.loginfo("Using world limits: " + json.dumps(limits)) except FileNotFoundError: warn_msg = "Tried to load environment limits: '%s', but the file was not found. Using default limits." \ % map_limits rospy.logwarn(warn_msg) map_limits = None else: map_limits = None return map_limits def straight_trajectory(quad, discretization_dt, speed): n = 3 pos_traj = np.zeros((n, 3)) pos_traj[:, 1] = np.linspace(-3, 3, n) pos_traj[:, 2] = 1 att_traj = np.zeros_like(pos_traj) av_dist = np.mean(np.sqrt(np.sum(np.diff(pos_traj, axis=0) ** 2, axis=1))) av_dt = av_dist / speed poly_pos_traj = fit_multi_segment_polynomial_trajectory(pos_traj.T, att_traj[:, -1].T) traj, yaw, t_ref = get_full_traj(poly_pos_traj, target_dt=av_dt, int_dt=discretization_dt) reference_traj, t_ref, reference_u = minimum_snap_trajectory_generator(traj, yaw, t_ref, quad, None, False) return reference_traj, t_ref, reference_u def flyover_trajectory_collect(quad, discretization_dt, seed, speed, flyover_box_name): np.random.seed(seed) flyover_box = load_map_limits_from_file(flyover_box_name) box_x_len = flyover_box['x'][1] - flyover_box['x'][0] box_y_len = flyover_box['y'][1] - flyover_box['y'][0] box_z_len = flyover_box['z'][1] - flyover_box['z'][0] box_center = np.array([flyover_box['x'][0] + box_x_len / 2, flyover_box['y'][0] + box_y_len / 2]) box_diag = (box_x_len ** 2 + box_y_len ** 2) ** 0.5 pos_traj = np.zeros((50, 3)) for i in range(0, 50, 5): rand_direction = (np.random.rand(2) - 0.5) * 2 rand_direction = rand_direction / np.linalg.norm(rand_direction) height_start = flyover_box['z'][0] + np.random.rand() * box_z_len start_point = box_center + rand_direction * 0.6 * box_diag direction = (box_center - start_point) direction = direction / np.linalg.norm(direction) height_flyover = flyover_box['z'][0] + np.random.rand() * box_z_len flyover_point = box_center - direction * 0.25 * box_diag height_flyover_center = 0.55 flyover_center_point = box_center height_flyover2 = flyover_box['z'][0] + np.random.rand() * box_z_len flyover_point2 = box_center + direction * 0.25 * box_diag height_target = flyover_box['z'][0] + np.random.rand() * box_z_len target_point = box_center + direction * 0.7 * box_diag pos_traj[i, :2] = start_point pos_traj[i, 2] = height_start pos_traj[i + 1, :2] = flyover_point pos_traj[i + 1, 2] = height_flyover pos_traj[i + 2, :2] = flyover_center_point pos_traj[i + 2, 2] = height_flyover_center pos_traj[i + 3, :2] = flyover_point2 pos_traj[i + 3, 2] = height_flyover2 pos_traj[i + 4, :2] = target_point pos_traj[i + 4, 2] = height_target att_traj = np.zeros_like(pos_traj) av_dist = np.mean(np.sqrt(np.sum(np.diff(pos_traj, axis=0) ** 2, axis=1))) av_dt = av_dist / speed poly_pos_traj = fit_multi_segment_polynomial_trajectory(pos_traj.T, att_traj[:, -1].T) traj, yaw, t_ref = get_full_traj(poly_pos_traj, target_dt=av_dt, int_dt=discretization_dt) reference_traj, t_ref, reference_u = minimum_snap_trajectory_generator(traj, yaw, t_ref, quad, None, False, adjust_pos=False) return reference_traj, t_ref, reference_u def flyover_trajectory(quad, discretization_dt, seed, speed, flyover_box_name): np.random.seed(seed) flyover_box = load_map_limits_from_file(flyover_box_name) box_x_len = flyover_box['x'][1] - flyover_box['x'][0] box_y_len = flyover_box['y'][1] - flyover_box['y'][0] box_z_len = flyover_box['z'][1] - flyover_box['z'][0] box_center = np.array([flyover_box['x'][0] + box_x_len / 2, flyover_box['y'][0] + box_y_len / 2]) box_diag = (box_x_len ** 2 + box_y_len ** 2) ** 0.5 pos_traj = np.zeros((17, 3)) rand_direction = (np.random.rand(2) - 0.5) * 2 rand_direction = rand_direction / np.linalg.norm(rand_direction) height = flyover_box['z'][0] + np.random.rand() * box_z_len start_point = box_center + rand_direction * 0.6 * box_diag pos_traj[0, :2] = start_point pos_traj[0, 2] = height for i in range(1, 16, 3): rand_direction = (np.random.rand(2) - 0.5) * 2 rand_direction = rand_direction / np.linalg.norm(rand_direction) height = flyover_box['z'][0] + np.random.rand() * box_z_len start_point = box_center + rand_direction * 0.6 * box_diag direction = (box_center - start_point) direction = direction / np.linalg.norm(direction) flyover_center_point = box_center target_point = box_center + direction * 0.6 * box_diag pos_traj[i, :2] = start_point pos_traj[i, 2] = height pos_traj[i + 1, :2] = flyover_center_point pos_traj[i + 1, 2] = height pos_traj[i + 2, :2] = target_point pos_traj[i + 2, 2] = height target_point = box_center + direction * 1.2 * box_diag pos_traj[i + 3, :2] = target_point pos_traj[i + 3, 2] = height att_traj = np.zeros_like(pos_traj) av_dist = np.mean(np.sqrt(np.sum(np.diff(pos_traj, axis=0) ** 2, axis=1))) av_dt = av_dist / speed poly_pos_traj = fit_multi_segment_polynomial_trajectory(pos_traj.T, att_traj[:, -1].T) traj, yaw, t_ref = get_full_traj(poly_pos_traj, target_dt=av_dt, int_dt=discretization_dt) reference_traj, t_ref, reference_u = minimum_snap_trajectory_generator(traj, yaw, t_ref, quad, None, False, adjust_pos=False) return reference_traj, t_ref, reference_u def random_trajectory(quad, discretization_dt, seed, speed, map_name=None, plot=False): map_limits = load_map_limits_from_file(map_name) # Get a random smooth position trajectory pos_traj, att_traj = random_periodical_trajectory(random_state=seed, map_limits=map_limits, plot=False) if map_limits is None: # Locate starting point right at x=0 and y=0. pos_traj[:, 0] -= pos_traj[0, 0] pos_traj[:, 1] -= pos_traj[0, 1] # Ensure no negative z position min_z_pos = np.min(pos_traj[:, -1]) if min_z_pos < 0: pos_traj[:, -1] = pos_traj[:, -1] + 2 * min_z_pos * np.sign(min_z_pos) # Calculate feasible time based on the average distance between two position track points: av_dist = np.mean(np.sqrt(np.sum(np.diff(pos_traj, axis=0) ** 2, axis=1))) av_dt = av_dist / speed # Calculate the polynomial fit to the position trajectory, and compute the full kinematic reference. This # trajectory is sampled according to the frequency that will be needed for the MPC. poly_pos_traj = fit_multi_segment_polynomial_trajectory(pos_traj.T, att_traj[:, -1].T) traj, yaw, t_ref = get_full_traj(poly_pos_traj, av_dt, discretization_dt) reference_traj, t_ref, reference_u = minimum_snap_trajectory_generator(traj, yaw, t_ref, quad, map_limits, False) if plot: draw_poly(reference_traj, reference_u, t_ref, pos_traj.T) return reference_traj, t_ref, reference_u def loop_trajectory(quad, discretization_dt, radius, z, lin_acc, clockwise, yawing, v_max, map_name, plot): """ Creates a circular trajectory on the x-y plane that increases speed by 1m/s at every revolution. :param quad: Quadrotor model :param discretization_dt: Sampling period of the trajectory. :param radius: radius of loop trajectory in meters :param z: z position of loop plane in meters :param lin_acc: linear acceleration of trajectory (and successive deceleration) in m/s^2 :param clockwise: True if the rotation will be done clockwise. :param yawing: True if the quadrotor yaws along the trajectory. False for 0 yaw trajectory. :param v_max: Maximum speed at peak velocity. Revolutions needed will be calculated automatically. :param map_name: Name of map to load its limits :param plot: Whether to plot an analysis of the planned trajectory or not. :return: The full 13-DoF trajectory with time and input vectors """ # Apply map limits to radius map_limits = load_map_limits_from_file(map_name) if map_limits is not None: x_max_range = map_limits["x"][1] - map_limits["x"][0] y_max_range = map_limits["y"][1] - map_limits["y"][0] max_radius = min(x_max_range / 2, y_max_range / 2) radius = min(radius, max_radius) assert z > 0 ramp_up_t = 2 # s # Calculate simulation time to achieve desired maximum velocity with specified acceleration t_total = 2 * v_max / lin_acc + 2 * ramp_up_t # Transform to angular acceleration alpha_acc = lin_acc / radius # rad/s^2 # Generate time and angular acceleration sequences # Ramp up sequence ramp_t_vec = np.arange(0, ramp_up_t, discretization_dt) ramp_up_alpha = alpha_acc * np.sin(np.pi / (2 * ramp_up_t) * ramp_t_vec) ** 2 # Acceleration phase coasting_duration = (t_total - 4 * ramp_up_t) / 2 coasting_t_vec = ramp_up_t + np.arange(0, coasting_duration, discretization_dt) coasting_alpha = np.ones_like(coasting_t_vec) * alpha_acc # Transition phase: decelerate transition_t_vec = np.arange(0, 2 * ramp_up_t, discretization_dt) transition_alpha = alpha_acc * np.cos(np.pi / (2 * ramp_up_t) * transition_t_vec) transition_t_vec += coasting_t_vec[-1] + discretization_dt # Deceleration phase down_coasting_t_vec = transition_t_vec[-1] + np.arange(0, coasting_duration, discretization_dt) + discretization_dt down_coasting_alpha = -np.ones_like(down_coasting_t_vec) * alpha_acc # Bring to rest phase ramp_up_t_vec = down_coasting_t_vec[-1] + np.arange(0, ramp_up_t, discretization_dt) + discretization_dt ramp_up_alpha_end = ramp_up_alpha - alpha_acc # Concatenate all sequences t_ref = np.concatenate((ramp_t_vec, coasting_t_vec, transition_t_vec, down_coasting_t_vec, ramp_up_t_vec)) alpha_vec = np.concatenate(( ramp_up_alpha, coasting_alpha, transition_alpha, down_coasting_alpha, ramp_up_alpha_end)) # Calculate derivative of angular acceleration (alpha_vec) ramp_up_alpha_dt = alpha_acc * np.pi / (2 * ramp_up_t) * np.sin(np.pi / ramp_up_t * ramp_t_vec) coasting_alpha_dt = np.zeros_like(coasting_alpha) transition_alpha_dt = - alpha_acc * np.pi / (2 * ramp_up_t) * np.sin(np.pi / (2 * ramp_up_t) * transition_t_vec) alpha_dt = np.concatenate(( ramp_up_alpha_dt, coasting_alpha_dt, transition_alpha_dt, coasting_alpha_dt, ramp_up_alpha_dt)) if not clockwise: alpha_vec *= -1 alpha_dt *= -1 # Compute angular integrals w_vec = np.cumsum(alpha_vec) * discretization_dt angle_vec = np.cumsum(w_vec) * discretization_dt # Compute position, velocity, acceleration, jerk pos_traj_x = radius * np.sin(angle_vec)[np.newaxis, np.newaxis, :] pos_traj_y = radius * np.cos(angle_vec)[np.newaxis, np.newaxis, :] pos_traj_z = np.ones_like(pos_traj_x) * z vel_traj_x = (radius * w_vec * np.cos(angle_vec))[np.newaxis, np.newaxis, :] vel_traj_y = - (radius * w_vec * np.sin(angle_vec))[np.newaxis, np.newaxis, :] acc_traj_x = radius * (alpha_vec * np.cos(angle_vec) - w_vec ** 2 * np.sin(angle_vec))[np.newaxis, np.newaxis, :] acc_traj_y = - radius * (alpha_vec * np.sin(angle_vec) + w_vec ** 2 * np.cos(angle_vec))[np.newaxis, np.newaxis, :] jerk_traj_x = radius * (alpha_dt * np.cos(angle_vec) - alpha_vec * np.sin(angle_vec) * w_vec - np.cos(angle_vec) * w_vec ** 3 - 2 * np.sin(angle_vec) * w_vec * alpha_vec) jerk_traj_y = - radius * (np.cos(angle_vec) * w_vec * alpha_vec + np.sin(angle_vec) * alpha_dt - np.sin(angle_vec) * w_vec ** 3 + 2 * np.cos(angle_vec) * w_vec * alpha_vec) jerk_traj_x = jerk_traj_x[np.newaxis, np.newaxis, :] jerk_traj_y = jerk_traj_y[np.newaxis, np.newaxis, :] if yawing: yaw_traj = -angle_vec else: yaw_traj = np.zeros_like(angle_vec) traj = np.concatenate(( np.concatenate((pos_traj_x, pos_traj_y, pos_traj_z), 1), np.concatenate((vel_traj_x, vel_traj_y, np.zeros_like(vel_traj_x)), 1), np.concatenate((acc_traj_x, acc_traj_y, np.zeros_like(acc_traj_x)), 1), np.concatenate((jerk_traj_x, jerk_traj_y, np.zeros_like(jerk_traj_x)), 1)), 0) yaw = np.concatenate((yaw_traj[np.newaxis, :], w_vec[np.newaxis, :]), 0) return minimum_snap_trajectory_generator(traj, yaw, t_ref, quad, map_limits, plot) def lemniscate_trajectory(quad, discretization_dt, radius, z, lin_acc, clockwise, yawing, v_max, map_name, plot): """ :param quad: :param discretization_dt: :param radius: :param z: :param lin_acc: :param clockwise: :param yawing: :param v_max: :param map_name: :param plot: :return: """ # Apply map limits to radius map_limits = load_map_limits_from_file(map_name) if map_limits is not None: x_max_range = map_limits["x"][1] - map_limits["x"][0] y_max_range = map_limits["y"][1] - map_limits["y"][0] max_radius = min(x_max_range / 2, y_max_range / 2) radius = min(radius, max_radius) assert z > 0 ramp_up_t = 2 # s # Calculate simulation time to achieve desired maximum velocity with specified acceleration t_total = 2 * v_max / lin_acc + 2 * ramp_up_t # Transform to angular acceleration alpha_acc = lin_acc / radius # rad/s^2 # Generate time and angular acceleration sequences # Ramp up sequence ramp_t_vec = np.arange(0, ramp_up_t, discretization_dt) ramp_up_alpha = alpha_acc * np.sin(np.pi / (2 * ramp_up_t) * ramp_t_vec) ** 2 # Acceleration phase coasting_duration = (t_total - 4 * ramp_up_t) / 2 coasting_t_vec = ramp_up_t + np.arange(0, coasting_duration, discretization_dt) coasting_alpha = np.ones_like(coasting_t_vec) * alpha_acc # Transition phase: decelerate transition_t_vec = np.arange(0, 2 * ramp_up_t, discretization_dt) transition_alpha = alpha_acc * np.cos(np.pi / (2 * ramp_up_t) * transition_t_vec) transition_t_vec += coasting_t_vec[-1] + discretization_dt # Deceleration phase down_coasting_t_vec = transition_t_vec[-1] + np.arange(0, coasting_duration, discretization_dt) + discretization_dt down_coasting_alpha = -np.ones_like(down_coasting_t_vec) * alpha_acc # Bring to rest phase ramp_up_t_vec = down_coasting_t_vec[-1] + np.arange(0, ramp_up_t, discretization_dt) + discretization_dt ramp_up_alpha_end = ramp_up_alpha - alpha_acc # Concatenate all sequences t_ref = np.concatenate((ramp_t_vec, coasting_t_vec, transition_t_vec, down_coasting_t_vec, ramp_up_t_vec)) alpha_vec = np.concatenate(( ramp_up_alpha, coasting_alpha, transition_alpha, down_coasting_alpha, ramp_up_alpha_end)) # Compute angular integrals w_vec = np.cumsum(alpha_vec) * discretization_dt angle_vec = np.cumsum(w_vec) * discretization_dt # Adaption: we achieve the highest spikes in the bodyrates when passing through the 'center' part of the figure-8 # This leads to negative reference thrusts. # Let's see if we can alleviate this by adapting the z-reference in these parts to add some acceleration in the # z-component z_dim = 0.0 # Compute position, velocity, acceleration, jerk pos_traj_x = radius * np.cos(angle_vec)[np.newaxis, np.newaxis, :] pos_traj_y = radius * (np.sin(angle_vec) * np.cos(angle_vec))[np.newaxis, np.newaxis, :] pos_traj_z = - z_dim * np.cos(4.0 * angle_vec)[np.newaxis, np.newaxis, :] + z vel_traj_x = -radius * (w_vec * np.sin(angle_vec))[np.newaxis, np.newaxis, :] vel_traj_y = radius * (w_vec * np.cos(angle_vec) ** 2 - w_vec * np.sin(angle_vec) ** 2)[np.newaxis, np.newaxis, :] vel_traj_z = 4.0 * z_dim * w_vec * np.sin(4.0 * angle_vec)[np.newaxis, np.newaxis, :] acc_traj_x = -radius * (alpha_vec * np.sin(angle_vec) + w_vec ** 2 * np.cos(angle_vec)) acc_traj_y = radius * (alpha_vec * np.cos(angle_vec) ** 2 - 2.0 * w_vec ** 2 * np.cos(angle_vec) * np.sin( angle_vec) - alpha_vec * np.sin(angle_vec) ** 2 - 2.0 * w_vec ** 2 * np.sin(angle_vec) * np.cos(angle_vec)) acc_traj_z = 16.0 * z_dim * (w_vec ** 2 * np.cos(4.0 * angle_vec) + alpha_vec * np.sin(4.0 * angle_vec)) acc_traj_x = acc_traj_x[np.newaxis, np.newaxis, :] acc_traj_y = acc_traj_y[np.newaxis, np.newaxis, :] acc_traj_z = acc_traj_z[np.newaxis, np.newaxis, :] traj = np.concatenate(( np.concatenate((pos_traj_x, pos_traj_y, pos_traj_z), 1), np.concatenate((vel_traj_x, vel_traj_y, vel_traj_z), 1), np.concatenate((acc_traj_x, acc_traj_y, acc_traj_z), 1)), 0) yaw = np.zeros_like(traj) return minimum_snap_trajectory_generator(traj, yaw, t_ref, quad, map_limits, plot) ================================================ FILE: ros_dd_mpc/src/utils/trajectory_generator.py ================================================ import numpy as np import matplotlib.pyplot as plt def draw_poly(traj, u_traj, t, target_points=None, target_t=None): """ Plots the generated trajectory of length n with the used keypoints. :param traj: Full generated reference trajectory. Numpy array of shape nx13 :param u_traj: Generated reference inputs. Numpy array of shape nx4 :param t: Timestamps of the references. Numpy array of length n :param target_points: m position keypoints used for trajectory generation. Numpy array of shape 3 x m. :param target_t: Timestamps of the reference position keypoints. If not passed, then they are extracted from the t vector, assuming constant time separation. """ ders = 2 dims = 3 y_labels = [r'pos $[m]$', r'vel $[m/s]$', r'acc $[m/s^2]$', r'jer $[m/s^3]$'] dim_legends = ['x', 'y', 'z'] if target_t is None and target_points is not None: target_t = np.linspace(0, t[-1], target_points.shape[1]) p_traj = traj[:, :3] a_traj = traj[:, 3:7] v_traj = traj[:, 7:10] r_traj = traj[:, 10:] plt_traj = [p_traj, v_traj] fig = plt.figure() for d_ord in range(ders): plt.subplot(ders + 2, 2, d_ord * 2 + 1) for dim in range(dims): plt.plot(t, plt_traj[d_ord][:, dim], label=dim_legends[dim]) if d_ord == 0 and target_points is not None: plt.plot(target_t, target_points[dim, :], 'bo') plt.gca().set_xticklabels([]) plt.legend() plt.grid() plt.ylabel(y_labels[d_ord]) dim_legends = [['w', 'x', 'y', 'z'], ['x', 'y', 'z']] y_labels = [r'att $[quat]$', r'rate $[rad/s]$'] plt_traj = [a_traj, r_traj] for d_ord in range(ders): plt.subplot(ders + 2, 2, d_ord * 2 + 1 + ders * 2) for dim in range(plt_traj[d_ord].shape[1]): plt.plot(t, plt_traj[d_ord][:, dim], label=dim_legends[d_ord][dim]) plt.legend() plt.grid() plt.ylabel(y_labels[d_ord]) if d_ord == ders - 1: plt.xlabel(r'time $[s]$') else: plt.gca().set_xticklabels([]) ax = fig.add_subplot(2, 2, 2, projection='3d') plt.plot(p_traj[:, 0], p_traj[:, 1], p_traj[:, 2]) if target_points is not None: plt.plot(target_points[0, :], target_points[1, :], target_points[2, :], 'bo') plt.title('Target position trajectory') ax.set_xlabel(r'$p_x [m]$') ax.set_ylabel(r'$p_y [m]$') ax.set_zlabel(r'$p_z [m]$') plt.subplot(ders + 1, 2, (ders + 1) * 2) for i in range(u_traj.shape[1]): plt.plot(t, u_traj[:, i], label=r'$u_{}$'.format(i)) plt.grid() plt.legend() plt.gca().yaxis.set_label_position("right") plt.gca().yaxis.tick_right() plt.xlabel(r'time $[s]$') plt.ylabel(r'single thrusts $[N]$') plt.title('Control inputs') plt.suptitle('Generated polynomial trajectory') plt.show() def get_full_traj(poly_coeffs, target_dt, int_dt): dims = poly_coeffs.shape[-1] full_traj = np.zeros((4, dims, 0)) t_total = np.zeros((0,)) if isinstance(target_dt, float): # Adjust target_dt to make it divisible by int_dt target_dt = round(target_dt / int_dt) * int_dt # Assign target time for each keypoint using homogeneous spacing t_vec = np.arange(0, target_dt * (poly_coeffs.shape[0] + 1) - 1e-5, target_dt) else: # The time between each pair of points is assigned independently # First, also adjust each value of the target_dt vector to make it divisible by int_dt for i, dt in enumerate(target_dt): target_dt[i] = round(dt / int_dt) * int_dt t_vec = np.append(np.zeros(1), np.cumsum(target_dt[:-1])) for seg in range(len(t_vec) - 1): # Select time sampling (linear or quadratic) mode tau_dt = np.arange(t_vec[seg], t_vec[seg + 1] + 1e-5, int_dt) # Re-normalize time sampling vector between -1 and 1 t1 = (tau_dt - t_vec[seg]) / (t_vec[seg + 1] - t_vec[seg]) * 2 - 1 # Compression ratio compress = 2 / np.diff(t_vec)[seg] # Integrate current segment of trajectory traj = np.zeros((4, dims, len(t1))) for der_order in range(4): for i in range(dims): traj[der_order, i, :] = np.polyval(np.polyder(poly_coeffs[seg, :, i], der_order), t1) * (compress ** der_order) if seg < len(t_vec) - 2: # Remove last sample (will be the initial point of next segment) traj = traj[:, :, :-1] t_seg = tau_dt[:-1] else: t_seg = tau_dt full_traj = np.concatenate((full_traj, traj), axis=-1) t_total = np.concatenate((t_total, t_seg)) # Separate into p_xyz and yaw trajectories yaw_traj = full_traj[:, -1, :] full_traj = full_traj[:, :-1, :] return full_traj, yaw_traj, t_total def fit_multi_segment_polynomial_trajectory(p_targets, yaw_targets): p_targets = np.concatenate((p_targets, yaw_targets[np.newaxis, :]), 0) m = multiple_waypoints(p_targets.shape[1] - 1) dims = p_targets.shape[0] n_segments = p_targets.shape[1] poly_coefficients = np.zeros((n_segments - 1, 8, dims)) for dim in range(dims): b = rhs_generation(p_targets[dim, :]) poly_coefficients[:, :, dim] = np.fliplr(np.linalg.solve(m, b).reshape(n_segments - 1, 8)) return poly_coefficients def matrix_generation(ts): b = np.array([[1, ts, ts**2, ts**3, ts**4, ts**5, ts**6, ts**7], [0, 1, 2*ts, 3*ts**2, 4*ts**3, 5*ts**4, 6*ts**5, 7*ts**6], [0, 0, 2, 6*ts, 12*ts**2, 20*ts**3, 30*ts**4, 42*ts**5], [0, 0, 0, 6, 24*ts, 60*ts**2, 120*ts**3, 210*ts**4], [0, 0, 0, 0, 24, 120*ts, 360*ts**2, 840*ts**3], [0, 0, 0, 0, 0, 120, 720*ts, 2520*ts**2], [0, 0, 0, 0, 0, 0, 720, 5040*ts], [0, 0, 0, 0, 0, 0, 0, 5040]]) return b def multiple_waypoints(n_segments): m = np.zeros((8 * n_segments, 8 * n_segments)) for i in range(n_segments): if i == 0: # initial condition of the first curve b = matrix_generation(-1.0) m[8 * i:8 * i + 4, 8 * i:8 * i + 8] = b[:4, :] # intermediary condition of the first curve b = matrix_generation(1.0) m[8 * i + 4:8 * i + 7 + 4, 8 * i:8 * i + 8] = b[:-1, :] # starting condition of the second curve position and derivatives b = matrix_generation(-1.0) m[8 * i + 4 + 1:8 * i + 4 + 7, 8 * (i + 1):8 * (i + 1) + 8] = -b[1:-1, :] m[8 * i + 4 + 7:8 * i + 4 + 8, 8 * (i + 1):8 * (i + 1) + 8] = b[0, :] elif i != n_segments - 1: # starting condition of the ith curve position and derivatives b = matrix_generation(1.0) m[8 * i + 4:8 * i + 7 + 4, 8 * i:8 * i + 8] = b[:-1, :] # end condition of the ith curve position and derivatives b = matrix_generation(-1.0) m[8 * i + 4 + 1:8 * i + 4 + 7, 8 * (i + 1):8 * (i + 1) + 8] = -b[1:-1, :] m[8 * i + 4 + 7:8 * i + 4 + 8, 8 * (i + 1):8 * (i + 1) + 8] = b[0, :] if i == n_segments - 1: # end condition of the final curve position and derivatives (4 boundary conditions) b = matrix_generation(1.0) m[8 * i + 4:8 * i + 4 + 4, 8 * i:8 * i + 8] = b[:4, :] return m def fit_single_segment(p_start, p_end, v_start=None, v_end=None, a_start=None, a_end=None, j_start=None, j_end=None): if v_start is None: v_start = np.array([0, 0]) if v_end is None: v_end = np.array([0, 0]) if a_start is None: a_start = np.array([0, 0]) if a_end is None: a_end = np.array([0, 0]) if j_start is None: j_start = np.array([0, 0]) if j_end is None: j_end = np.array([0, 0]) poly_coefficients = np.zeros((8, len(p_start))) tf = 1 ti = -1 A = np.array(([ [1 * tf ** 7, 1 * tf ** 6, 1 * tf ** 5, 1 * tf ** 4, 1 * tf ** 3, 1 * tf ** 2, 1 * tf ** 1, 1], [7 * tf ** 6, 6 * tf ** 5, 5 * tf ** 4, 4 * tf ** 3, 3 * tf ** 2, 2 * tf ** 1, 1, 0], [42 * tf ** 5, 30 * tf ** 4, 20 * tf ** 3, 12 * tf ** 2, 6 * tf ** 1, 2, 0, 0], [210 * tf ** 4, 120 * tf ** 3, 60 * tf ** 2, 24 * tf ** 1, 6, 0, 0, 0], [1 * ti ** 7, 1 * ti ** 6, 1 * ti ** 5, 1 * ti ** 4, 1 * ti ** 3, 1 * ti ** 2, 1 * ti ** 1, 1], [7 * ti ** 6, 6 * ti ** 5, 5 * ti ** 4, 4 * ti ** 3, 3 * ti ** 2, 2 * ti ** 1, 1, 0], [42 * ti ** 5, 30 * ti ** 4, 20 * ti ** 3, 12 * ti ** 2, 6 * ti ** 1, 2, 0, 0], [210 * ti ** 4, 120 * ti ** 3, 60 * ti ** 2, 24 * ti ** 1, 6, 0, 0, 0]])) A = np.tile(A[:, :, np.newaxis], (1, 1, len(p_start))) b = np.concatenate((p_end, v_end, a_end, j_end, p_start, v_start, a_start, j_start)).reshape(8, -1) for i in range(len(p_start)): poly_coefficients[:, i] = np.linalg.inv(A[:, :, i]).dot(np.array(b[:, i])) return np.expand_dims(poly_coefficients, 0) def rhs_generation(x): n = x.shape[0] - 1 big_x = np.zeros((8 * n)) big_x[:4] = np.array([x[0], 0, 0, 0]).T big_x[-4:] = np.array([x[-1], 0, 0, 0]).T for i in range(1, n): big_x[8 * (i - 1) + 4:8 * (i - 1) + 8 + 4] = np.array([x[i], 0, 0, 0, 0, 0, 0, x[i]]).T return big_x ================================================ FILE: ros_dd_mpc/src/utils/utils.py ================================================ """ Miscellaneous utility functions. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import os import math import json import errno import shutil import joblib import random import pyquaternion import numpy as np import casadi as cs from sklearn import preprocessing from sklearn.decomposition import PCA from sklearn.cluster import KMeans from scipy.interpolate.interpolate import interp1d from config.configuration_parameters import DirectoryConfig as GPConfig import matplotlib.pyplot as plt import xml.etree.ElementTree as XMLtree def safe_mkdir_recursive(directory, overwrite=False): if not os.path.exists(directory): try: os.makedirs(directory) except OSError as exc: if exc.errno == errno.EEXIST and os.path.isdir(directory): pass else: raise else: if overwrite: try: shutil.rmtree(directory) except: print('Error while removing directory: {0}'.format(directory)) def safe_mknode_recursive(destiny_dir, node_name, overwrite): safe_mkdir_recursive(destiny_dir) if overwrite and os.path.exists(os.path.join(destiny_dir, node_name)): os.remove(os.path.join(destiny_dir, node_name)) if not os.path.exists(os.path.join(destiny_dir, node_name)): open(os.path.join(destiny_dir, node_name), 'w').close() #os.mknod(os.path.join(destiny_dir, node_name)) return False return True def jsonify(array): if isinstance(array, np.ndarray): return array.tolist() if isinstance(array, list): return array return array def undo_jsonify(array): x = [] for elem in array: a = elem.split('[')[1].split(']')[0].split(',') a = [float(num) for num in a] x = x + [a] return np.array(x) def get_data_dir_and_file(ds_name, training_split, params, read_only=False): """ Returns the directory and file name where to store the next simulation-based dataset. """ # Data folder directory rec_file_dir = GPConfig.DATA_DIR # Look for file f = [] d = [] for (dir_path, dir_names, file_names) in os.walk(rec_file_dir): f.extend(file_names) d.extend(dir_names) break split = "train" if training_split else "test" # Check if current dataset folder already exists. Create it otherwise if ds_name in d and os.path.exists(os.path.join(rec_file_dir, ds_name, split)): dataset_instances = [] for (_, _, file_names) in os.walk(os.path.join(rec_file_dir, ds_name, split)): dataset_instances.extend([os.path.splitext(file)[0] for file in file_names if not file.startswith('.')]) else: if read_only: return None safe_mkdir_recursive(os.path.join(rec_file_dir, ds_name, split)) dataset_instances = [] json_file_name = os.path.join(rec_file_dir, 'metadata.json') if not f: if read_only: return None # Metadata file not found -> make new one with open(json_file_name, 'w') as json_file: ds_instance_name = "dataset_001" json.dump({ds_name: {split: {ds_instance_name: params}}}, json_file, indent=4) else: # Metadata file existing with open(json_file_name) as json_file: metadata = json.load(json_file) # Check if current dataset name with data split exists in metadata: if ds_name in metadata.keys() and split in metadata[ds_name].keys(): existing_instance_idx = -1 for i, instance in enumerate(dataset_instances): if metadata[ds_name][split][instance] == params: existing_instance_idx = i if not read_only: print("This configuration already exists in the dataset with the same name.") if existing_instance_idx == -1: if read_only: return None if dataset_instances: # Dataset exists but this is a new instance existing_instances = [int(instance.split("_")[1]) for instance in dataset_instances] max_instance_number = max(existing_instances) ds_instance_name = "dataset_" + str(max_instance_number + 1).zfill(3) # Add to metadata dictionary the new instance metadata[ds_name][split][ds_instance_name] = params else: # Edge case where, for some error, there was something added to the metadata file but no actual # datasets were recorded. Remove entries from metadata and add them again. ds_instance_name = "dataset_001" metadata[ds_name][split] = {} metadata[ds_name][split][ds_instance_name] = params else: # Dataset exists and there is an instance with the same configuration ds_instance_name = dataset_instances[existing_instance_idx] else: if read_only: return None ds_instance_name = "dataset_001" safe_mkdir_recursive(os.path.join(rec_file_dir, ds_name, split)) # Add to metadata dictionary the new instance if ds_name in metadata.keys(): metadata[ds_name][split] = {ds_instance_name: params} else: metadata[ds_name] = {split: {ds_instance_name: params}} if not read_only: with open(json_file_name, 'w') as json_file: json.dump(metadata, json_file, indent=4) return os.path.join(rec_file_dir, ds_name, split), ds_instance_name + '.csv' def get_model_dir_and_file(model_options): directory = os.path.join(GPConfig.SAVE_DIR, str(model_options["git"]), str(model_options["model_name"])) model_params = model_options["params"] file_name = '' model_vars = list(model_params.keys()) model_vars.sort() for i, param in enumerate(model_vars): if i > 0: file_name += '__' file_name += 'no_' if not model_params[param] else '' file_name += param return directory, file_name def load_pickled_models(directory='', file_name='', model_options=None): """ Loads a pre-trained model from the specified directory, contained in a given pickle filename. Otherwise, if the model_options dictionary is given, use its contents to reconstruct the directory location of the pre-trained model fitting the requirements. :param directory: directory where the model file is located :param file_name: file name of the pre-trained model :param model_options: dictionary with the keys: "noisy" (bool), "drag" (bool), "git" (string), "training_samples" (int), "payload" (bool). :return: a dictionary with the recovered models from the pickle files. """ if model_options is not None: directory, file_name = get_model_dir_and_file(model_options) try: pickled_files = os.listdir(directory) except FileNotFoundError: return None loaded_models = [] try: for file in pickled_files: if not file.startswith(file_name) and file != 'feats.csv': continue if '.pt' in file: directory, file_name = get_model_dir_and_file(load_ops) saved_dict = torch.load(os.path.join(directory, f"{file_name}.pt")) mlp_model = dc.nn.MultiLayerPerceptron(saved_dict['input_size'], saved_dict['hidden_size'], saved_dict['output_size'], saved_dict['hidden_layers'], 'Tanh') model = NormalizedMLP(mlp_model, torch.tensor(np.zeros((saved_dict['input_size'],))).float(), torch.tensor(np.zeros((saved_dict['input_size'],))).float(), torch.tensor(np.zeros((saved_dict['output_size'],))).float(), torch.tensor(np.zeros((saved_dict['output_size'],))).float()) model.load_state_dict(saved_dict['state_dict']) model.eval() pre_trained_models = model loaded_models.append(joblib.load(os.path.join(directory, file))) continue if '.pkl' not in file and '.csv' not in file: continue if '.pkl' in file: loaded_models.append(joblib.load(os.path.join(directory, file))) except IsADirectoryError: raise FileNotFoundError("Tried to load file from directory %s, but it was not found." % directory) if loaded_models is not None: if loaded_models: pre_trained_models = {"models": loaded_models} else: pre_trained_models = None else: pre_trained_models = None return pre_trained_models def interpol_mse(t_1, x_1, t_2, x_2, n_interp_samples=1000): if np.all(t_1 == t_2): return np.mean(np.sqrt(np.sum((x_1 - x_2) ** 2, axis=1))) assert x_1.shape[1] == x_2.shape[1] t_min = max(t_1[0], t_2[0]) t_max = min(t_1[-1], t_2[-1]) t_interp = np.linspace(t_min, t_max, n_interp_samples) err = np.zeros((n_interp_samples, x_1.shape[1])) for dim in range(x_1.shape[1]): x1_interp = interp1d(t_1, x_1[:, dim], kind='cubic') x2_interp = interp1d(t_2, x_2[:, dim], kind='cubic') x1_sample = x1_interp(t_interp) x2_sample = x2_interp(t_interp) err[:, dim] = x1_sample - x2_sample return np.mean(np.sqrt(np.sum(err ** 2, axis=1))) def euclidean_dist(x, y, thresh=None): """ Measures the euclidean distance between points x and y. If a threshold value is provided, this function returns True if the calculated distance is smaller than the threshold, or False otherwise. If no threshold is provided, this function returns the computed distance. :param x: n-dimension point :param y: n-dimension point :param thresh: threshold :type thresh: float :return: If thresh is not None: whether x and y are closer to each other than the threshold. If thresh is None: the distance between x and y """ dist = np.sqrt(np.sum((x - y) ** 2)) if thresh is None: return dist return dist < thresh def euler_to_quaternion(roll, pitch, yaw): qx = np.sin(roll / 2) * np.cos(pitch / 2) * np.cos(yaw / 2) - np.cos(roll / 2) * np.sin(pitch / 2) * np.sin(yaw / 2) qy = np.cos(roll / 2) * np.sin(pitch / 2) * np.cos(yaw / 2) + np.sin(roll / 2) * np.cos(pitch / 2) * np.sin(yaw / 2) qz = np.cos(roll / 2) * np.cos(pitch / 2) * np.sin(yaw / 2) - np.sin(roll / 2) * np.sin(pitch / 2) * np.cos(yaw / 2) qw = np.cos(roll / 2) * np.cos(pitch / 2) * np.cos(yaw / 2) + np.sin(roll / 2) * np.sin(pitch / 2) * np.sin(yaw / 2) return np.array([qw, qx, qy, qz]) def quaternion_to_euler(q): q = pyquaternion.Quaternion(w=q[0], x=q[1], y=q[2], z=q[3]) yaw, pitch, roll = q.yaw_pitch_roll return [roll, pitch, yaw] def unit_quat(q): """ Normalizes a quaternion to be unit modulus. :param q: 4-dimensional numpy array or CasADi object :return: the unit quaternion in the same data format as the original one """ if isinstance(q, np.ndarray): # if (q == np.zeros(4)).all(): # q = np.array([1, 0, 0, 0]) q_norm = np.sqrt(np.sum(q ** 2)) else: q_norm = cs.sqrt(cs.sumsqr(q)) return 1 / q_norm * q def v_dot_q(v, q): rot_mat = q_to_rot_mat(q) if isinstance(q, np.ndarray): return rot_mat.dot(v) return cs.mtimes(rot_mat, v) def q_to_rot_mat(q): qw, qx, qy, qz = q[0], q[1], q[2], q[3] if isinstance(q, np.ndarray): rot_mat = np.array([ [1 - 2 * (qy ** 2 + qz ** 2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)], [2 * (qx * qy + qw * qz), 1 - 2 * (qx ** 2 + qz ** 2), 2 * (qy * qz - qw * qx)], [2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx ** 2 + qy ** 2)]]) else: rot_mat = cs.vertcat( cs.horzcat(1 - 2 * (qy ** 2 + qz ** 2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)), cs.horzcat(2 * (qx * qy + qw * qz), 1 - 2 * (qx ** 2 + qz ** 2), 2 * (qy * qz - qw * qx)), cs.horzcat(2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx ** 2 + qy ** 2))) return rot_mat def q_dot_q(q, r): """ Applies the rotation of quaternion r to quaternion q. In order words, rotates quaternion q by r. Quaternion format: wxyz. :param q: 4-length numpy array or CasADi MX. Initial rotation :param r: 4-length numpy array or CasADi MX. Applied rotation :return: The quaternion q rotated by r, with the same format as in the input. """ qw, qx, qy, qz = q[0], q[1], q[2], q[3] rw, rx, ry, rz = r[0], r[1], r[2], r[3] t0 = rw * qw - rx * qx - ry * qy - rz * qz t1 = rw * qx + rx * qw - ry * qz + rz * qy t2 = rw * qy + rx * qz + ry * qw - rz * qx t3 = rw * qz - rx * qy + ry * qx + rz * qw if isinstance(q, np.ndarray): return np.array([t0, t1, t2, t3]) else: return cs.vertcat(t0, t1, t2, t3) def rotation_matrix_to_quat(rot): """ Calculate a quaternion from a 3x3 rotation matrix. :param rot: 3x3 numpy array, representing a valid rotation matrix :return: a quaternion corresponding to the 3D rotation described by the input matrix. Quaternion format: wxyz """ q = pyquaternion.Quaternion(matrix=rot) return np.array([q.w, q.x, q.y, q.z]) def undo_quaternion_flip(q_past, q_current): """ Detects if q_current generated a quaternion jump and corrects it. Requires knowledge of the previous quaternion in the series, q_past :param q_past: 4-dimensional vector representing a quaternion in wxyz form. :param q_current: 4-dimensional vector representing a quaternion in wxyz form. Will be corrected if it generates a flip wrt q_past. :return: q_current with the flip removed if necessary """ if np.sqrt(np.sum((q_past - q_current) ** 2)) > np.sqrt(np.sum((q_past + q_current) ** 2)): return -q_current return q_current def skew_symmetric(v): """ Computes the skew-symmetric matrix of a 3D vector (PAMPC version) :param v: 3D numpy vector or CasADi MX :return: the corresponding skew-symmetric matrix of v with the same data type as v """ if isinstance(v, np.ndarray): return np.array([[0, -v[0], -v[1], -v[2]], [v[0], 0, v[2], -v[1]], [v[1], -v[2], 0, v[0]], [v[2], v[1], -v[0], 0]]) return cs.vertcat( cs.horzcat(0, -v[0], -v[1], -v[2]), cs.horzcat(v[0], 0, v[2], -v[1]), cs.horzcat(v[1], -v[2], 0, v[0]), cs.horzcat(v[2], v[1], -v[0], 0)) def decompose_quaternion(q): """ Decomposes a quaternion into a z rotation and an xy rotation :param q: 4-dimensional numpy array of CasADi MX (format qw, qx, qy, qz) :return: two 4-dimensional arrays (same format as input), where the first contains the xy rotation and the second the z rotation, in quaternion forms. """ w, x, y, z = q[0], q[1], q[2], q[3] if isinstance(q, cs.MX): qz = unit_quat(cs.vertcat(w, 0, 0, z)) else: qz = unit_quat(np.array([w, 0, 0, z])) qxy = q_dot_q(q, quaternion_inverse(qz)) return qxy, qz def quaternion_inverse(q): w, x, y, z = q[0], q[1], q[2], q[3] if isinstance(q, np.ndarray): return np.array([w, -x, -y, -z]) else: return cs.vertcat(w, -x, -y, -z) def rotation_matrix_to_euler(r_mat): sy = math.sqrt(r_mat[0, 0] * r_mat[0, 0] + r_mat[1, 0] * r_mat[1, 0]) singular = sy < 1e-6 if not singular: x = math.atan2(r_mat[2, 1], r_mat[2, 2]) y = math.atan2(-r_mat[2, 0], sy) z = math.atan2(r_mat[1, 0], r_mat[0, 0]) else: x = math.atan2(-r_mat[1, 2], r_mat[1, 1]) y = math.atan2(-r_mat[2, 0], sy) z = 0 return np.array([x, y, z]) def prune_dataset(x, y, x_cap, bins, thresh, plot, labels=None): """ Prunes the collected model error dataset with two filters. First, remove values where the input values (velocities) exceed 10. Second, create an histogram for each of the three axial velocity errors (y) with the specified number of bins and remove any data where the total amount of samples in that bin is less than the specified threshold ratio. :param x: dataset of input GP features (velocities). Dimensions N x n (N entries and n dimensions) :param y: dataset of errors. Dimensions N x m (N entries and m dimensions) :param x_cap: remove values from dataset if x > x_cap or x < -x_cap :param bins: number of bins used for histogram :param thresh: threshold ratio below which data from that bin will be removed :param plot: make a plot of the pruning :param labels: Labels to use for the plot :return: The indices kept after the pruning """ n_bins = bins original_length = x.shape[0] plot_bins = [] if plot: plt.figure() for i in range(y.shape[1]): plt.subplot(y.shape[1] + 1, 1, i + 1) h, bins = np.histogram(y[:, i], bins=n_bins) plot_bins.append(bins) plt.bar(bins[:-1], h, np.ones_like(h) * (bins[1] - bins[0]), align='edge', label='discarded') if labels is not None: plt.ylabel(labels[i]) pruned_idx_unique = np.zeros(0, dtype=int) # Prune velocities (max axial velocity = x_cap m/s). if x_cap is not None: for i in range(x.shape[1]): pruned_idx = np.where(np.abs(x[:, i]) > x_cap)[0] pruned_idx_unique = np.unique(np.append(pruned_idx, pruned_idx_unique)) # Prune by error histogram dimension wise (discard bins with less than 1% of the data) for i in range(y.shape[1]): h, bins = np.histogram(y[:, i], bins=n_bins) for j in range(len(h)): if h[j] / np.sum(h) < thresh: pruned_idx = np.where(np.logical_and(bins[j + 1] >= y[:, i], y[:, i] >= bins[j])) pruned_idx_unique = np.unique(np.append(pruned_idx, pruned_idx_unique)) y_norm = np.sqrt(np.sum(y ** 2, 1)) # Prune by error histogram norm h, error_bins = np.histogram(y_norm, bins=n_bins) h = h / np.sum(h) if plot: plt.subplot(y.shape[1] + 1, 1, y.shape[1] + 1) plt.bar(error_bins[:-1], h, np.ones_like(h) * (error_bins[1] - error_bins[0]), align='edge', label='discarded') for j in range(len(h)): if h[j] / np.sum(h) < thresh: pruned_idx = np.where(np.logical_and(error_bins[j + 1] >= y_norm, y_norm >= error_bins[j])) pruned_idx_unique = np.unique(np.append(pruned_idx, pruned_idx_unique)) y = np.delete(y, pruned_idx_unique, axis=0) if plot: for i in range(y.shape[1]): plt.subplot(y.shape[1] + 1, 1, i + 1) h, bins = np.histogram(y[:, i], bins=plot_bins[i]) plt.bar(bins[:-1], h, np.ones_like(h) * (bins[1] - bins[0]), align='edge', label='kept') plt.legend() plt.subplot(y.shape[1] + 1, 1, y.shape[1] + 1) h, bins = np.histogram(np.sqrt(np.sum(y ** 2, 1)), bins=error_bins) h = h / sum(h) plt.bar(bins[:-1], h, np.ones_like(h) * (bins[1] - bins[0]), align='edge', label='kept') plt.ylabel('Error norm') kept_idx = np.delete(np.arange(0, original_length), pruned_idx_unique) return kept_idx def distance_maximizing_points_1d(points, n_train_points, dense_gp=None): """ Heuristic function for sampling training points in 1D (one input feature and one output prediction dimensions) :param points: dataset points for the current cluster. Array of shape Nx1 :param n_train_points: Integer. number of training points to sample. :param dense_gp: A GP object to sample the points from, or None of the points will be taken directly from the data. :return: """ closest_points = np.zeros(n_train_points, dtype=int if dense_gp is None else float) if dense_gp is not None: n_train_points -= 1 # Fit histogram in data with as many bins as the number of training points a, b = np.histogram(points, bins=n_train_points) hist_indices = np.digitize(points, b) - 1 # Pick as training value the median or mean value of each bin for i in range(n_train_points): bin_values = points[np.where(hist_indices == i)] if len(bin_values) < 1: closest_points[i] = np.random.choice(np.arange(len(points)), 1) continue if divmod(len(bin_values), 2)[1] == 0: bin_values = bin_values[:-1] if dense_gp is None: # If no dense GP, sample median points in each bin from training set bin_median = np.median(bin_values) median_point_id = np.where(points == bin_median)[0] if len(median_point_id) > 1: closest_points[i] = median_point_id[0] else: closest_points[i] = median_point_id else: # If with GP, sample mean points in each bin from GP bin_mean = np.min(bin_values) closest_points[i] = bin_mean if dense_gp is not None: # Add dimension axis 0 closest_points[-1] = np.max(points) closest_points = closest_points[np.newaxis, :] return closest_points def distance_maximizing_points_2d(points, n_train_points, dense_gp, plot=False): if n_train_points > 30: n_clusters = max(int(n_train_points / 10), 30) n_samples = int(np.floor(n_train_points / n_clusters)) else: n_clusters = n_train_points n_samples = 1 kmeans = KMeans(n_clusters).fit_predict(points) closest_points = [] for i in range(n_clusters): closest_points += np.random.choice(np.where(kmeans == i)[0], n_samples).tolist() # Remove exceeding points for _ in range(len(closest_points) - n_train_points): rnd_item = random.choice(closest_points) closest_points.remove(rnd_item) closest_points = np.array(closest_points) if plot: plt.figure() for i in range(n_clusters): cluster_points = points[np.where(kmeans == i)] plt.scatter(cluster_points[:, 0], cluster_points[:, 1]) plt.scatter(points[closest_points, 0], points[closest_points, 1], marker='*', facecolors='w', edgecolors='k', s=100, label='selected') plt.legend() plt.show() if dense_gp is None: return closest_points closest_points = points[closest_points].T return closest_points def distance_maximizing_points(x_values, center, n_train_points=7, dense_gp=None, plot=False): if x_values.shape[1] == 1: return distance_maximizing_points_1d(x_values, n_train_points, dense_gp) if x_values.shape[1] >= 2: return distance_maximizing_points_2d(x_values, n_train_points, dense_gp, plot) # Compute PCA of data to find variability maximizing axes pca = PCA(n_components=3) pca.fit(x_values) pca_axes = pca.components_ data_center = center # Apply PCA transformation points_pca = (x_values - data_center).dot(pca_axes.T) center = (center - data_center).dot(pca_axes.T) # Compute the corners of the cube containing the data in the PCA space p_min = center - (center - np.min(points_pca, 0)) p_max = (np.max(points_pca, 0) - center) + center centroids = np.array([[center[0], center[1], center[2]]]) pyramids = np.array([[p_max[0], center[1], center[2]], [center[0], p_max[1], center[2]], [center[0], center[1], p_max[2]], [p_min[0], center[1], center[2]], [center[0], p_min[1], center[2]], [center[0], center[1], p_min[2]]]) cuboid = np.array([[p_max[0], p_max[1], p_max[2]], [p_max[0], p_max[1], p_min[2]], [p_max[0], p_min[1], p_max[2]], [p_max[0], p_min[1], p_min[2]], [p_min[0], p_max[1], p_max[2]], [p_min[0], p_max[1], p_min[2]], [p_min[0], p_min[1], p_max[2]], [p_min[0], p_min[1], p_min[2]]]) if n_train_points >= 15: centroids = np.concatenate((centroids, pyramids, cuboid), axis=0) elif n_train_points >= 9: centroids = np.concatenate((centroids, cuboid), axis=0) elif n_train_points >= 7: centroids = np.concatenate((centroids, pyramids), axis=0) else: centroids = centroids if dense_gp is None: closest_points = np.ones(centroids.shape[0], dtype=int) * -1 # Find the closest points to all centroids for i in range(centroids.shape[0]): centroid = centroids[i, :] dist = np.sqrt(np.sum((points_pca - centroid) ** 2, 1)) closest_point = int(np.argmin(dist)) # Assert no repeated points. If repeated, find next closest one and check again while closest_point in closest_points: dist[closest_point] = np.inf closest_point = int(np.argmin(dist)) closest_points[i] = closest_point # Convert centroids to data space centroids_ = centroids.dot(pca_axes) + data_center if not plot: return closest_points fig = plt.figure() ax = fig.add_subplot(122, projection='3d') ax.scatter(points_pca[:, 0], points_pca[:, 1], points_pca[:, 2], 'b', label='data') ax.scatter(centroids[0, 0], centroids[0, 1], centroids[0, 2], s=50, label='center') ax.scatter(centroids[1:, 0], centroids[1:, 1], centroids[1:, 2], s=50, label='centroids') ax.scatter(points_pca[closest_points, 0], points_pca[closest_points, 1], points_pca[closest_points, 2], s=50, label='selected') ax.set_title('PCA space') ax.legend() ax = fig.add_subplot(121, projection='3d') ax.scatter(x_values[:, 0], x_values[:, 1], x_values[:, 2], 'b', label='data') closest_points_x = x_values[closest_points] ax.scatter(centroids_[0, 0], centroids_[0, 1], centroids_[0, 2], s=50, label='center') ax.scatter(centroids_[1:, 0], centroids_[1:, 1], centroids_[1:, 2], s=50, label='centroids') ax.scatter(closest_points_x[:, 0], closest_points_x[:, 1], closest_points_x[:, 2], s=50, label='selected') ax.set_title('Data space') ax.legend() plt.show() # Returns the indices of the closest points from the dataset to the ideal ones return closest_points else: # TODO: Use GP covariance for sampling # Convert centroids to data space centroids_ = centroids.dot(pca_axes) + data_center return centroids_.T def sample_random_points(points, used_idx, points_to_sample, dense_gp=None): bins = min(10, int(len(points) / 10)) bins = max(bins, 2) # Add remaining points as random points free_points = np.arange(0, points.shape[0], 1) gp_i_free_points = np.delete(free_points, used_idx) n_samples = min(points_to_sample, len(gp_i_free_points)) # Compute histogram of data a, b = np.histogramdd(points[gp_i_free_points, :], bins) assignments = [np.minimum(np.digitize(points[gp_i_free_points, j], bins=b[j]) - 1, bins - 1) for j in range(points.shape[1])] # Compute probability distribution based on inverse histogram probs = np.max(a) - a[tuple(assignments)] probs = probs / sum(probs) try: gp_i_free_points = np.random.choice(gp_i_free_points, n_samples, p=probs, replace=False) except ValueError: print('a') used_idx = np.append(used_idx, gp_i_free_points) return used_idx def parse_xacro_file(xacro): """ Reads a .xacro file describing a robot for Gazebo and returns a dictionary with its properties. :param xacro: full path of .xacro file to read :return: a dictionary of robot attributes """ tree = XMLtree.parse(xacro) attrib_dict = {} for node in tree.getroot().getchildren(): # Get attributes attributes = node.attrib if 'value' in attributes.keys(): attrib_dict[attributes['name']] = attributes['value'] if node.getchildren(): try: attrib_dict[attributes['name']] = [child.attrib for child in node.getchildren()] except: continue return attrib_dict def make_bx_matrix(x_dims, y_feats): """ Generates the Bx matrix for the GP augmented MPC. :param x_dims: dimensionality of the state vector :param y_feats: array with the indices of the state vector x that are augmented by the GP regressor :return: The Bx matrix to map the GP output to the state space. """ bx = np.zeros((x_dims, len(y_feats))) for i in range(len(y_feats)): bx[y_feats[i], i] = 1 return bx def make_bz_matrix(x_dims, u_dims, x_feats, u_feats): """ Generates the Bz matrix for the GP augmented MPC. :param x_dims: dimensionality of the state vector :param u_dims: dimensionality of the input vector :param x_feats: array with the indices of the state vector x used to make the first part of the GP feature vector z :param u_feats: array with the indices of the input vector u used to make the second part of the GP feature vector z :return: The Bz matrix to map from input x and u features to the z feature vector. """ bz = np.zeros((len(x_feats), x_dims)) for i in range(len(x_feats)): bz[i, x_feats[i]] = 1 bzu = np.zeros((len(u_feats), u_dims)) for i in range(len(u_feats)): bzu[i, u_feats[i]] = 1 bz = np.concatenate((bz, np.zeros((len(x_feats), u_dims))), axis=1) bzu = np.concatenate((np.zeros((len(u_feats), x_dims)), bzu), axis=1) bz = np.concatenate((bz, bzu), axis=0) return bz def quaternion_state_mse(x, x_ref, mask): """ Calculates the MSE of the 13-dimensional state (p_xyz, q_wxyz, v_xyz, r_xyz) wrt. the reference state. The MSE of the quaternions are treated axes-wise. :param x: 13-dimensional state :param x_ref: 13-dimensional reference state :param mask: 12-dimensional masking for weighted MSE (p_xyz, q_xyz, v_xyz, r_xyz) :return: the mean squared error of both """ q_error = q_dot_q(x[3:7], quaternion_inverse(x_ref[3:7])) e = np.concatenate((x[:3] - x_ref[:3], q_error[1:], x[7:10] - x_ref[7:10], x[10:] - x_ref[10:])) return np.sqrt((e * np.array(mask)).dot(e)) def separate_variables(traj): """ Reshapes a trajectory into expected format. :param traj: N x 13 array representing the reference trajectory :return: A list with the components: Nx3 position trajectory array, Nx4 quaternion trajectory array, Nx3 velocity trajectory array, Nx3 body rate trajectory array """ p_traj = traj[:, :3] a_traj = traj[:, 3:7] v_traj = traj[:, 7:10] r_traj = traj[:, 10:] return [p_traj, a_traj, v_traj, r_traj] ================================================ FILE: ros_dd_mpc/src/utils/visualization.py ================================================ """ Miscellaneous visualization functions. This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import json # import tikzplotlib import matplotlib.pyplot as plt import numpy as np import matplotlib.animation as animation from matplotlib.colors import LinearSegmentedColormap, BoundaryNorm from matplotlib.colorbar import ColorbarBase from matplotlib import cm from mpl_toolkits.mplot3d import Axes3D from config.configuration_parameters import DirectoryConfig as PathConfig from src.utils.utils import v_dot_q, quaternion_to_euler, quaternion_inverse, q_dot_q, safe_mknode_recursive, \ safe_mkdir_recursive import os def angle_to_rot_mat(angle): """ Computes the 2x2 rotation matrix from the scalar angle :param angle: scalar angle in radians :return: the corresponding 2x2 rotation matrix """ s = np.sin(angle) c = np.cos(angle) return np.array([[c, -s], [s, c]]) def draw_arrow(x_base, y_base, x_body, y_body): """ Returns the coordinates for drawing a 2D arrow given its origin point and its length. :param x_base: x coordinate of the arrow origin :param y_base: y coordinate of the arrow origin :param x_body: x length of the arrow :param y_body: y length of the arrow :return: a tuple of x, y coordinates to plot the arrow """ len_arrow = np.sqrt(x_body ** 2 + y_body ** 2) beta = np.arctan2(y_body, x_body) beta_rot = angle_to_rot_mat(beta) lower_arrow = beta_rot.dot(np.array([[-np.cos(np.pi / 6)], [-np.sin(np.pi / 6)]]) * len_arrow / 3) upper_arrow = beta_rot.dot(np.array([[-np.cos(np.pi / 6)], [np.sin(np.pi / 6)]]) * len_arrow / 3) return ([x_base, x_base + x_body, x_base + x_body + lower_arrow[0, 0], x_base + x_body, x_base + x_body + upper_arrow[0, 0]], [y_base, y_base + y_body, y_base + y_body + lower_arrow[1, 0], y_base + y_body, y_base + y_body + upper_arrow[1, 0]]) def draw_drone(pos, q_rot, x_f, y_f): # Define quadrotor extremities in body reference frame x1 = np.array([x_f[0], y_f[0], 0]) x2 = np.array([x_f[1], y_f[1], 0]) x3 = np.array([x_f[2], y_f[2], 0]) x4 = np.array([x_f[3], y_f[3], 0]) # Convert to world reference frame and add quadrotor center point: x1 = v_dot_q(x1, q_rot) + pos x2 = v_dot_q(x2, q_rot) + pos x3 = v_dot_q(x3, q_rot) + pos x4 = v_dot_q(x4, q_rot) + pos # Build set of coordinates for plotting return ([x1[0], x3[0], pos[0], x2[0], x4[0]], [x1[1], x3[1], pos[1], x2[1], x4[1]], [x1[2], x3[2], pos[2], x2[2], x4[2]]) def draw_covariance_ellipsoid(center, covar): """ :param center: 3-dimensional array. Center of the ellipsoid :param covar: 3x3 covariance matrix. If the covariance is diagonal, the ellipsoid will have radii equal to the three diagonal axis along axes x, y, z respectively. :return: """ # find the rotation matrix and radii of the axes _, radii, rotation = np.linalg.svd(covar) # now carry on with EOL's answer u = np.linspace(0.0, 2.0 * np.pi, 20) v = np.linspace(0.0, np.pi, 20) x = radii[0] * np.outer(np.cos(u), np.sin(v)) y = radii[1] * np.outer(np.sin(u), np.sin(v)) z = radii[2] * np.outer(np.ones_like(u), np.cos(v)) for i in range(len(x)): for j in range(len(x)): [x[i, j], y[i, j], z[i, j]] = np.dot([x[i, j], y[i, j], z[i, j]], rotation) + center x = np.reshape(x, -1) y = np.reshape(y, -1) z = np.reshape(z, -1) return x, y, z def visualize_data_distribution(x_data, y_data, clusters, x_pruned, y_pruned): """ Visualizes the distribution of the training dataset and the assignation of the GP prediction clusters. :param x_data: numpy array of shape N x 3, where N is the number of training points. Feature variables. :param y_data: numpy array of shape N x 3, where N is the number of training points. Regressed variables. :param x_pruned: numpy array of shape M x 3, where M is the number of pruned training points. Feature variables. :param y_pruned: numpy array of shape M x 3, where M is the number of pruned training points. Regressed variables. :param clusters: A dictionary where each entry is indexed by the cluster number, and contains a list of all the indices of the points in x_pruned belonging to that cluster. """ if x_data.shape[1] < 3: return fig = plt.figure() ax = fig.add_subplot(131, projection='3d') c = np.sqrt(np.sum(y_data ** 2, 1)) scatter = ax.scatter(x_data[:, 0], x_data[:, 1], x_data[:, 2], c=c, alpha=0.6) ax.set_title('Raw data: Correction magnitude') ax.set_xlabel(r'$v_x\: [m/s]$') ax.set_ylabel(r'$v_y\: [m/s]$') ax.set_zlabel(r'$v_z\: [m/s]$') fig.colorbar(scatter, ax=ax, orientation='vertical', shrink=0.75) ax = fig.add_subplot(132, projection='3d') c = np.sqrt(np.sum(y_pruned ** 2, 1)) scatter = ax.scatter(x_pruned[:, 0], x_pruned[:, 1], x_pruned[:, 2], c=c, alpha=0.6) ax.set_title('Pruned data: Correction magnitude') ax.set_xlabel(r'$v_x\: [m/s]$') ax.set_ylabel(r'$v_y\: [m/s]$') ax.set_zlabel(r'$v_z\: [m/s]$') fig.colorbar(scatter, ax=ax, orientation='vertical', shrink=0.75) n_clusters = len(clusters.keys()) ax = fig.add_subplot(133, projection='3d') for i in range(int(n_clusters)): ax.scatter(x_pruned[clusters[i], 0], x_pruned[clusters[i], 1], x_pruned[clusters[i], 2], alpha=0.6) ax.set_title('Cluster assignations') ax.set_xlabel(r'$v_x\: [m/s]$') ax.set_ylabel(r'$v_y\: [m/s]$') ax.set_zlabel(r'$v_z\: [m/s]$') plt.show() def visualize_gp_inference(x_data, u_data, y_data, gp_ensemble, vis_features_x, y_dims, labels): # WARNING: This function is extremely limited to the case where the regression is performed using just the # velocity state as input features and as output dimensions. predictions = gp_ensemble.predict(x_data.T, u_data.T) predictions = np.atleast_2d(np.atleast_2d(predictions["pred"])[y_dims]) if len(vis_features_x) > 1: y_pred = np.sqrt(np.sum(predictions ** 2, 0)) y_mse = np.sqrt(np.sum(y_data ** 2, 1)) else: y_pred = predictions[0, :] y_mse = y_data[:, 0] v_min = min(np.min(y_pred), np.min(y_mse)) v_max = max(np.max(y_pred), np.max(y_mse)) fig = plt.figure() font_size = 16 if len(vis_features_x) == 1: # Feature dimension is only 1 # Compute windowed average n_bins = 20 _, b = np.histogram(x_data[:, vis_features_x], bins=n_bins) hist_indices = np.digitize(x_data[:, vis_features_x], b) win_average = np.zeros(n_bins) for i in range(n_bins): win_average[i] = np.mean(y_mse[np.where(hist_indices == i + 1)[0]]) bin_midpoints = b[:-1] + np.diff(b)[0] / 2 ax = [fig.add_subplot(121), fig.add_subplot(122)] ax[0].scatter(x_data[:, vis_features_x], y_mse) ax[0].set_xlabel(labels[0]) ax[0].set_ylabel('RMSE') ax[0].set_title('Post-processed dataset') ax[1].scatter(x_data[:, vis_features_x], y_pred, label='GP') ax[1].plot(bin_midpoints, win_average, label='window average') ax[1].set_xlabel(labels[0]) ax[1].set_title('Predictions') ax[1].legend() return elif len(vis_features_x) >= 3: ax = [fig.add_subplot(121, projection='3d'), fig.add_subplot(122, projection='3d')] im = ax[0].scatter(x_data[:, vis_features_x[0]], x_data[:, vis_features_x[1]], x_data[:, vis_features_x[2]], c=y_mse, cmap='viridis', alpha=0.6, vmin=v_min, vmax=v_max) ax[0].set_xlabel(labels[0], size=font_size - 4, labelpad=10) ax[0].set_ylabel(labels[1], size=font_size - 4, labelpad=10) ax[0].set_zlabel(labels[2], size=font_size - 4, labelpad=10) ax[0].set_title(r'Nominal MPC error $\|\mathbf{a}^e\|$', size=font_size) ax[0].view_init(65, 15) ax[1].scatter(x_data[:, vis_features_x[0]], x_data[:, vis_features_x[1]], x_data[:, vis_features_x[2]], c=y_pred, cmap='viridis', alpha=0.6, vmin=v_min, vmax=v_max) ax[1].set_xlabel(labels[0], size=font_size - 4, labelpad=10) ax[1].set_ylabel(labels[1], size=font_size - 4, labelpad=10) ax[1].set_zlabel(labels[2], size=font_size - 4, labelpad=10) ax[1].set_title(r'GP prediction mangnitude $\|\tilde{\mathbf{a}}^e\|$', size=font_size) ax[1].view_init(65, 15) plt.tight_layout() fig.subplots_adjust(right=0.85) cbar = fig.add_axes([0.90, 0.05, 0.03, 0.8]) fig.colorbar(im, cax=cbar) cbar.get_yaxis().labelpad = 15 cbar.set_ylabel(r'$\|\mathbf{a}^e\|\left[\frac{m}{s^2}\right]$', size=font_size, labelpad=20, rotation=270) cbar.tick_params(labelsize=font_size - 4) # Create values for the regressed variables x = np.linspace(min(x_data[:, vis_features_x[0]]), max(x_data[:, vis_features_x[0]]), 100) y = np.linspace(min(x_data[:, vis_features_x[1]]), max(x_data[:, vis_features_x[1]]), 100) # x = np.linspace(-8, 8, 50) # y = np.linspace(-8, 8, 50) x_mesh, y_mesh = np.meshgrid(x, y) x = np.reshape(x_mesh, (-1, 1)) y = np.reshape(y_mesh, (-1, 1)) z = np.zeros_like(x) x_sample = np.concatenate((x, y, z), 1) # Generate complete mock x features. Only vis_features are non-zero x_mock = np.tile(np.zeros_like(z), (1, x_data.shape[1])) x_mock[:, np.array(vis_features_x)] = x_sample # Also create mock u features u_mock = np.tile(np.zeros_like(z), (1, u_data.shape[1])) if len(vis_features_x) != 3: plt.show() return # Generate animated plot showing prediction of the multiple clusters. # Cluster coloring only possible if all the output dimensions have exactly the same clusters. fig = plt.figure() ax = fig.add_subplot(111, projection='3d') print("Grid sampling...") outs = gp_ensemble.predict(x_mock.T, u_mock.T, return_gp_id=True, progress_bar=True) y_pred = np.atleast_2d(np.atleast_2d(outs["pred"])[y_dims]) gp_ids = outs["gp_id"] y_sample = np.sqrt(np.sum(y_pred ** 2, 0)) y_sample = np.reshape(y_sample, x_mesh.shape) gp_ids = np.reshape(gp_ids[next(iter(gp_ids))], x_mesh.shape) def init(): # create the new map cmap = cm.get_cmap('jet') cmaplist = [cmap(j) for j in range(cmap.N)] cmap = LinearSegmentedColormap.from_list('Custom cmap', cmaplist, cmap.N) # define the bins and normalize capped_n_clusters = min(np.amax(gp_ids) + 2, 20) bounds = np.linspace(0, np.amax(gp_ids) + 1, capped_n_clusters) norm = BoundaryNorm(bounds, cmap.N) my_col = cm.get_cmap('jet')(gp_ids / (np.amax(gp_ids) + 1)) ax.plot_surface(x_mesh, y_mesh, y_sample, facecolors=my_col, linewidth=0, rstride=1, cstride=1, antialiased=False, alpha=0.7, cmap=cmap, norm=norm) ax2 = fig.add_axes([0.90, 0.2, 0.03, 0.6]) ColorbarBase(ax2, cmap=cmap, norm=norm, spacing='proportional', ticks=bounds, boundaries=bounds, format='%1i') ax2.set_ylabel('Cluster assignment ID', size=14) ax2.tick_params(labelsize=16) ax.tick_params(labelsize=14) ax.set_xlabel(labels[0], size=16, labelpad=10) ax.set_ylabel(labels[1], size=16, labelpad=10) ax.set_zlabel(r'$\|\tilde{\mathbf{a}}^e\|\: \left[\frac{m}{s^2}\right]$', size=16, labelpad=10) ax.set_title(r'GP correction. Slice $v_z=0 \:\: \left[\frac{m}{s}\right]$', size=18) return fig, def animate(i): ax.view_init(elev=30., azim=i*3) return fig _ = animation.FuncAnimation(fig, animate, init_func=init, frames=360, interval=20, blit=False) plt.show() def initialize_drone_plotter(world_rad, quad_rad, n_props, full_traj=None): fig = plt.figure(figsize=(10, 10), dpi=96) fig.show() mng = plt.get_current_fig_manager() mng.resize(*mng.window.maxsize()) ax = fig.add_subplot(111, projection='3d') if full_traj is not None: ax.plot(full_traj[:, 0], full_traj[:, 1], full_traj[:, 2], '--', color='tab:blue', alpha=0.5) ax.set_xlim([ax.get_xlim()[0] - 2 * quad_rad, ax.get_xlim()[1] + 2 * quad_rad]) ax.set_ylim([ax.get_ylim()[0] - 2 * quad_rad, ax.get_ylim()[1] + 2 * quad_rad]) ax.set_zlim([ax.get_zlim()[0] - 2 * quad_rad, ax.get_zlim()[1] + 2 * quad_rad]) else: ax.set_xlim([-world_rad, world_rad]) ax.set_ylim([-world_rad, world_rad]) ax.set_zlim([-world_rad, world_rad]) ax.set_xlabel('x [m]') ax.set_ylabel('y [m]') ax.set_zlabel('z [m]') fig.canvas.draw() plt.draw() # cache the background background = fig.canvas.copy_from_bbox(ax.bbox) artists = { "trajectory": ax.plot([], [])[0], "drone": ax.plot([], [], 'o-')[0], "drone_x": ax.plot([], [], 'o-', color='r')[0], "missing_targets": ax.plot([], [], [], color='r', marker='o', linestyle='None', markersize=12)[0], "reached_targets": ax.plot([], [], [], color='g', marker='o', linestyle='None', markersize=12)[0], "sim_trajectory": [ax.plot([], [], [], '-', color='tab:blue', alpha=0.9 - i * 0.2 / n_props)[0] for i in range(n_props)], "int_trajectory": [ax.plot([], [], [], '-', color='tab:orange', alpha=0.9 - i * 0.5 / n_props)[0] for i in range(n_props + 1)], "prop_trajectory": [ax.plot([], [], [], '-', color='tab:red', alpha=0.9 - i * 0.2 / n_props)[0] for i in range(n_props)], "prop_covariance": [ax.plot([], [], [], color='r', alpha=0.5 - i * 0.45 / n_props)[0] for i in range(n_props)], "projection_lines": [ax.plot([], [], [], '-', color='tab:blue', alpha=0.2)[0], ax.plot([], [], [], '-', color='tab:blue', alpha=0.2)[0]], "projection_target": [ax.plot([], [], [], marker='o', color='r', linestyle='None', alpha=0.2)[0], ax.plot([], [], [], marker='o', color='r', linestyle='None', alpha=0.2)[0]]} art_pack = fig, ax, artists, background, world_rad return art_pack def draw_drone_simulation(art_pack, x_trajectory, quad, targets, targets_reached, sim_traj=None, int_traj=None, pred_traj=None, x_pred_cov=None, follow_quad=False): fig, ax, artists, background, world_rad = art_pack trajectories_artist = artists["trajectory"] if "trajectory" in artists.keys() else [] projected_traj_artists = artists["projection_lines"] if "projection_lines" in artists.keys() else [] drone_sketch_artist = artists["drone"] if "drone" in artists.keys() else [] drone_sketch_artist_x_motor = artists["drone_x"] if "drone_x" in artists.keys() else[] targets_artist = artists["missing_targets"] if "missing_targets" in artists.keys() else [] reached_targets_artist = artists["reached_targets"] if "reached_targets" in artists.keys() else [] projected_tar_artists = artists["projection_target"] if "projection_target" in artists.keys() else [] sim_traj_artists = artists["sim_trajectory"] if "sim_trajectory" in artists.keys() else [] int_traj_artists = artists["int_trajectory"] if "int_trajectory" in artists.keys() else [] pred_traj_artists = artists["prop_trajectory"] if "prop_trajectory" in artists.keys() else [] cov_artists = artists["prop_covariance"] if "prop_covariance" in artists.keys() else [] # restore background fig.canvas.restore_region(background) def draw_fading_traj(traj, traj_artists): traj = np.squeeze(np.array(traj)) for j in range(min(traj.shape[0] - 1, len(traj_artists))): traj_artists[j].set_data([traj[j, 0], traj[j + 1, 0]], [traj[j, 1], traj[j + 1, 1]]) traj_artists[j].set_3d_properties([traj[j, 2], traj[j + 1, 2]]) # Draw missing and reached targets if targets is not None and targets_reached is not None: reached = targets[targets_reached, :] reached = reached[-2:, :] missing = targets[targets_reached == False, :] reached_targets_artist.set_data(reached[:, 0], reached[:, 1]) reached_targets_artist.set_3d_properties(reached[:, 2]) targets_artist.set_data(missing[:, 0], missing[:, 1]) targets_artist.set_3d_properties(missing[:, 2]) ax.draw_artist(targets_artist) ax.draw_artist(reached_targets_artist) # Draw projected target if missing.any(): projected_tar_artists[0].set_data([missing[0, 0]], [ax.get_ylim()[1]]) projected_tar_artists[0].set_3d_properties([missing[0, 2]]) projected_tar_artists[1].set_data([ax.get_xlim()[0]], [missing[0, 1]]) projected_tar_artists[1].set_3d_properties([missing[0, 2]]) [ax.draw_artist(projected_tar_artist) for projected_tar_artist in projected_tar_artists] # Draw quadrotor trajectory trajectory_start_pt = max(len(x_trajectory) - 100, 0) trajectories_artist.set_data(x_trajectory[trajectory_start_pt:, 0], x_trajectory[trajectory_start_pt:, 1]) trajectories_artist.set_3d_properties(x_trajectory[trajectory_start_pt:, 2]) ax.draw_artist(trajectories_artist) # Draw projected trajectory projected_traj_artists[0].set_data(x_trajectory[trajectory_start_pt:, 0], ax.get_ylim()[1]) projected_traj_artists[0].set_3d_properties(x_trajectory[trajectory_start_pt:, 2]) projected_traj_artists[1].set_data(ax.get_xlim()[0], x_trajectory[trajectory_start_pt:, 1]) projected_traj_artists[1].set_3d_properties(x_trajectory[trajectory_start_pt:, 2]) [ax.draw_artist(projected_traj_artist) for projected_traj_artist in projected_traj_artists] # Draw drone art drone_art = draw_drone(x_trajectory[-1, 0:3], x_trajectory[-1, 3:7], quad.x_f, quad.y_f) drone_sketch_artist_x_motor.set_data(drone_art[0][0], drone_art[1][0]) drone_sketch_artist_x_motor.set_3d_properties(drone_art[2][0]) drone_sketch_artist.set_data(drone_art[0], drone_art[1]) drone_sketch_artist.set_3d_properties(drone_art[2]) ax.draw_artist(drone_sketch_artist) ax.draw_artist(drone_sketch_artist_x_motor) if int_traj is not None: draw_fading_traj(int_traj, int_traj_artists) for int_traj_artist in int_traj_artists: ax.draw_artist(int_traj_artist) if sim_traj is not None: draw_fading_traj(sim_traj, sim_traj_artists) for sim_traj_artist in sim_traj_artists: ax.draw_artist(sim_traj_artist) if pred_traj is not None: draw_fading_traj(pred_traj, pred_traj_artists) for pred_traj_artist in pred_traj_artists: ax.draw_artist(pred_traj_artist) if x_pred_cov is not None: n_std = 3 x_std = np.sqrt(x_pred_cov[:, 0, 0]) * n_std y_std = np.sqrt(x_pred_cov[:, 1, 1]) * n_std z_std = np.sqrt(x_pred_cov[:, 2, 2]) * n_std for i, cov_artist in enumerate(cov_artists): center = pred_traj[i+1, 0:3] radii = np.diag(np.array([x_std[i], y_std[i], z_std[i]])) x, y, z = draw_covariance_ellipsoid(center, radii) cov_artist.set_data(x, y) cov_artist.set_3d_properties(z) for cov_artist in cov_artists: ax.draw_artist(cov_artist) if follow_quad: ax.set_xlim([x_trajectory[-1, 0] - world_rad, x_trajectory[-1, 0] + world_rad]) ax.set_ylim([x_trajectory[-1, 1] - world_rad, x_trajectory[-1, 1] + world_rad]) ax.set_zlim([x_trajectory[-1, 2] - world_rad, x_trajectory[-1, 2] + world_rad]) # fill in the axes rectangle fig.canvas.blit(ax.bbox) def trajectory_tracking_results(t_ref, x_ref, x_executed, u_ref, u_executed, title, w_control=None, legend_labels=None, quat_error=True): if legend_labels is None: legend_labels = ['reference', 'simulated'] with_ref = True if x_ref is not None else False fig, ax = plt.subplots(3, 4, sharex='all', figsize=(7, 9)) SMALL_SIZE = 8 MEDIUM_SIZE = 10 BIGGER_SIZE = 12 plt.rc('font', size=SMALL_SIZE) # controls default text sizes plt.rc('axes', titlesize=SMALL_SIZE) # fontsize of the axes title plt.rc('axes', labelsize=MEDIUM_SIZE) # fontsize of the x and y labels plt.rc('xtick', labelsize=SMALL_SIZE) # fontsize of the tick labels plt.rc('ytick', labelsize=SMALL_SIZE) # fontsize of the tick labels plt.rc('legend', fontsize=SMALL_SIZE) # legend fontsize plt.rc('figure', titlesize=BIGGER_SIZE) # fontsize of the figure title labels = ['x', 'y', 'z'] for i in range(3): ax[i, 0].plot(t_ref, x_executed[:, i], label=legend_labels[1]) if with_ref: ax[i, 0].plot(t_ref, x_ref[:, i], label=legend_labels[0]) ax[i, 0].legend() ax[i, 0].set_ylabel(labels[i]) ax[0, 0].set_title(r'$p\:[m]$') ax[2, 0].set_xlabel(r'$t [s]$') q_euler = np.stack([quaternion_to_euler(x_executed[j, 3:7]) for j in range(x_executed.shape[0])]) for i in range(3): ax[i, 1].plot(t_ref, q_euler[:, i], label=legend_labels[1]) if with_ref: ref_euler = np.stack([quaternion_to_euler(x_ref[j, 3:7]) for j in range(x_ref.shape[0])]) q_err = [] for i in range(t_ref.shape[0]): q_err.append(q_dot_q(x_executed[i, 3:7], quaternion_inverse(x_ref[i, 3:7]))) q_err = np.stack(q_err) for i in range(3): ax[i, 1].plot(t_ref, ref_euler[:, i], label=legend_labels[0]) if quat_error: ax[i, 1].plot(t_ref, q_err[:, i + 1], label='quat error') for i in range(3): ax[i, 1].legend() ax[0, 1].set_title(r'$\theta\:[rad]$') ax[2, 1].set_xlabel(r'$t [s]$') for i in range(3): ax[i, 2].plot(t_ref, x_executed[:, i + 7], label=legend_labels[1]) if with_ref: ax[i, 2].plot(t_ref, x_ref[:, i + 7], label=legend_labels[0]) ax[i, 2].legend() ax[0, 2].set_title(r'$v\:[m/s]$') ax[2, 2].set_xlabel(r'$t [s]$') for i in range(3): ax[i, 3].plot(t_ref, x_executed[:, i + 10], label=legend_labels[1]) if with_ref: ax[i, 3].plot(t_ref, x_ref[:, i + 10], label=legend_labels[0]) if w_control is not None: ax[i, 3].plot(t_ref, w_control[:, i], label='control') ax[i, 3].legend() ax[0, 3].set_title(r'$\omega\:[rad/s]$') ax[2, 3].set_xlabel(r'$t [s]$') plt.suptitle(title) if u_ref is not None and u_executed is not None: ax = plt.subplots(1, 4, sharex="all", sharey="all")[1] for i in range(4): ax[i].plot(t_ref, u_ref[:, i], label='ref') ax[i].plot(t_ref, u_executed[:, i], label='simulated') ax[i].set_xlabel(r'$t [s]$') tit = 'Control %d' % (i + 1) ax[i].set_title(tit) ax[i].legend() dir_path = os.path.dirname(os.path.realpath(__file__)) img_save_dir = dir_path + '/../../results/images/' safe_mkdir_recursive(img_save_dir, overwrite=False) fig.savefig(img_save_dir + 'mse_exp', dpi=None, facecolor='w', edgecolor='w', orientation='portrait', transparent=False, pad_inches=0.1) def mse_tracking_experiment_plot(v_max, mse, traj_type_vec, train_samples_vec, legends, y_labels, t_opt=None, font_size=16): # Check if there is the variants dimension in the data if len(mse.shape) == 4: variants_dim = mse.shape[3] else: variants_dim = 1 fig, axes = plt.subplots(variants_dim, len(traj_type_vec), sharex='col', sharey='none', figsize=(17, 2.5 * variants_dim + 2)) if variants_dim == 1 and len(traj_type_vec) > 1: axes = axes[np.newaxis, :] elif variants_dim == 1: axes = np.expand_dims(axes, 0) axes = np.expand_dims(axes, 0) elif len(traj_type_vec) == 1: axes = axes[:, np.newaxis] for seed_id, track_seed in enumerate(traj_type_vec): for j in range(variants_dim): for i, _ in enumerate(train_samples_vec): mse_data = mse[seed_id, :, i, j] if len(mse.shape) == 4 else mse[seed_id, :, i] label = legends[i] if seed_id == 0 and j == 0 else None if legends[i] == 'perfect': axes[j, seed_id].plot(v_max[seed_id, :], mse_data, '--o', linewidth=4, label=label) else: axes[j, seed_id].plot(v_max[seed_id, :], mse_data, '--o', label=label) if seed_id == 0: axes[j, seed_id].set_ylabel(y_labels[j], size=font_size) if j == 0: axes[j, seed_id].set_title('RMSE [m] | ' + str(track_seed), size=font_size+2) axes[j, seed_id].grid() axes[j, seed_id].tick_params(labelsize=font_size) axes[variants_dim - 1, seed_id].set_xlabel('max vel [m/s]', size=font_size) legend_cols = len(train_samples_vec) fig.legend(loc="upper center", fancybox=True, borderaxespad=0.05, ncol=legend_cols, mode="expand", fontsize=font_size - 4) plt.tight_layout(h_pad=1.4) plt.subplots_adjust(top=0.7 + 0.05 * variants_dim) dir_path = os.path.dirname(os.path.realpath(__file__)) img_save_dir = dir_path + '/../../results/images/' safe_mkdir_recursive(img_save_dir, overwrite=False) try: tikzplotlib.save(img_save_dir + "mse.tex") except: pass fig.savefig(img_save_dir + 'mse', dpi=None, facecolor='w', edgecolor='w', orientation='portrait', transparent=False, pad_inches=0.1) if t_opt is None: return v = v_max.reshape(-1) ind_v = np.argsort(v, axis=0) fig = plt.figure(figsize=(17, 4.5)) for i, n_train in enumerate(train_samples_vec): plt.plot(v[ind_v], t_opt.reshape(t_opt.shape[0] * t_opt.shape[1], -1)[ind_v, i], label=legends[i]) fig.legend(loc="upper center", fancybox=True, borderaxespad=0.05, ncol=legend_cols, mode="expand", fontsize=font_size) plt.ylabel('Mean MPC loop time (s)', fontsize=font_size) plt.xlabel('Max vel [m/s]', fontsize=font_size) try: tikzplotlib.save(img_save_dir + "t_opt.tex") except: pass fig.savefig(img_save_dir + 't_opt', dpi=None, facecolor='w', edgecolor='w', orientation='portrait', transparent=False, bbox_inches=None, pad_inches=0.1) def load_past_experiments(): metadata_file, mse_file, v_file, t_opt_file = get_experiment_files() try: with open(metadata_file) as json_file: metadata = json.load(json_file) except: metadata = None mse = np.load(mse_file) v = np.load(v_file) t_opt = np.load(t_opt_file) return metadata, mse, v, t_opt def get_experiment_files(): results_path = PathConfig.RESULTS_DIR metadata_file = os.path.join(results_path, 'experiments', 'metadata.json') mse_file = os.path.join(results_path, 'experiments', 'mse.npy') mean_v_file = os.path.join(results_path, 'experiments', 'mean_v.npy') t_opt_file = os.path.join(results_path, 'experiments', 't_opt.npy') if not os.path.exists(metadata_file): safe_mknode_recursive(os.path.join(results_path, 'experiments'), 'metadata.json', overwrite=False) if not os.path.exists(mse_file): safe_mknode_recursive(os.path.join(results_path, 'experiments'), 'mse.npy', overwrite=False) if not os.path.exists(mean_v_file): safe_mknode_recursive(os.path.join(results_path, 'experiments'), 'mean_v.npy', overwrite=False) if not os.path.exists(t_opt_file): safe_mknode_recursive(os.path.join(results_path, 'experiments'), 't_opt.npy', overwrite=False) return metadata_file, mse_file, mean_v_file, t_opt_file