├── .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


/.gitignore:
--------------------------------------------------------------------------------
 1 | venv/*
 2 | data/
 3 | results/
 4 | .idea/
 5 | *.pyc
 6 | ros_gp_mpc/acados_models/*
 7 | c_generated_code/
 8 | *.egg-info/
 9 | dist/
10 | build/
11 | acados_models/


--------------------------------------------------------------------------------
/.gitmodules:
--------------------------------------------------------------------------------
1 | [submodule "ml-casadi"]
2 | 	path = ml-casadi
3 | 	url = https://github.com/TUM-AAS/ml-casadi
4 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | # Real-Time Neural MPC
  2 | 
  3 | This repository contains the code for experiments associated to our paper 
  4 | 
  5 | ```
  6 | Real-time Neural-MPC: Deep Learning Model Predictive Control for Quadrotors and Agile Robotic Platforms
  7 | ```
  8 | [Arxiv Link](https://arxiv.org/pdf/2203.07747)
  9 | 
 10 | If you are looking for the ML-CasADi framework code you can find it [here](https://github.com/TUM-AAS/ml-casadi).
 11 | 
 12 | ## Installation
 13 | ### Checkout Submodules
 14 | ```
 15 | git submodule update --init --recursive
 16 | ```
 17 | ### Acados
 18 | - Follow the [installation instructions](https://docs.acados.org/installation/index.html).
 19 | - Install the [Python Interface](https://docs.acados.org/python_interface/index.html).
 20 | - Ensure that `LD_LIBRARY_PATH` is set correctly (`DYLD_LIBRARY_PATH`on MacOS).
 21 | - Ensure that `ACADOS_SOURCE_DIR` is set correctly.
 22 | 
 23 | ### Further Requirements
 24 | ```
 25 | pip install -r requirements.txt
 26 | ```
 27 | 
 28 | Make sure the ML-CasADi framework is part of the python path.
 29 | ```
 30 | export PYTHONPATH="${PYTHONPATH}:<path-to-git>/ml-casadi"
 31 | ```
 32 | Python 3.9 is recommended.
 33 | 
 34 | # Experiments
 35 | 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.
 36 | 
 37 | Change the working directory to
 38 | ```
 39 | cd ros_dd_mpc
 40 | ```
 41 | ## Simulation
 42 | ### Data Collection
 43 | Run the following script to collect a few minutes of flight samples
 44 | ```
 45 | python src/experiments/point_tracking_and_record.py --recording --dataset_name simplified_sim_dataset --simulation_time 300
 46 | ```
 47 | 
 48 | ### Fitting a MLP Model
 49 | 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.
 50 | 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:
 51 | ```
 52 |     # ## Dataset loading ## #
 53 |     ds_name = "simplified_sim_dataset"
 54 |     ds_metadata = {
 55 |         "noisy": True,
 56 |         "drag": True,
 57 |         "payload": False,
 58 |         "motor_noise": True
 59 |     }
 60 | ```
 61 | 
 62 | 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).
 63 | We assign a name to the model for future referencing, e.g.: `simple_sim_mlp`
 64 | ```
 65 | 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
 66 | ```
 67 | The model will be saved under the directory `ros_dd_mpc/results/model_fitting/<git_hash>/`
 68 | 
 69 | ### Fitting GP and RDRv
 70 | For instructions on how to fit a GP or RDRv model for comparison see the [here](https://github.com/uzh-rpg/data_driven_mpc)
 71 | 
 72 | ### Test the Fitted Model
 73 | ```
 74 | python src/experiments/trajectory_test.py --model_version <git_hash> --model_name simple_sim_mlp --model_typ mlp_approx
 75 | ```
 76 | where the `model_type` argument can be one of `mlp_approx`(Real-time Neural MPC), `mlp` (Naive Integration), `gp` (Gaussian Process Model).
 77 | 
 78 | For a baseline comparison result run the same script without model parameters:
 79 | ```
 80 | python src/experiments/trajectory_test.py
 81 | ```
 82 | 
 83 | Multiple models can be compared at once via
 84 | ```
 85 | python src/experiments/comparative_experiment.py --model_version <git_hash_1 git_hash_2 ...> --model_name <name_1 name_2 ...> --model_type <type_1 type_2> --fast
 86 | ```
 87 | 
 88 | Results are saved in the `results/` folder.
 89 | 
 90 | ### Citing
 91 | 
 92 | If you use this code in an academic context, please cite the following publication:
 93 | 
 94 | ```
 95 | @article{salzmann2023neural,
 96 |   title={Real-time Neural-MPC: Deep Learning Model Predictive Control for Quadrotors and Agile Robotic Platforms},
 97 |   author={Salzmann, Tim and Kaufmann, Elia and Arrizabalaga, Jon and Pavone, Marco and Scaramuzza, Davide and Ryll, Markus},
 98 |   journal={IEEE Robotics and Automation Letters},
 99 |   doi={10.1109/LRA.2023.3246839},
100 |   year={2023}
101 | }
102 | ```


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
 1 | numpy==2.2
 2 | scipy==1.15
 3 | tqdm
 4 | matplotlib==3.9
 5 | scikit-learn==1.6
 6 | casadi==3.6
 7 | pyquaternion
 8 | joblib
 9 | pandas
10 | PyYAML
11 | rospkg>=1.3
12 | torch==2.6
13 | 
14 | -i https://rospypi.github.io/simple/
15 | rospy


--------------------------------------------------------------------------------
/ros_dd_mpc/CMakeLists.txt:
--------------------------------------------------------------------------------
  1 | cmake_minimum_required(VERSION 2.8.3)
  2 | project(ros_dd_mpc)
  3 | 
  4 | ## Compile as C++11, supported in ROS Kinetic and newer
  5 | # add_compile_options(-std=c++11)
  6 | 
  7 | ## Find catkin macros and libraries
  8 | ## if COMPONENTS list like find_package(catkin REQUIRED COMPONENTS xyz)
  9 | ## is used, also find other catkin packages
 10 | find_package(catkin REQUIRED COMPONENTS
 11 |   rospy
 12 |   std_msgs
 13 |   message_generation
 14 | )
 15 | 
 16 | ## System dependencies are found with CMake's conventions
 17 | # find_package(Boost REQUIRED COMPONENTS system)
 18 | 
 19 | 
 20 | ## Uncomment this if the package has a setup.py. This macro ensures
 21 | ## modules and global scripts declared therein get installed
 22 | ## See http://ros.org/doc/api/catkin/html/user_guide/setup_dot_py.html
 23 | # catkin_python_setup()
 24 | 
 25 | ################################################
 26 | ## Declare ROS messages, services and actions ##
 27 | ################################################
 28 | 
 29 | ## To declare and build messages, services or actions from within this
 30 | ## package, follow these steps:
 31 | ## * Let MSG_DEP_SET be the set of packages whose message types you use in
 32 | ##   your messages/services/actions (e.g. std_msgs, actionlib_msgs, ...).
 33 | ## * In the file package.xml:
 34 | ##   * add a build_depend tag for "message_generation"
 35 | ##   * add a build_depend and a exec_depend tag for each package in MSG_DEP_SET
 36 | ##   * If MSG_DEP_SET isn't empty the following dependency has been pulled in
 37 | ##     but can be declared for certainty nonetheless:
 38 | ##     * add a exec_depend tag for "message_runtime"
 39 | ## * In this file (CMakeLists.txt):
 40 | ##   * add "message_generation" and every package in MSG_DEP_SET to
 41 | ##     find_package(catkin REQUIRED COMPONENTS ...)
 42 | ##   * add "message_runtime" and every package in MSG_DEP_SET to
 43 | ##     catkin_package(CATKIN_DEPENDS ...)
 44 | ##   * uncomment the add_*_files sections below as needed
 45 | ##     and list every .msg/.srv/.action file to be processed
 46 | ##   * uncomment the generate_messages entry below
 47 | ##   * add every package in MSG_DEP_SET to generate_messages(DEPENDENCIES ...)
 48 | 
 49 | ## Generate messages in the 'msg' folder
 50 | add_message_files(
 51 |   FILES
 52 |   ReferenceTrajectory.msg
 53 | )
 54 | 
 55 | ## Generate services in the 'srv' folder
 56 | # add_service_files(
 57 | #   FILES
 58 | #   Service1.srv
 59 | #   Service2.srv
 60 | # )
 61 | 
 62 | ## Generate actions in the 'action' folder
 63 | # add_action_files(
 64 | #   FILES
 65 | #   Action1.action
 66 | #   Action2.action
 67 | # )
 68 | 
 69 | ## Generate added messages and services with any dependencies listed here
 70 | generate_messages(
 71 |   DEPENDENCIES
 72 |   std_msgs  # Or other packages containing msgs
 73 | )
 74 | 
 75 | ################################################
 76 | ## Declare ROS dynamic reconfigure parameters ##
 77 | ################################################
 78 | 
 79 | ## To declare and build dynamic reconfigure parameters within this
 80 | ## package, follow these steps:
 81 | ## * In the file package.xml:
 82 | ##   * add a build_depend and a exec_depend tag for "dynamic_reconfigure"
 83 | ## * In this file (CMakeLists.txt):
 84 | ##   * add "dynamic_reconfigure" to
 85 | ##     find_package(catkin REQUIRED COMPONENTS ...)
 86 | ##   * uncomment the "generate_dynamic_reconfigure_options" section below
 87 | ##     and list every .cfg file to be processed
 88 | 
 89 | ## Generate dynamic reconfigure parameters in the 'cfg' folder
 90 | # generate_dynamic_reconfigure_options(
 91 | #   cfg/DynReconf1.cfg
 92 | #   cfg/DynReconf2.cfg
 93 | # )
 94 | 
 95 | ###################################
 96 | ## catkin specific configuration ##
 97 | ###################################
 98 | ## The catkin_package macro generates cmake config files for your package
 99 | ## Declare things to be passed to dependent projects
100 | ## INCLUDE_DIRS: uncomment this if your package contains header files
101 | ## LIBRARIES: libraries you create in this project that dependent projects also need
102 | ## CATKIN_DEPENDS: catkin_packages dependent projects also need
103 | ## DEPENDS: system dependencies of this project that dependent projects also need
104 | catkin_package(
105 | #  INCLUDE_DIRS include
106 | #  LIBRARIES ros_python3_issues
107 |    CATKIN_DEPENDS message_runtime
108 | #  DEPENDS system_lib
109 | )
110 | 
111 | ###########
112 | ## Build ##
113 | ###########
114 | 
115 | ## Specify additional locations of header files
116 | ## Your package locations should be listed before other locations
117 | include_directories(
118 | # include
119 |   ${catkin_INCLUDE_DIRS}
120 | )
121 | 
122 | ## Declare a C++ library
123 | # add_library(${PROJECT_NAME}
124 | #   src/${PROJECT_NAME}/ros_python3_issues.cpp
125 | # )
126 | 
127 | ## Add cmake target dependencies of the library
128 | ## as an example, code may need to be generated before libraries
129 | ## either from message generation or dynamic reconfigure
130 | # add_dependencies(${PROJECT_NAME} ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS})
131 | 
132 | ## Declare a C++ executable
133 | ## With catkin_make all packages are built within a single CMake context
134 | ## The recommended prefix ensures that target names across packages don't collide
135 | # add_executable(${PROJECT_NAME}_node src/ros_python3_issues_node.cpp)
136 | 
137 | ## Rename C++ executable without prefix
138 | ## The above recommended prefix causes long target names, the following renames the
139 | ## target back to the shorter version for ease of user use
140 | ## e.g. "rosrun someones_pkg node" instead of "rosrun someones_pkg someones_pkg_node"
141 | # set_target_properties(${PROJECT_NAME}_node PROPERTIES OUTPUT_NAME node PREFIX "")
142 | 
143 | ## Add cmake target dependencies of the executable
144 | ## same as for the library above
145 | # add_dependencies(${PROJECT_NAME}_node ${${PROJECT_NAME}_EXPORTED_TARGETS} ${catkin_EXPORTED_TARGETS})
146 | 
147 | ## Specify libraries to link a library or executable target against
148 | # target_link_libraries(${PROJECT_NAME}_node
149 | #   ${catkin_LIBRARIES}
150 | # )
151 | 
152 | #############
153 | ## Install ##
154 | #############
155 | 
156 | # all install targets should use catkin DESTINATION variables
157 | # See http://ros.org/doc/api/catkin/html/adv_user_guide/variables.html
158 | 
159 | ## Mark executable scripts (Python etc.) for installation
160 | ## in contrast to setup.py, you can choose the destination
161 | # install(PROGRAMS
162 | #   scripts/my_python_script
163 | #   DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION}
164 | # )
165 | 
166 | ## Mark executables and/or libraries for installation
167 | # install(TARGETS ${PROJECT_NAME} ${PROJECT_NAME}_node
168 | #   ARCHIVE DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION}
169 | #   LIBRARY DESTINATION ${CATKIN_PACKAGE_LIB_DESTINATION}
170 | #   RUNTIME DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION}
171 | # )
172 | 
173 | ## Mark cpp header files for installation
174 | # install(DIRECTORY include/${PROJECT_NAME}/
175 | #   DESTINATION ${CATKIN_PACKAGE_INCLUDE_DESTINATION}
176 | #   FILES_MATCHING PATTERN "*.h"
177 | #   PATTERN ".svn" EXCLUDE
178 | # )
179 | 
180 | ## Mark other files for installation (e.g. launch and bag files, etc.)
181 | # install(FILES
182 | #   # myfile1
183 | #   # myfile2
184 | #   DESTINATION ${CATKIN_PACKAGE_SHARE_DESTINATION}
185 | # )
186 | 
187 | #############
188 | ## Testing ##
189 | #############
190 | 
191 | ## Add gtest based cpp test target and link libraries
192 | # catkin_add_gtest(${PROJECT_NAME}-test test/test_ros_python3_issues.cpp)
193 | # if(TARGET ${PROJECT_NAME}-test)
194 | #   target_link_libraries(${PROJECT_NAME}-test ${PROJECT_NAME})
195 | # endif()
196 | 
197 | ## Add folders to be run by python nosetests
198 | # catkin_add_nosetests(test)
199 | 
200 | # if (CATKIN_ENABLE_TESTING)
201 | #   find_package(rostest REQUIRED)
202 | #   add_rostest(test/test_issue_rosunit.launch)
203 | # endif()
204 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/config/agisim_simulation_run.yaml:
--------------------------------------------------------------------------------
1 | quad_name: 'kingfisher'
2 | world_limits: None
3 | use_ekf_synchronization: False
4 | control_freq_factor: 5
5 | environment: 'agisim'


--------------------------------------------------------------------------------
/ros_dd_mpc/config/arena_limits.yaml:
--------------------------------------------------------------------------------
1 | x_min: -6.0
2 | x_max: 12.0
3 | y_min: -7.0
4 | y_max: 7.0
5 | z_min: 1.0
6 | z_max: 4.0


--------------------------------------------------------------------------------
/ros_dd_mpc/config/arena_run.yaml:
--------------------------------------------------------------------------------
1 | quad_name: 'kingfisher'
2 | world_limits: arena_limits
3 | use_ekf_synchronization: False
4 | control_freq_factor: 5
5 | environment: 'arena'


--------------------------------------------------------------------------------
/ros_dd_mpc/config/basic.world:
--------------------------------------------------------------------------------
 1 | <?xml version="1.0" ?>
 2 | <sdf version="1.4">
 3 |   <world name="default">
 4 |     <include>
 5 |       <uri>model://ground_plane</uri>
 6 |     </include>
 7 |     <include>
 8 |       <uri>model://sun</uri>
 9 |     </include>
10 | 
11 |     <!-- Only one ROS interface plugin is required per world, as any other plugin can connect a Gazebo
12 |          topic to a ROS topic (or vise versa). -->
13 |     <plugin name="ros_interface_plugin" filename="librotors_gazebo_ros_interface_plugin.so"/>
14 | 
15 |     <spherical_coordinates>
16 |       <surface_model>EARTH_WGS84</surface_model>
17 |       <latitude_deg>47.3667</latitude_deg>
18 |       <longitude_deg>8.5500</longitude_deg>
19 |       <elevation>500.0</elevation>
20 |       <heading_deg>0</heading_deg>
21 |     </spherical_coordinates>
22 |     <physics type='ode'>
23 |       <ode>
24 |         <solver>
25 |           <type>quick</type>
26 |           <iters>1000</iters>
27 |           <sor>1.3</sor>
28 |         </solver>
29 |         <constraints>
30 |           <cfm>0</cfm>
31 |           <erp>0.2</erp>
32 |           <contact_max_correcting_vel>100</contact_max_correcting_vel>
33 |           <contact_surface_layer>0.001</contact_surface_layer>
34 |         </constraints>
35 |       </ode>
36 |       <max_step_size>0.01</max_step_size>
37 |       <real_time_factor>0.5</real_time_factor>
38 |       <real_time_update_rate>50</real_time_update_rate>
39 |       <gravity>0 0 -9.8</gravity>
40 |     </physics>
41 |   </world>
42 | </sdf>
43 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/config/circle_and_lemniscate_options.yaml:
--------------------------------------------------------------------------------
1 | # Parameters for the loop and lemniscate trajectories at successively increasing speed
2 | loop_z: 2.5              # Z position of loop plane [m]
3 | loop_r: 5.0            # Radius of loop [m]
4 | loop_v_max: 10.0       # Maximum speed achieved (approx.) [m/s]
5 | loop_lin_a: 0.25       # Linear acceleration and deceleration of trajectory [m/s^2]
6 | loop_clockwise: False  # Rotation direction of the quad
7 | loop_yawing: False     # Yaw quadrotor along the circle (only for loop)


--------------------------------------------------------------------------------
/ros_dd_mpc/config/configuration_parameters.py:
--------------------------------------------------------------------------------
  1 | """ Set of tunable parameters for the Simplified Simulator and model fitting.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | import os
 15 | 
 16 | 
 17 | class DirectoryConfig:
 18 |     """
 19 |     Class for storing directories within the package
 20 |     """
 21 | 
 22 |     _dir_path = os.path.dirname(os.path.realpath(__file__))
 23 |     SAVE_DIR = _dir_path + '/../results/model_fitting'
 24 |     RESULTS_DIR = _dir_path + '/../results'
 25 |     CONFIG_DIR = _dir_path + ''
 26 |     DATA_DIR = _dir_path + '/../data'
 27 | 
 28 | 
 29 | class SimpleSimConfig:
 30 |     """
 31 |     Class for storing the Simplified Simulator configurations.
 32 |     """
 33 | 
 34 |     # Set to True to show a real-time Matplotlib animation of the experiments for the Simplified Simulator. Execution
 35 |     # will be slower if the GUI is turned on. Note: setting to True may require some further library installation work.
 36 |     custom_sim_gui = False
 37 | 
 38 |     # Set to True to display a plot describing the trajectory tracking results after the execution.
 39 |     result_plots = True
 40 | 
 41 |     # Set to True to show the trajectory that will be executed before the execution time
 42 |     pre_run_debug_plots = True
 43 | 
 44 |     # Choice of disturbances modeled in our Simplified Simulator. For more details about the parameters used refer to
 45 |     # the script: src/quad_mpc/quad_3d.py.
 46 |     simulation_disturbances = {
 47 |         "noisy": True,                       # Thrust and torque gaussian noises
 48 |         "drag": True,                        # 2nd order polynomial aerodynamic drag effect
 49 |         "payload": False,                    # Payload force in the Z axis
 50 |         "motor_noise": True                  # Asymmetric voltage noise in the motors
 51 |     }
 52 | 
 53 | 
 54 | class ModelFitConfig:
 55 |     """
 56 |     Class for storing flags for the model fitting scripts.
 57 |     """
 58 | 
 59 |     # Dataset loading ## #
 60 |     ds_name = "simplified_sim_dataset"
 61 |     ds_metadata = {
 62 |         "noisy": True,
 63 |         "drag": True,
 64 |         "payload": False,
 65 |         "motor_noise": True
 66 |     }
 67 | 
 68 |     # ds_name = "agisim_dataset"
 69 |     # ds_metadata = {
 70 |     #     "agisim": "default",
 71 |     # }
 72 | 
 73 |     # ds_name = "arena_dataset"
 74 |     # ds_metadata = {
 75 |     #     "arena": "default",
 76 |     # }
 77 | 
 78 |     # ds_name = "neurobem_dataset"
 79 |     # ds_metadata = {
 80 |     #     "arena": "default",
 81 |     # }
 82 | 
 83 |     # ## Visualization ## #
 84 |     # Training mode
 85 |     visualize_training_result = True
 86 |     visualize_data = False
 87 | 
 88 |     # Visualization mode
 89 |     grid_sampling_viz = True
 90 |     x_viz = [7, 8, 9]
 91 |     u_viz = []
 92 |     y_viz = [7, 8, 9]
 93 | 
 94 |     # ## Data post-processing ## #
 95 |     histogram_bins = 40              # Cluster data using histogram binning
 96 |     histogram_threshold = 0.001      # Remove bins where the total ratio of data is lower than this threshold
 97 |     velocity_cap = 16                # Also remove datasets point if abs(velocity) > x_cap
 98 | 
 99 |     # ############# Experimental ############# #
100 | 
101 |     # ## Use fit model to generate synthetic data ## #
102 |     use_dense_model = False
103 |     dense_model_version = ""
104 |     dense_model_name = ""
105 |     dense_training_points = 200
106 | 
107 |     # ## Clustering for multidimensional models ## #
108 |     clusters = 1
109 |     load_clusters = False
110 | 
111 | 
112 | class GroundEffectMapConfig:
113 |     """
114 |     Class for storing parameters for the ground effect map.
115 |     """
116 |     resolution = 0.1
117 |     origin = (-4, 9)
118 |     horizon = ((-7, 7), (-7, 7))
119 |     box_min = (-4.25, 9.37)
120 |     box_max = (-2.76, 10.13)
121 |     box_height = 0.7
122 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/config/flyingroom_limits.yaml:
--------------------------------------------------------------------------------
1 | x_min: -0.4
2 | x_max: 4.3
3 | y_min: -2.4
4 | y_max: 2.3
5 | z_min: 0.5
6 | z_max: 2.0


--------------------------------------------------------------------------------
/ros_dd_mpc/config/ground_effect_limits.yaml:
--------------------------------------------------------------------------------
1 | x_min: -4.25
2 | x_max: -2.76
3 | y_min: 9.37
4 | y_max: 10.13
5 | z_min: 0.8 # 0.95
6 | z_max: 0.85


--------------------------------------------------------------------------------
/ros_dd_mpc/config/kingfisher.yaml:
--------------------------------------------------------------------------------
 1 | mass:               0.752                    # [kg]
 2 | tbm_fr:             [ 0.075, -0.10, 0.0]          # [m]
 3 | tbm_bl:             [-0.075,  0.10, 0.0]         # [m]
 4 | tbm_br:             [-0.075, -0.10, 0.0]         # [m]
 5 | tbm_fl:             [ 0.075,  0.10, 0.0]          # [m]
 6 | inertia:            [0.0025, 0.0021, 0.0043]     # [kgm^2]
 7 | motor_omega_min:    150.0                     # [rad/s]
 8 | motor_omega_max:    2800.0                    # [rad/s]
 9 | motor_tau:          0.033                     # [s]
10 | omega_max:          [10.0, 10.0, 4.0]         # [rad/s]
11 | comm_delay:         0.04                      # [s]
12 | 
13 | thrust_map:         [1.562522e-6, 0.0,  0.0]
14 | kappa:              0.022                    # [Nm/N]
15 | thrust_min:         0.0                       # [N]
16 | thrust_max:         8.5                      # [N] per motor
17 | rotors_config:      "cross"
18 | 
19 | aero_coeff_1:       [0.0, 0.0, 0.0] # [0.26, 0.28, 0.42] # [N/(v/m)]
20 | aero_coeff_3:       [0.0, 0.0, 0.0] # [N/(v/m)]
21 | aero_coeff_h:        0.0            # 0.01 [N/(v^2/m^2)]
22 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/config/simulation_run.yaml:
--------------------------------------------------------------------------------
1 | quad_name: 'kingfisher'
2 | world_limits: None
3 | use_ekf_synchronization: False
4 | control_freq_factor: 5
5 | environment: 'gazebo'


--------------------------------------------------------------------------------
/ros_dd_mpc/launch/dd_mpc_wrapper.launch:
--------------------------------------------------------------------------------
  1 | <?xml version="1.0"?>
  2 | <launch>
  3 |     <!-- true if running the nodes in the gazebo simulator environment. false if running on the real platform-->
  4 |     <arg name="environment" default="flying_room"/>
  5 | 
  6 |     <!-- Recording parameters -->
  7 |     <arg name="recording" default="false"/>
  8 |     <arg name="dataset_name" default="gazebo_dataset"/>
  9 |     <arg name="overwrite" default="true"/>
 10 |     <arg name="train_split" default="true"/>
 11 |     <arg name="record_raw_optitrack" default="false"/>
 12 | 
 13 |     <!-- Reference trajectory parameters -->
 14 |     <arg name="n_seeds" default="1"/> <!-- How many random seeds to use to generate trajectories -->
 15 |     <arg name="trajectory_speeds" default="[4.0]"/>
 16 |     <arg name="flight_mode" default="random"/> <!-- one of: random, hover, loop -->
 17 |     <arg name="plot" default="false"/>
 18 |     <arg name="world_limits" default="None"/>
 19 | 
 20 |     <!-- Stable version: Use time horizon of 1 second, 10 control nodes and control @ 50 Hz (control_freq_factor=5) -->
 21 |     <arg name="t_horizon" default="1"/>
 22 |     <arg name="n_nodes" default="10"/>
 23 | 
 24 |     <!-- Model loading parameters -->
 25 |     <arg name="model_version" default=""/>
 26 |     <arg name="model_name" default=""/>
 27 |     <arg name="model_type" default="gp"/>
 28 | 
 29 |     <!-- Trajectory tracking experiment reset -->
 30 |     <arg name="reset_experiment" default="true"/>
 31 | 
 32 |     <group ns="gp_mpc">
 33 | 
 34 |         <group if="$(eval arg('environment')=='agisim')">
 35 |             <!-- Node that runs the MPC -->
 36 |             <node pkg="ros_dd_mpc" type="dd_mpc_node.py" name="mpc_wrapper" output="screen">
 37 |                 <rosparam file="$(find ros_dd_mpc)/config/agisim_simulation_run.yaml"/>
 38 | 
 39 |                 <param name="recording" value="$(arg recording)"/>
 40 |                 <param name="dataset_name" value="$(arg dataset_name)"/>
 41 |                 <param name="overwrite" value="$(arg overwrite)"/>
 42 |                 <param name="training_split" value="$(arg train_split)"/>
 43 |                 <param name="record_raw_optitrack" value="$(arg record_raw_optitrack)"/>
 44 | 
 45 |                 <param name="model_type" value="$(arg model_type)"/>
 46 |                 <param name="model_version" value="$(arg model_version)"/>
 47 |                 <param name="model_name" value="$(arg model_name)"/>
 48 | 
 49 |                 <param name="plot" value="$(arg plot)"/>
 50 |                 <param name="t_horizon" value="$(arg t_horizon)"/>
 51 |                 <param name="n_nodes" value="$(arg n_nodes)"/>
 52 | 
 53 |                 <param name="reset_experiment" value="$(arg reset_experiment)"/>
 54 |             </node>
 55 | 
 56 |             <!-- Random trajectory generator -->
 57 |             <node pkg="ros_dd_mpc" type="reference_publisher_node.py" name="ref_gen" output="screen">
 58 |                 <rosparam file="$(find ros_dd_mpc)/config/agisim_simulation_run.yaml"/>
 59 |                 <rosparam file="$(find ros_dd_mpc)/config/circle_and_lemniscate_options.yaml"/>
 60 | 
 61 |                 <param name="n_seeds" value="$(arg n_seeds)"/>
 62 |                 <param name="v_list" value="$(arg trajectory_speeds)"/>
 63 |                 <param name="mode" value="$(arg flight_mode)"/>
 64 |                 <param name="world_limits" value="$(arg world_limits)"/>
 65 | 
 66 |                 <param name="t_horizon" value="$(arg t_horizon)"/>
 67 |                 <param name="n_nodes" value="$(arg n_nodes)"/>
 68 | 
 69 |                 <param name="plot" value="$(arg plot)"/>
 70 |             </node>
 71 |         </group>
 72 | 
 73 |         <group if="$(eval arg('environment')=='gazebo')">
 74 |             <!-- Node that runs the MPC -->
 75 |             <node pkg="ros_dd_mpc" type="dd_mpc_node.py" name="mpc_wrapper" output="screen">
 76 |                 <rosparam file="$(find ros_dd_mpc)/config/simulation_run.yaml"/>
 77 | 
 78 |                 <param name="recording" value="$(arg recording)"/>
 79 |                 <param name="dataset_name" value="$(arg dataset_name)"/>
 80 |                 <param name="overwrite" value="$(arg overwrite)"/>
 81 |                 <param name="training_split" value="$(arg train_split)"/>
 82 |                 <param name="record_raw_optitrack" value="$(arg record_raw_optitrack)"/>
 83 | 
 84 |                 <param name="model_type" value="$(arg model_type)"/>
 85 |                 <param name="model_version" value="$(arg model_version)"/>
 86 |                 <param name="model_name" value="$(arg model_name)"/>
 87 | 
 88 |                 <param name="plot" value="$(arg plot)"/>
 89 |                 <param name="t_horizon" value="$(arg t_horizon)"/>
 90 |                 <param name="n_nodes" value="$(arg n_nodes)"/>
 91 | 
 92 |                 <param name="reset_experiment" value="$(arg reset_experiment)"/>
 93 |             </node>
 94 | 
 95 |             <!-- Random trajectory generator -->
 96 |             <node pkg="ros_dd_mpc" type="reference_publisher_node.py" name="ref_gen" output="screen">
 97 |                 <rosparam file="$(find ros_dd_mpc)/config/simulation_run.yaml"/>
 98 |                 <rosparam file="$(find ros_dd_mpc)/config/circle_and_lemniscate_options.yaml"/>
 99 | 
100 |                 <param name="n_seeds" value="$(arg n_seeds)"/>
101 |                 <param name="v_list" value="$(arg trajectory_speeds)"/>
102 |                 <param name="mode" value="$(arg flight_mode)"/>
103 |                 <param name="world_limits" value="$(arg world_limits)"/>
104 | 
105 |                 <param name="t_horizon" value="$(arg t_horizon)"/>
106 |                 <param name="n_nodes" value="$(arg n_nodes)"/>
107 | 
108 |                 <param name="plot" value="$(arg plot)"/>
109 |             </node>
110 |         </group>
111 | 
112 |         <group if="$(eval arg('environment')=='flying_room')">
113 |             <!-- Node that runs the MPC -->
114 |             <node pkg="ros_dd_mpc" type="dd_mpc_node.py" name="mpc_wrapper" output="screen">
115 |                 <rosparam file="$(find ros_dd_mpc)/config/flying_room_run.yaml"/>
116 | 
117 |                 <param name="recording" value="$(arg recording)"/>
118 |                 <param name="dataset_name" value="$(arg dataset_name)"/>
119 |                 <param name="overwrite" value="$(arg overwrite)"/>
120 |                 <param name="training_split" value="$(arg train_split)"/>
121 |                 <param name="record_raw_optitrack" value="$(arg record_raw_optitrack)"/>
122 | 
123 |                 <param name="model_type" value="$(arg model_type)"/>
124 |                 <param name="model_version" value="$(arg model_version)"/>
125 |                 <param name="model_name" value="$(arg model_name)"/>
126 | 
127 |                 <param name="plot" value="$(arg plot)"/>
128 |                 <param name="t_horizon" value="$(arg t_horizon)"/>
129 |                 <param name="n_nodes" value="$(arg n_nodes)"/>
130 | 
131 |                 <param name="reset_experiment" value="$(arg reset_experiment)"/>
132 |             </node>
133 | 
134 |             <!-- Random trajectory generator -->
135 |             <node pkg="ros_dd_mpc" type="reference_publisher_node.py" name="ref_gen" output="screen">
136 |                 <rosparam file="$(find ros_dd_mpc)/config/flying_room_run.yaml"/>
137 |                 <rosparam file="$(find ros_dd_mpc)/config/circle_and_lemniscate_options.yaml"/>
138 | 
139 |                 <param name="n_seeds" value="$(arg n_seeds)"/>
140 |                 <param name="v_list" value="$(arg trajectory_speeds)"/>
141 |                 <param name="mode" value="$(arg flight_mode)"/>
142 |                 <param name="world_limits" value="$(arg world_limits)"/>
143 | 
144 |                 <param name="t_horizon" value="$(arg t_horizon)"/>
145 |                 <param name="n_nodes" value="$(arg n_nodes)"/>
146 | 
147 |                 <param name="plot" value="$(arg plot)"/>
148 |             </node>
149 |         </group>
150 | 
151 |         <group if="$(eval arg('environment')=='arena')">
152 |             <!-- Node that runs the MPC -->
153 |             <node pkg="ros_dd_mpc" type="dd_mpc_node.py" name="mpc_wrapper" output="screen">
154 |                 <rosparam file="$(find ros_dd_mpc)/config/arena_run.yaml"/>
155 | 
156 |                 <param name="recording" value="$(arg recording)"/>
157 |                 <param name="dataset_name" value="$(arg dataset_name)"/>
158 |                 <param name="overwrite" value="$(arg overwrite)"/>
159 |                 <param name="training_split" value="$(arg train_split)"/>
160 |                 <param name="record_raw_optitrack" value="$(arg record_raw_optitrack)"/>
161 | 
162 |                 <param name="model_type" value="$(arg model_type)"/>
163 |                 <param name="model_version" value="$(arg model_version)"/>
164 |                 <param name="model_name" value="$(arg model_name)"/>
165 | 
166 |                 <param name="plot" value="$(arg plot)"/>
167 |                 <param name="t_horizon" value="$(arg t_horizon)"/>
168 |                 <param name="n_nodes" value="$(arg n_nodes)"/>
169 | 
170 |                 <param name="reset_experiment" value="$(arg reset_experiment)"/>
171 |             </node>
172 | 
173 |             <!-- Random trajectory generator -->
174 |             <node pkg="ros_dd_mpc" type="reference_publisher_node.py" name="ref_gen" output="screen">
175 |                 <rosparam file="$(find ros_dd_mpc)/config/arena_run.yaml"/>
176 |                 <rosparam file="$(find ros_dd_mpc)/config/circle_and_lemniscate_options.yaml"/>
177 | 
178 |                 <param name="n_seeds" value="$(arg n_seeds)"/>
179 |                 <param name="v_list" value="$(arg trajectory_speeds)"/>
180 |                 <param name="mode" value="$(arg flight_mode)"/>
181 |                 <param name="world_limits" value="$(arg world_limits)"/>
182 | 
183 |                 <param name="t_horizon" value="$(arg t_horizon)"/>
184 |                 <param name="n_nodes" value="$(arg n_nodes)"/>
185 | 
186 |                 <param name="plot" value="$(arg plot)"/>
187 |             </node>
188 |         </group>
189 | 
190 |     </group>
191 | </launch>


--------------------------------------------------------------------------------
/ros_dd_mpc/msg/ReferenceTrajectory.msg:
--------------------------------------------------------------------------------
1 | int32 seq_len
2 | string traj_name
3 | float64 v_input
4 | float64[] trajectory
5 | float64[] dt
6 | float64[] inputs
7 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/nodes/reference_publisher_node.py:
--------------------------------------------------------------------------------
  1 | #!/usr/bin/env python3.6
  2 | """ Node wrapper for publishing trajectories for the MPC pipeline to track.
  3 | 
  4 | This program is free software: you can redistribute it and/or modify it under
  5 | the terms of the GNU General Public License as published by the Free Software
  6 | Foundation, either version 3 of the License, or (at your option) any later
  7 | version.
  8 | This program is distributed in the hope that it will be useful, but WITHOUT
  9 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 10 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 11 | You should have received a copy of the GNU General Public License along with
 12 | this program. If not, see <http://www.gnu.org/licenses/>.
 13 | """
 14 | 
 15 | from std_msgs.msg import Bool
 16 | from ros_dd_mpc.msg import ReferenceTrajectory
 17 | from src.quad_mpc.create_ros_dd_mpc import custom_quad_param_loader
 18 | from src.utils.trajectories import loop_trajectory, random_trajectory, lemniscate_trajectory, flyover_trajectory, \
 19 |     flyover_trajectory_collect
 20 | import numpy as np
 21 | import rospy
 22 | 
 23 | 
 24 | class ReferenceGenerator:
 25 | 
 26 |     def __init__(self):
 27 | 
 28 |         self.gp_mpc_busy = True
 29 | 
 30 |         rospy.init_node("reference_generator")
 31 | 
 32 |         plot = rospy.get_param('~plot', default=True)
 33 | 
 34 |         quad_name = rospy.get_param('~quad_name', default='hummingbird')
 35 |         quad = custom_quad_param_loader(quad_name)
 36 | 
 37 |         # Configuration for random flight mode
 38 |         n_seeds = rospy.get_param('~n_seeds', default=1)
 39 |         v_list = rospy.get_param('~v_list', default=[3.5])
 40 |         if isinstance(v_list, str):
 41 |             v_list = v_list.split('[')[1].split(']')[0]
 42 |             v_list = [float(v.strip()) for v in v_list.split(',')]
 43 | 
 44 |         # Select if generate "random" trajectories, "hover" mode or increasing speed "loop" mode
 45 |         mode = rospy.get_param('~mode', default="random")
 46 |         if mode != "random" and mode != "flyover":
 47 |             n_seeds = 1
 48 | 
 49 |         # Load parameters of loop trajectory
 50 |         loop_r = rospy.get_param('~loop_r', default=1.5)
 51 |         loop_z = rospy.get_param('~loop_z', default=1)
 52 |         loop_v_max = rospy.get_param('~loop_v_max', default=3)
 53 |         loop_a = rospy.get_param('~loop_lin_a', default=0.075)
 54 |         loop_cc = rospy.get_param('~loop_clockwise', default=True)
 55 |         loop_yawing = rospy.get_param('~loop_yawing', default=True)
 56 | 
 57 |         # Load world limits if any
 58 |         map_limits = rospy.get_param('~world_limits', default=None)
 59 | 
 60 |         # Control at 50 hz. Use time horizon=1 and 10 nodes
 61 |         n_mpc_nodes = rospy.get_param('~n_nodes', default=10)
 62 |         t_horizon = rospy.get_param('~t_horizon', default=1.0)
 63 |         control_freq_factor = rospy.get_param('~control_freq_factor', default=5 if quad_name == "hummingbird" else 6)
 64 |         opt_dt = t_horizon / (n_mpc_nodes * control_freq_factor)
 65 | 
 66 |         reference_topic = "reference"
 67 |         status_topic = "busy"
 68 |         reference_pub = rospy.Publisher(reference_topic, ReferenceTrajectory, queue_size=1)
 69 |         rospy.Subscriber(status_topic, Bool, self.status_callback)
 70 | 
 71 |         v_ind = 0
 72 |         seed = 0
 73 | 
 74 |         # Calculate total number of trajectories
 75 |         n_trajectories = n_seeds * len(v_list)
 76 |         curr_trajectory_ind = 0
 77 | 
 78 |         rate = rospy.Rate(0.2)
 79 |         while not rospy.is_shutdown():
 80 |             if not self.gp_mpc_busy and mode == "hover":
 81 |                 rospy.loginfo("Sending hover-in-place command")
 82 |                 msg = ReferenceTrajectory()
 83 |                 reference_pub.publish(msg)
 84 |                 rospy.signal_shutdown("All trajectories were sent to the MPC")
 85 |                 break
 86 | 
 87 |             if not self.gp_mpc_busy and curr_trajectory_ind == n_trajectories:
 88 |                 msg = ReferenceTrajectory()
 89 |                 reference_pub.publish(msg)
 90 |                 rospy.signal_shutdown("All trajectories were sent to the MPC")
 91 |                 break
 92 | 
 93 |             if not self.gp_mpc_busy and mode == "loop":
 94 |                 rospy.loginfo("Sending increasing speed loop trajectory")
 95 |                 x_ref, t_ref, u_ref = loop_trajectory(quad, opt_dt, v_max=loop_v_max, radius=loop_r, z=loop_z,
 96 |                                                       lin_acc=loop_a, clockwise=loop_cc, map_name=map_limits,
 97 |                                                       yawing=loop_yawing, plot=plot)
 98 | 
 99 |                 msg = ReferenceTrajectory()
100 |                 msg.traj_name = "circle"
101 |                 msg.v_input = loop_v_max
102 |                 msg.seq_len = x_ref.shape[0]
103 |                 msg.trajectory = np.reshape(x_ref, (-1,)).tolist()
104 |                 msg.dt = t_ref.tolist()
105 |                 msg.inputs = np.reshape(u_ref, (-1,)).tolist()
106 | 
107 |                 reference_pub.publish(msg)
108 |                 curr_trajectory_ind += 1
109 |                 self.gp_mpc_busy = True
110 | 
111 |             elif not self.gp_mpc_busy and mode == "lemniscate":
112 |                 rospy.loginfo("Sending increasing speed lemniscate trajectory")
113 |                 x_ref, t_ref, u_ref = lemniscate_trajectory(quad, opt_dt, v_max=loop_v_max, radius=loop_r, z=loop_z,
114 |                                                             lin_acc=loop_a, clockwise=loop_cc, map_name=map_limits,
115 |                                                             yawing=loop_yawing, plot=plot)
116 | 
117 |                 msg = ReferenceTrajectory()
118 |                 msg.traj_name = "lemniscate"
119 |                 msg.v_input = loop_v_max
120 |                 msg.seq_len = x_ref.shape[0]
121 |                 msg.trajectory = np.reshape(x_ref, (-1,)).tolist()
122 |                 msg.dt = t_ref.tolist()
123 |                 msg.inputs = np.reshape(u_ref, (-1,)).tolist()
124 | 
125 |                 reference_pub.publish(msg)
126 |                 curr_trajectory_ind += 1
127 |                 self.gp_mpc_busy = True
128 | 
129 |             elif not self.gp_mpc_busy and mode == "random":
130 | 
131 |                 speed = v_list[v_ind]
132 |                 log_msg = "Random trajectory generator %d/%d. Seed: %d. Mean vel: %.3f m/s" % \
133 |                           (curr_trajectory_ind + 1, n_trajectories, seed, speed)
134 |                 rospy.loginfo(log_msg)
135 | 
136 |                 x_ref, t_ref, u_ref = random_trajectory(quad, opt_dt, seed=seed, speed=speed, map_name=map_limits,
137 |                                                         plot=plot)
138 |                 msg = ReferenceTrajectory()
139 |                 msg.traj_name = "random_" + str(seed)
140 |                 msg.v_input = speed
141 |                 msg.seq_len = x_ref.shape[0]
142 |                 msg.trajectory = np.reshape(x_ref, (-1, )).tolist()
143 |                 msg.dt = t_ref.tolist()
144 |                 msg.inputs = np.reshape(u_ref, (-1, )).tolist()
145 | 
146 |                 reference_pub.publish(msg)
147 |                 curr_trajectory_ind += 1
148 |                 self.gp_mpc_busy = True
149 | 
150 |                 if v_ind + 1 < len(v_list):
151 |                     v_ind += 1
152 |                 else:
153 |                     seed += 1
154 |                     v_ind = 0
155 | 
156 |             elif not self.gp_mpc_busy and mode == "flyover_collect":
157 | 
158 |                 speed = v_list[v_ind]
159 |                 log_msg = "Flyover trajectory generator %d/%d. Seed: %d. Mean vel: %.3f m/s" % \
160 |                           (curr_trajectory_ind + 1, n_trajectories, seed, speed)
161 |                 rospy.loginfo(log_msg)
162 | 
163 |                 x_ref, t_ref, u_ref = flyover_trajectory_collect(quad, opt_dt, seed=seed, speed=speed,
164 |                                                                  flyover_box_name=map_limits)
165 |                 msg = ReferenceTrajectory()
166 |                 msg.traj_name = "flyover_" + str(seed)
167 |                 msg.v_input = speed
168 |                 msg.seq_len = x_ref.shape[0]
169 |                 msg.trajectory = np.reshape(x_ref, (-1, )).tolist()
170 |                 msg.dt = t_ref.tolist()
171 |                 msg.inputs = np.reshape(u_ref, (-1, )).tolist()
172 | 
173 |                 reference_pub.publish(msg)
174 |                 curr_trajectory_ind += 1
175 |                 self.gp_mpc_busy = True
176 | 
177 |                 if v_ind + 1 < len(v_list):
178 |                     v_ind += 1
179 |                 else:
180 |                     seed += 1
181 |                     v_ind = 0
182 | 
183 |             elif not self.gp_mpc_busy and mode == "flyover":
184 | 
185 |                 speed = v_list[v_ind]
186 |                 log_msg = "Flyover trajectory generator %d/%d. Seed: %d. Mean vel: %.3f m/s" % \
187 |                           (curr_trajectory_ind + 1, n_trajectories, seed, speed)
188 |                 rospy.loginfo(log_msg)
189 | 
190 |                 x_ref, t_ref, u_ref = flyover_trajectory(quad, opt_dt, seed=seed, speed=speed,
191 |                                                          flyover_box_name=map_limits)
192 |                 msg = ReferenceTrajectory()
193 |                 msg.traj_name = "flyover_" + str(seed)
194 |                 msg.v_input = speed
195 |                 msg.seq_len = x_ref.shape[0]
196 |                 msg.trajectory = np.reshape(x_ref, (-1, )).tolist()
197 |                 msg.dt = t_ref.tolist()
198 |                 msg.inputs = np.reshape(u_ref, (-1, )).tolist()
199 | 
200 |                 reference_pub.publish(msg)
201 |                 curr_trajectory_ind += 1
202 |                 self.gp_mpc_busy = True
203 | 
204 |                 if v_ind + 1 < len(v_list):
205 |                     v_ind += 1
206 |                 else:
207 |                     seed += 1
208 |                     v_ind = 0
209 | 
210 |             elif not self.gp_mpc_busy:
211 |                 raise ValueError("Unknown trajectory type: %s" % mode)
212 | 
213 |             rate.sleep()
214 | 
215 |     def status_callback(self, msg):
216 |         """
217 |         Callback function for tracking if the dd_mpc node is busy
218 |         :param msg: Message from the subscriber
219 |         :type msg: Bool
220 |         """
221 |         self.gp_mpc_busy = msg.data
222 | 
223 | 
224 | if __name__ == "__main__":
225 | 
226 |     ReferenceGenerator()
227 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/package.xml:
--------------------------------------------------------------------------------
 1 | <?xml version="1.0"?>
 2 | <package format="2">
 3 |   <name>ros_dd_mpc</name>
 4 |   <version>0.0.0</version>
 5 |   <description>The ros_dd_mpc package</description>
 6 | 
 7 |   <!-- One maintainer tag required, multiple allowed, one person per tag -->
 8 |   <!-- Example:  -->
 9 |   <!-- <maintainer email="jane.doe@example.com">Jane Doe</maintainer> -->
10 |   <maintainer email="guillemtorrente@hotmail.com">Guillem Torrente</maintainer>
11 | 
12 | 
13 |   <!-- One license tag required, multiple allowed, one license per tag -->
14 |   <!-- Commonly used license strings: -->
15 |   <!--   BSD, MIT, Boost Software License, GPLv2, GPLv3, LGPLv2.1, LGPLv3 -->
16 |   <license>TODO</license>
17 | 
18 | 
19 |   <!-- Url tags are optional, but multiple are allowed, one per tag -->
20 |   <!-- Optional attribute type can be: website, bugtracker, or repository -->
21 |   <!-- Example: -->
22 |   <!-- <url type="website">http://wiki.ros.org/ros_python3_issues</url> -->
23 | 
24 | 
25 |   <!-- Author tags are optional, multiple are allowed, one per tag -->
26 |   <!-- Authors do not have to be maintainers, but could be -->
27 |   <!-- Example: -->
28 |   <!-- <author email="jane.doe@example.com">Jane Doe</author> -->
29 | 
30 | 
31 |   <!-- The *depend tags are used to specify dependencies -->
32 |   <!-- Dependencies can be catkin packages or system dependencies -->
33 |   <!-- Examples: -->
34 |   <!-- Use depend as a shortcut for packages that are both build and exec dependencies -->
35 |   <!--   <depend>roscpp</depend> -->
36 |   <!--   Note that this is equivalent to the following: -->
37 |   <!--   <build_depend>roscpp</build_depend> -->
38 |   <!--   <exec_depend>roscpp</exec_depend> -->
39 |   <!-- Use build_depend for packages you need at compile time: -->
40 |   <build_depend>message_generation</build_depend>
41 |   <!-- Use build_export_depend for packages you need in order to build against this package: -->
42 |   <!--   <build_export_depend>message_generation</build_export_depend> -->
43 |   <!-- Use buildtool_depend for build tool packages: -->
44 |   <!--   <buildtool_depend>catkin</buildtool_depend> -->
45 |   <!-- Use exec_depend for packages you need at runtime: -->
46 |   <exec_depend>message_runtime</exec_depend>
47 |   <!-- Use test_depend for packages you need only for testing: -->
48 |   <!--   <test_depend>gtest</test_depend> -->
49 |   <!-- Use doc_depend for packages you need only for building documentation: -->
50 |   <!--   <doc_depend>doxygen</doc_depend> -->
51 |   <buildtool_depend>catkin</buildtool_depend>
52 |   <build_depend>rospy</build_depend>
53 |   <build_export_depend>rospy</build_export_depend>
54 |   <exec_depend>rospy</exec_depend>
55 |   <exec_depend>agiros_msgs</exec_depend>
56 |   <exec_depend>deep_casadi</exec_depend>
57 | 
58 | 
59 |   <!-- The export tag contains other, unspecified, tags -->
60 |   <export>
61 |     <!-- Other tools can request additional information be placed here -->
62 | 
63 |   </export>
64 | </package>
65 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/TUM-AAS/neural-mpc/5f15661c2d1a3b6072911d8993214c50a68e11eb/ros_dd_mpc/src/__init__.py


--------------------------------------------------------------------------------
/ros_dd_mpc/src/experiments/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/TUM-AAS/neural-mpc/5f15661c2d1a3b6072911d8993214c50a68e11eb/ros_dd_mpc/src/experiments/__init__.py


--------------------------------------------------------------------------------
/ros_dd_mpc/src/experiments/point_tracking_and_record.py:
--------------------------------------------------------------------------------
  1 | """ Executes aggressive maneuvers for collecting flight data on the Simplified Simulator to later train models on.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | 
 15 | import os
 16 | import sys
 17 | import time
 18 | import copy
 19 | import argparse
 20 | import itertools
 21 | 
 22 | import pandas as pd
 23 | import numpy as np
 24 | import casadi as cs
 25 | 
 26 | from src.utils.visualization import draw_drone_simulation, initialize_drone_plotter
 27 | from src.experiments.comparative_experiment import prepare_quadrotor_mpc
 28 | from src.utils.utils import safe_mknode_recursive, jsonify, euclidean_dist, get_data_dir_and_file
 29 | from config.configuration_parameters import SimpleSimConfig
 30 | 
 31 | 
 32 | def get_record_file_and_dir(record_dict_template, recording_options, simulation_setup, overwrite=True):
 33 |     dataset_name = recording_options["dataset_name"]
 34 |     training_split = recording_options["training_split"]
 35 | 
 36 |     # Directory and file name for data recording
 37 |     rec_file_dir, rec_file_name = get_data_dir_and_file(dataset_name, training_split, simulation_setup)
 38 | 
 39 |     overwritten = safe_mknode_recursive(rec_file_dir, rec_file_name, overwrite=overwrite)
 40 | 
 41 |     rec_dict = copy.deepcopy(record_dict_template)
 42 |     rec_file = os.path.join(rec_file_dir, rec_file_name)
 43 |     if overwrite or (not overwrite and not overwritten):
 44 |         for key in rec_dict.keys():
 45 |             rec_dict[key] = jsonify(rec_dict[key])
 46 | 
 47 |         df = pd.DataFrame(rec_dict)
 48 |         df.to_csv(rec_file, index=False, header=True)
 49 | 
 50 |         rec_dict = copy.deepcopy(record_dict_template)
 51 | 
 52 |     return rec_dict, rec_file
 53 | 
 54 | 
 55 | def make_record_dict(state_dim):
 56 |     blank_recording_dict = {
 57 |             "state_in": np.zeros((0, state_dim)),
 58 |             "state_ref": np.zeros((0, state_dim)),
 59 |             "error": np.zeros((0, state_dim)),
 60 |             "input_in": np.zeros((0, 4)),
 61 |             "state_out": np.zeros((0, state_dim)),
 62 |             "state_pred": np.zeros((0, state_dim)),
 63 |             "timestamp": np.zeros((0, 1)),
 64 |             "dt": np.zeros((0, 1)),
 65 |     }
 66 |     return blank_recording_dict
 67 | 
 68 | 
 69 | def check_out_data(rec_dict, state_final, x_pred, w_opt, dt):
 70 |     rec_dict["dt"] = np.append(rec_dict["dt"], dt)
 71 |     rec_dict["input_in"] = np.append(rec_dict["input_in"], w_opt[np.newaxis, :4], axis=0)
 72 |     rec_dict["state_out"] = np.append(rec_dict["state_out"], state_final, 0)
 73 | 
 74 |     if x_pred is not None:
 75 |         err = state_final - x_pred
 76 |         rec_dict["error"] = np.append(rec_dict["error"], err, axis=0)
 77 |         rec_dict["state_pred"] = np.append(rec_dict["state_pred"], x_pred[np.newaxis, :], axis=0)
 78 | 
 79 |     return rec_dict
 80 | 
 81 | 
 82 | def sample_random_target(x_current, world_radius, aggressive=True):
 83 |     """
 84 |     Creates a new target point to reach.
 85 |     :param x_current: current position of the quad. Only used if aggressive=True
 86 |     :param world_radius: radius of the area where points are sampled from
 87 |     :param aggressive: if aggressive=True, points will be sampled away from the current position. If False, then points
 88 |     will be sampled uniformly from the world area.
 89 |     :return: new sampled target point. A 3-dimensional numpy array.
 90 |     """
 91 | 
 92 |     if aggressive:
 93 | 
 94 |         # Polar 3D coordinates
 95 |         theta = np.random.uniform(0, 2 * np.pi, 1)
 96 |         psi = np.random.uniform(0, 2 * np.pi, 1)
 97 |         r = 1 * world_radius + np.random.uniform(-0.5, 0.5, 1) * world_radius
 98 | 
 99 |         # Transform to cartesian
100 |         x = r * np.sin(theta) * np.cos(psi)
101 |         y = r * np.sin(theta) * np.sin(psi)
102 |         z = r * np.cos(theta)
103 | 
104 |         return x_current + np.array([x, y, z]).reshape((1, 3))
105 | 
106 |     else:
107 |         return np.random.uniform(-world_radius, world_radius, (1, 3))
108 | 
109 | 
110 | def main(model_options, recording_options, simulation_options, parameters):
111 | 
112 |     world_radius = 3
113 | 
114 |     if parameters["initial_state"] is None:
115 |         initial_state = [0.0, 0.0, 0.0] + [1, 0, 0, 0] + [0, 0, 0] + [0, 0, 0]
116 |     else:
117 |         initial_state = parameters["initial_state"]
118 |     sim_starting_pos = initial_state
119 |     quad_current_state = sim_starting_pos
120 | 
121 |     if parameters["preset_targets"] is not None:
122 |         targets = parameters["preset_targets"]
123 |     else:
124 |         targets = sample_random_target(np.array(initial_state[:3]), world_radius,
125 |                                        aggressive=recording_options["recording"])
126 | 
127 |     quad_mpc = prepare_quadrotor_mpc(simulation_options, **model_options, t_horizon=0.5,
128 |                                      q_mask=np.array([1, 1, 1, 0.01, 0.01, 0.01, 1, 1, 1, 0, 0, 0]))
129 | 
130 |     # Recover some necessary variables from the MPC object
131 |     my_quad = quad_mpc.quad
132 |     n_mpc_nodes = quad_mpc.n_nodes
133 |     t_horizon = quad_mpc.t_horizon
134 |     simulation_dt = quad_mpc.simulation_dt
135 |     reference_over_sampling = 3
136 |     control_period = t_horizon / (n_mpc_nodes * reference_over_sampling)
137 | 
138 |     my_quad.set_state(quad_current_state)
139 | 
140 |     # Real time plot params
141 |     n_forward_props = n_mpc_nodes
142 |     plot_sim_traj = False
143 | 
144 |     x_pred = None
145 |     w_opt = None
146 |     initial_guess = None
147 | 
148 |     # The optimization should be faster or equal than the duration of the optimization time step
149 |     assert control_period <= t_horizon / n_mpc_nodes
150 | 
151 |     state = quad_mpc.get_state()
152 | 
153 |     # ####### Recording mode code ####### #
154 |     recording = recording_options["recording"]
155 |     state_dim = state.shape[0]
156 |     blank_recording_dict = make_record_dict(state_dim)
157 | 
158 |     # Get recording file and directory
159 |     if recording:
160 |         if parameters["real_time_plot"]:
161 |             parameters["real_time_plot"] = False
162 |             print("Turned off real time plot during recording mode.")
163 | 
164 |         rec_dict, rec_file = get_record_file_and_dir(blank_recording_dict, recording_options, simulation_options)
165 | 
166 |     else:
167 |         rec_dict = rec_file = None
168 | 
169 |     # Generate necessary art pack for real time plot
170 |     if parameters["real_time_plot"]:
171 |         real_time_art_pack = initialize_drone_plotter(n_props=n_forward_props, quad_rad=my_quad.x_f,
172 |                                                       world_rad=world_radius)
173 |     else:
174 |         real_time_art_pack = None
175 | 
176 |     start_time = time.time()
177 |     simulation_time = 0.0
178 | 
179 |     # Simulation tracking stuff
180 |     targets_reached = np.array([False for _ in targets])
181 |     quad_trajectory = np.array(quad_current_state).reshape(1, -1)
182 | 
183 |     n_iteration_count = 0
184 | 
185 |     print("Targets reached: ", end='')
186 |     # All targets loop
187 |     while False in targets_reached and (time.time() - start_time) < parameters["max_sim_time"]:
188 |         current_target_i = np.where(targets_reached == False)[0][0]
189 |         current_target = targets[current_target_i]
190 |         current_target_reached = False
191 | 
192 |         quad_target_state = [list(current_target), [1, 0, 0, 0], [0, 0, 0], [0, 0, 0]]
193 |         model_ind = quad_mpc.set_reference(quad_target_state)
194 | 
195 |         # Provide an initial guess without the uncertainty prop.
196 |         if initial_guess is None:
197 |             initial_guess = quad_mpc.optimize(use_model=model_ind)
198 |             initial_guess = quad_mpc.reshape_input_sequence(initial_guess)
199 | 
200 |         # MPC loop
201 |         while not current_target_reached and (time.time() - start_time) < parameters["max_sim_time"]:
202 | 
203 |             # Emergency recovery (quad controller gone out of control lol)
204 |             if np.any(state[7:10] > 14) or n_iteration_count > 100:
205 |                 n_iteration_count = 0
206 |                 my_quad.set_state(list(itertools.chain.from_iterable(quad_target_state)))
207 | 
208 |             state = quad_mpc.get_state()
209 |             if recording:
210 |                 rec_dict["state_in"] = np.append(rec_dict["state_in"], state.T, 0)
211 |                 rec_dict["timestamp"] = np.append(rec_dict["timestamp"], time.time() - start_time)
212 |                 stacked_ref = np.array(list(itertools.chain.from_iterable(quad_target_state)))[np.newaxis, :]
213 |                 rec_dict["state_ref"] = np.append(rec_dict["state_ref"], stacked_ref, 0)
214 |                 if simulation_time != 0.0:
215 |                     rec_dict = check_out_data(rec_dict, state.T, x_pred, w_opt, simulation_time)
216 | 
217 |             simulation_time = 0.0
218 | 
219 |             # Optimize control input to reach pre-set target
220 |             w_opt, x_pred_horizon = quad_mpc.optimize(use_model=model_ind, return_x=True)
221 |             if np.any(w_opt > (my_quad.max_input_value + 0.01)) or np.any(w_opt < (my_quad.min_input_value - 0.01)):
222 |                 print("MPC constraints were violated")
223 |             initial_guess = quad_mpc.reshape_input_sequence(w_opt)
224 |             # Save initial guess for future optimization. It is a time-shift of the current optimized variables
225 |             initial_guess = np.array(cs.vertcat(initial_guess[1:, :], cs.DM.zeros(4).T))
226 | 
227 |             # Select first input (one for each motor) - MPC applies only first optimized input to the plant
228 |             ref_u = np.squeeze(np.array(w_opt[:4]))
229 | 
230 |             if recording:
231 |                 # Integrate first input. Will be used as nominal model prediction during next save
232 |                 x_pred, _ = quad_mpc.forward_prop(np.squeeze(state), w_opt=w_opt[:4],
233 |                                                   t_horizon=control_period, use_gp=False)
234 |                 x_pred = x_pred[-1, :]
235 | 
236 |             if parameters["real_time_plot"]:
237 |                 prop_params = {"x_0": np.squeeze(state), "w_opt": w_opt, "use_model": model_ind, "t_horizon": t_horizon}
238 |                 x_int, _ = quad_mpc.forward_prop(**prop_params, use_gp=False)
239 |                 if plot_sim_traj:
240 |                     x_sim = quad_mpc.simulate_plant(quad_mpc.reshape_input_sequence(w_opt))
241 |                 else:
242 |                     x_sim = None
243 |                 draw_drone_simulation(real_time_art_pack, quad_trajectory, my_quad, targets,
244 |                                       targets_reached, x_sim, x_int, x_pred_horizon, follow_quad=False)
245 | 
246 |             while simulation_time < control_period:
247 | 
248 |                 # Simulation runtime (inner loop)
249 |                 simulation_time += simulation_dt
250 |                 quad_mpc.simulate(ref_u)
251 | 
252 |                 quad_current_state = quad_mpc.get_state()
253 | 
254 |                 # Target is reached
255 |                 if euclidean_dist(current_target[0:3], quad_current_state[0:3, 0], thresh=0.05):
256 |                     print("*", end='')
257 |                     sys.stdout.flush()
258 |                     n_iteration_count = 0
259 | 
260 |                     # Check out data immediately as new target will be optimized in next step
261 |                     if recording and len(rec_dict['state_in']) > len(rec_dict['input_in']):
262 |                         x_pred, _ = quad_mpc.forward_prop(np.squeeze(state), w_opt=w_opt[:4], t_horizon=simulation_time,
263 |                                                           use_gp=False)
264 |                         x_pred = x_pred[-1, :]
265 |                         rec_dict = check_out_data(rec_dict, quad_mpc.get_state().T, x_pred, w_opt, simulation_time)
266 | 
267 |                     # Reset optimization time -> Ask for new optimization for next target in next dt
268 |                     simulation_time = 0.0
269 | 
270 |                     # Mark current target as reached
271 |                     current_target_reached = True
272 |                     targets_reached[current_target_i] = True
273 | 
274 |                     # Remove initial guess
275 |                     initial_guess = None
276 | 
277 |                     if parameters["preset_targets"] is None:
278 |                         new_target = sample_random_target(quad_current_state[:3], world_radius, aggressive=recording)
279 |                         targets = np.append(targets, new_target, axis=0)
280 |                         targets_reached = np.append(targets_reached, False)
281 | 
282 |                     # Reset PID integral and past errors
283 |                     quad_mpc.reset()
284 | 
285 |                     break
286 | 
287 |             n_iteration_count += 1
288 | 
289 |             if parameters["real_time_plot"]:
290 |                 quad_trajectory = np.append(quad_trajectory, quad_current_state.T, axis=0)
291 |                 if len(quad_trajectory) > 300:
292 |                     quad_trajectory = np.delete(quad_trajectory, 0, 0)
293 | 
294 |         # Current target was reached. Remove incomplete recordings
295 |         if recording:
296 |             if len(rec_dict['state_in']) > len(rec_dict['input_in']):
297 |                 rec_dict["state_in"] = rec_dict["state_in"][:-1]
298 |                 rec_dict["timestamp"] = rec_dict["timestamp"][:-1]
299 |                 rec_dict["state_ref"] = rec_dict["state_ref"][:-1]
300 | 
301 |             for key in rec_dict.keys():
302 |                 rec_dict[key] = jsonify(rec_dict[key])
303 | 
304 |             df = pd.DataFrame(rec_dict)
305 |             df.to_csv(rec_file, index=True, mode='a', header=False)
306 | 
307 |             rec_dict = copy.deepcopy(blank_recording_dict)
308 | 
309 | 
310 | if __name__ == '__main__':
311 | 
312 |     parser = argparse.ArgumentParser()
313 | 
314 |     parser.add_argument("--model_version", type=str, default="",
315 |                         help="Version to load for the regression models. By default it is an 8 digit git hash.")
316 | 
317 |     parser.add_argument("--model_name", type=str, default="",
318 |                         help="Name of the regression model within the specified <model_version> folder.")
319 | 
320 |     parser.add_argument("--model_type", type=str, default="gp", choices=["gp", "rdrv"],
321 |                         help="Type of regression model (GP or RDRv linear)")
322 | 
323 |     parser.add_argument("--recording", dest="recording", action="store_true",
324 |                         help="Set to True to enable recording mode.")
325 |     parser.set_defaults(recording=False)
326 | 
327 |     parser.add_argument("--dataset_name", type=str, default="simplified_sim_dataset",
328 |                         help="Name for the generated dataset.")
329 | 
330 |     parser.add_argument("--test_split", dest="test_split", action="store_true",
331 |                         help="If the data set is the training or test split.")
332 |     parser.set_defaults(test_split=False)
333 | 
334 |     parser.add_argument("--simulation_time", type=float, default=300,
335 |                         help="Total duration of the simulation in seconds.")
336 | 
337 |     args = parser.parse_args()
338 | 
339 |     np.random.seed(123 if args.test_split else 456)
340 | 
341 |     acados_config = {
342 |         "solver_type": "SQP",
343 |         "terminal_cost": True
344 |     }
345 | 
346 |     run_options = {
347 |         "model_options": {
348 |             "version": args.model_version,
349 |             "name": args.model_name,
350 |             "reg_type": args.model_type,
351 |             "quad_name": "my_quad"
352 |         },
353 |         "recording_options": {
354 |             "recording": args.recording,
355 |             "dataset_name": args.dataset_name,
356 |             "training_split": not args.test_split,
357 |         },
358 |         "simulation_options": SimpleSimConfig.simulation_disturbances,
359 |         "parameters": {
360 |             "real_time_plot": SimpleSimConfig.custom_sim_gui,
361 |             "max_sim_time": args.simulation_time,
362 |             "preset_targets": None,
363 |             "initial_state": None,
364 |             "acados_options": acados_config
365 |         }
366 |     }
367 | 
368 |     main(**run_options)
369 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/experiments/trajectory_test.py:
--------------------------------------------------------------------------------
  1 | """ Tracks a specified trajectory on the simplified simulator using the data-augmented MPC.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | 
 15 | import time
 16 | import json
 17 | import argparse
 18 | import numpy as np
 19 | from tqdm import tqdm
 20 | from src.utils.utils import separate_variables
 21 | from src.utils.quad_3d_opt_utils import get_reference_chunk
 22 | from src.utils.trajectories import loop_trajectory, lemniscate_trajectory, check_trajectory
 23 | from src.utils.visualization import initialize_drone_plotter, draw_drone_simulation, trajectory_tracking_results
 24 | from src.experiments.comparative_experiment import prepare_quadrotor_mpc
 25 | from config.configuration_parameters import SimpleSimConfig
 26 | 
 27 | 
 28 | def main(args):
 29 |     params = {
 30 |         "version": args.model_version,
 31 |         "name": args.model_name,
 32 |         "reg_type": args.model_type,
 33 |         "quad_name": "my_quad"
 34 |     }
 35 | 
 36 |     # Load the disturbances for the custom offline simulator.
 37 |     simulation_options = SimpleSimConfig.simulation_disturbances
 38 | 
 39 |     debug_plots = SimpleSimConfig.pre_run_debug_plots
 40 |     tracking_results_plot = SimpleSimConfig.result_plots
 41 |     sim_gui = SimpleSimConfig.custom_sim_gui
 42 | 
 43 |     quad_mpc = prepare_quadrotor_mpc(simulation_options, **params)
 44 | 
 45 |     # Recover some necessary variables from the MPC object
 46 |     my_quad = quad_mpc.quad
 47 |     n_mpc_nodes = quad_mpc.n_nodes
 48 |     t_horizon = quad_mpc.t_horizon
 49 |     simulation_dt = quad_mpc.simulation_dt
 50 |     reference_over_sampling = 5
 51 |     control_period = t_horizon / (n_mpc_nodes * reference_over_sampling)
 52 | 
 53 |     if args.trajectory == "loop":
 54 |         reference_traj, reference_timestamps, reference_u = loop_trajectory(
 55 |             my_quad, control_period, radius=args.trajectory_radius, z=1, lin_acc=args.acceleration, clockwise=True,
 56 |             yawing=False, v_max=args.max_speed, map_name=None, plot=debug_plots)
 57 | 
 58 |     elif args.trajectory == "lemniscate":
 59 |         reference_traj, reference_timestamps, reference_u = lemniscate_trajectory(
 60 |             my_quad, control_period, radius=args.trajectory_radius, z=1, lin_acc=args.acceleration, clockwise=True,
 61 |             yawing=False, v_max=args.max_speed, map_name=None, plot=debug_plots)
 62 | 
 63 |     else:
 64 |         raise ValueError("Unknown trajectory {}. Options are `lemniscate` and `loop`".format(args.trajectory))
 65 | 
 66 |     if not check_trajectory(reference_traj, reference_u, reference_timestamps, debug_plots):
 67 |         return
 68 | 
 69 |     # Set quad initial state equal to the initial reference trajectory state
 70 |     quad_current_state = reference_traj[0, :].tolist()
 71 |     my_quad.set_state(quad_current_state)
 72 | 
 73 |     real_time_artists = None
 74 |     if sim_gui:
 75 |         # Initialize real time plot stuff
 76 |         world_radius = np.max(np.abs(reference_traj[:, :2])) * 1.2
 77 |         real_time_artists = initialize_drone_plotter(n_props=n_mpc_nodes, quad_rad=my_quad.length,
 78 |                                                      world_rad=world_radius, full_traj=reference_traj)
 79 | 
 80 |     ref_u = reference_u[0, :]
 81 |     quad_trajectory = np.zeros((len(reference_timestamps), len(quad_current_state)))
 82 |     u_optimized_seq = np.zeros((len(reference_timestamps), 4))
 83 | 
 84 |     # Sliding reference trajectory initial index
 85 |     current_idx = 0
 86 | 
 87 |     # Measure the MPC optimization time
 88 |     mean_opt_time = 0.0
 89 | 
 90 |     # Measure total simulation time
 91 |     total_sim_time = 0.0
 92 | 
 93 |     print("\nRunning simulation...")
 94 |     for current_idx in tqdm(range(reference_traj.shape[0])):
 95 | 
 96 |         quad_current_state = my_quad.get_state(quaternion=True, stacked=True)
 97 | 
 98 |         quad_trajectory[current_idx, :] = np.expand_dims(quad_current_state, axis=0)
 99 | 
100 |         # ##### Optimization runtime (outer loop) ##### #
101 |         # Get the chunk of trajectory required for the current optimization
102 |         ref_traj_chunk, ref_u_chunk = get_reference_chunk(
103 |             reference_traj, reference_u, current_idx, n_mpc_nodes, reference_over_sampling)
104 | 
105 |         # Set the reference for the OCP
106 |         model_ind = quad_mpc.set_reference(x_reference=separate_variables(ref_traj_chunk), u_reference=ref_u_chunk)
107 | 
108 |         # Optimize control input to reach pre-set target
109 |         t_opt_init = time.time()
110 |         w_opt, x_pred = quad_mpc.optimize(use_model=model_ind, return_x=True)
111 |         mean_opt_time += time.time() - t_opt_init
112 | 
113 |         # Select first input (one for each motor) - MPC applies only first optimized input to the plant
114 |         ref_u = np.squeeze(np.array(w_opt[:4]))
115 |         u_optimized_seq[current_idx, :] = np.reshape(ref_u, (1, -1))
116 | 
117 |         if len(quad_trajectory) > 0 and sim_gui and current_idx > 0:
118 |             draw_drone_simulation(real_time_artists, quad_trajectory[:current_idx, :], my_quad, targets=None,
119 |                                   targets_reached=None, pred_traj=x_pred, x_pred_cov=None)
120 | 
121 |         simulation_time = 0.0
122 | 
123 |         # ##### Simulation runtime (inner loop) ##### #
124 |         while simulation_time < control_period:
125 |             simulation_time += simulation_dt
126 |             total_sim_time += simulation_dt
127 |             quad_mpc.simulate(ref_u)
128 | 
129 |     u_optimized_seq[current_idx, :] = np.reshape(ref_u, (1, -1))
130 | 
131 |     quad_current_state = my_quad.get_state(quaternion=True, stacked=True)
132 |     quad_trajectory[-1, :] = np.expand_dims(quad_current_state, axis=0)
133 |     u_optimized_seq[-1, :] = np.reshape(ref_u, (1, -1))
134 | 
135 |     # Average elapsed time per optimization
136 |     mean_opt_time = mean_opt_time / current_idx * 1000
137 |     tracking_rmse = np.mean(np.sqrt(np.sum((reference_traj[:, :3] - quad_trajectory[:, :3]) ** 2, axis=1)))
138 | 
139 |     if tracking_results_plot:
140 |         v_max = np.max(reference_traj[:, 7:10])
141 | 
142 |         with_gp = ' + GP ' if quad_mpc.gp_ensemble is not None else ' - GP '
143 |         title = r'$v_{max}$=%.2f m/s | RMSE: %.4f m | %s ' % (v_max, float(tracking_rmse), with_gp)
144 |         trajectory_tracking_results(reference_timestamps, reference_traj, quad_trajectory, reference_u, u_optimized_seq,
145 |                                     title)
146 | 
147 |     v_max_abs = np.max(np.sqrt(np.sum(reference_traj[:, 7:10] ** 2, 1)))
148 | 
149 |     print("\n:::::::::::::: SIMULATION SETUP ::::::::::::::\n")
150 |     print("Simulation: Applied disturbances: ")
151 |     print(json.dumps(simulation_options))
152 |     if quad_mpc.gp_ensemble is not None:
153 |         print("\nModel: Using GP regression model: ", params["version"] + '/' + params["name"])
154 |     elif quad_mpc.mlp is not None:
155 |         print("\nModel: Using MLP regression model: ", params["version"] + '/' + params["name"])
156 |     elif quad_mpc.rdrv is not None:
157 |         print("\nModel: Using RDRv regression model: ", params["version"] + '/' + params["name"])
158 |     else:
159 |         print("\nModel: No regression model loaded")
160 | 
161 |     print("\nReference: Executed trajectory", '`' + args.trajectory + '`', "with a peak axial velocity of",
162 |           args.max_speed, "m/s, and a maximum speed of %2.3f m/s" % v_max_abs)
163 | 
164 |     print("\n::::::::::::: SIMULATION RESULTS :::::::::::::\n")
165 |     print("Mean optimization time: %.3f ms" % mean_opt_time)
166 |     print("Tracking RMSE: %.4f m\n" % tracking_rmse)
167 | 
168 | 
169 | if __name__ == '__main__':
170 | 
171 |     parser = argparse.ArgumentParser()
172 | 
173 |     parser.add_argument("--model_version", type=str, default="",
174 |                         help="Version to load for the regression models. By default it is an 8 digit git hash.")
175 | 
176 |     parser.add_argument("--model_name", type=str, default="",
177 |                         help="Name of the regression model within the specified <model_version> folder.")
178 | 
179 |     parser.add_argument("--model_type", type=str, default="gp",
180 |                         help="Type of regression model (GP or RDRv linear)")
181 | 
182 |     parser.add_argument("--trajectory", type=str, default="loop", choices=["loop", "lemniscate"],
183 |                         help='path to other necessary data files (eg. vocabularies)')
184 | 
185 |     parser.add_argument("--max_speed", type=float, default=8,
186 |                         help="Maximum axial speed executed during the flight in m/s. For the `loop` trajectory, "
187 |                              "velocities are feasible up to 14 m/s, and for the `lemniscate` up to 8 m/s")
188 | 
189 |     parser.add_argument("--acceleration", type=float, default=1,
190 |                         help="Acceleration of the reference trajectory. Higher accelerations will shorten the execution"
191 |                              "time of the tracking")
192 | 
193 |     parser.add_argument("--trajectory_radius", type=float, default=5, help="Radius of the reference trajectories")
194 |     input_arguments = parser.parse_args()
195 | 
196 |     main(input_arguments)
197 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/model_fitting/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/TUM-AAS/neural-mpc/5f15661c2d1a3b6072911d8993214c50a68e11eb/ros_dd_mpc/src/model_fitting/__init__.py


--------------------------------------------------------------------------------
/ros_dd_mpc/src/model_fitting/gp_common.py:
--------------------------------------------------------------------------------
  1 | """ Contains a set of utility functions and classes for the GP training and inference.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | 
 15 | import copy
 16 | import os
 17 | 
 18 | import joblib
 19 | import numpy as np
 20 | import pandas as pd
 21 | from sklearn.mixture import GaussianMixture
 22 | 
 23 | from src.model_fitting.gp import GPEnsemble, CustomGPRegression as npGPRegression
 24 | from src.utils.utils import undo_jsonify, prune_dataset, safe_mknode_recursive, get_data_dir_and_file, \
 25 |     separate_variables, v_dot_q, quaternion_inverse
 26 | from src.utils.visualization import visualize_data_distribution
 27 | 
 28 | 
 29 | class GPDataset:
 30 |     def __init__(self, raw_ds=None,
 31 |                  x_features=None, u_features=None, y_dim=None,
 32 |                  cap=None, n_bins=None, thresh=None,
 33 |                  visualize_data=False):
 34 | 
 35 |         self.data_labels = [r'$p_x$', r'$p_y$', r'$p_z$', r'$q_w$', r'$q_x$', r'$q_y$', r'$q_z$',
 36 |                             r'$v_x$', r'$v_y$', r'$v_z$', r'$w_x$', r'$w_y$', r'$w_z$']
 37 | 
 38 |         # Raw dataset data
 39 |         self.x_raw = None
 40 |         self.x_out_raw = None
 41 |         self.u_raw = None
 42 |         self.y_raw = None
 43 |         self.x_pred_raw = None
 44 |         self.dt_raw = None
 45 | 
 46 |         self.x_features = x_features
 47 |         self.u_features = u_features
 48 |         self.y_dim = y_dim
 49 | 
 50 |         # Data pruning parameters
 51 |         self.cap = cap
 52 |         self.n_bins = n_bins
 53 |         self.thresh = thresh
 54 |         self.plot = visualize_data
 55 | 
 56 |         # GMM clustering
 57 |         self.pruned_idx = []
 58 |         self.cluster_agency = None
 59 |         self.centroids = None
 60 | 
 61 |         # Number of data clusters
 62 |         self.n_clusters = 1
 63 | 
 64 |         if raw_ds is not None:
 65 |             self.load_data(raw_ds)
 66 |             self.prune()
 67 | 
 68 |     def load_data(self, ds):
 69 |         if isinstance(ds, np.lib.npyio.NpzFile):
 70 |             x_raw = ds['state_in']
 71 |             x_out = ds['state_out']
 72 |             x_pred = ds['state_pred']
 73 |             u_raw = ds['input_in']
 74 |             dt = ds["dt"]
 75 |         else:
 76 |             x_raw = undo_jsonify(ds['state_in'].to_numpy())
 77 |             x_out = undo_jsonify(ds['state_out'].to_numpy())
 78 |             x_pred = undo_jsonify(ds['state_pred'].to_numpy())
 79 |             u_raw = undo_jsonify(ds['input_in'].to_numpy())
 80 |             dt = ds["dt"].to_numpy()
 81 | 
 82 |         invalid = np.where(dt == 0)
 83 | 
 84 |         # Remove invalid entries (dt = 0)
 85 |         x_raw = np.delete(x_raw, invalid, axis=0)
 86 |         x_out = np.delete(x_out, invalid, axis=0)
 87 |         x_pred = np.delete(x_pred, invalid, axis=0)
 88 |         u_raw = np.delete(u_raw, invalid, axis=0)
 89 |         dt = np.delete(dt, invalid, axis=0)
 90 | 
 91 |         # Rotate velocities to body frame and recompute errors
 92 |         x_raw = world_to_body_velocity_mapping(x_raw)
 93 |         x_pred = world_to_body_velocity_mapping(x_pred)
 94 |         x_out = world_to_body_velocity_mapping(x_out)
 95 |         y_err = x_out - x_pred
 96 | 
 97 |         # Normalize error by window time (i.e. predict error dynamics instead of error itself)
 98 |         y_err /= np.expand_dims(dt, 1)
 99 | 
100 |         # Select features
101 |         self.x_raw = x_raw
102 |         self.x_out_raw = x_out
103 |         self.u_raw = u_raw
104 |         self.y_raw = y_err
105 |         self.x_pred_raw = x_pred
106 |         self.dt_raw = dt
107 | 
108 |     def prune(self):
109 |         # Prune noisy data
110 |         if self.cap is not None and self.n_bins is not None and self.thresh is not None:
111 |             x_interest = np.array([7, 8, 9])
112 |             y_interest = np.array([7, 8, 9])
113 | 
114 |             labels = [self.data_labels[dim] for dim in np.atleast_1d(y_interest)]
115 | 
116 |             prune_x_data = self.get_x(pruned=False, raw=True)[:, x_interest]
117 |             prune_y_data = self.get_y(pruned=False, raw=True)[:, y_interest]
118 |             self.pruned_idx.append(prune_dataset(prune_x_data, prune_y_data, self.cap, self.n_bins, self.thresh,
119 |                                                  plot=self.plot, labels=labels))
120 | 
121 |     def get_x(self, cluster=None, pruned=True, raw=False):
122 | 
123 |         if cluster is not None:
124 |             assert pruned
125 | 
126 |         if raw:
127 |             return self.x_raw[tuple(self.pruned_idx)] if pruned else self.x_raw
128 | 
129 |         if self.x_features is not None and self.u_features is not None:
130 |             x_f = self.x_features
131 |             u_f = self.u_features
132 |             data = np.concatenate((self.x_raw[:, x_f], self.u_raw[:, u_f]), axis=1) if u_f else self.x_raw[:, x_f]
133 |         else:
134 |             data = np.concatenate((self.x_raw, self.u_raw), axis=1)
135 |         data = data[:, np.newaxis] if len(data.shape) == 1 else data
136 | 
137 |         if pruned or cluster is not None:
138 |             data = data[tuple(self.pruned_idx)]
139 |             data = data[self.cluster_agency[cluster]] if cluster is not None else data
140 | 
141 |         return data
142 | 
143 |     @property
144 |     def x(self):
145 |         return self.get_x()
146 | 
147 |     def get_x_out(self, cluster=None, pruned=True):
148 | 
149 |         if cluster is not None:
150 |             assert pruned
151 | 
152 |         if pruned or cluster is not None:
153 |             data = self.x_out_raw[tuple(self.pruned_idx)]
154 |             data = data[self.cluster_agency[cluster]] if cluster is not None else data
155 | 
156 |             return data
157 | 
158 |         return self.x_out_raw[tuple(self.pruned_idx)] if pruned else self.x_out_raw
159 | 
160 |     @property
161 |     def x_out(self):
162 |         return self.get_x_out()
163 | 
164 |     def get_u(self, cluster=None, pruned=True, raw=False):
165 | 
166 |         if cluster is not None:
167 |             assert pruned
168 | 
169 |         if raw:
170 |             return self.u_raw[tuple(self.pruned_idx)] if pruned else self.u_raw
171 | 
172 |         data = self.u_raw[:, self.u_features] if self.u_features is not None else self.u_raw
173 |         data = data[:, np.newaxis] if len(data.shape) == 1 else data
174 | 
175 |         if pruned or cluster is not None:
176 |             data = data[tuple(self.pruned_idx)]
177 |             data = data[self.cluster_agency[cluster]] if cluster is not None else data
178 | 
179 |         return data
180 | 
181 |     @property
182 |     def u(self):
183 |         return self.get_u()
184 | 
185 |     def get_y(self, cluster=None, pruned=True, raw=False):
186 | 
187 |         if cluster is not None:
188 |             assert pruned
189 | 
190 |         if raw:
191 |             return self.y_raw[self.pruned_idx] if pruned else self.y_raw
192 | 
193 |         data = self.y_raw[:, self.y_dim] if self.y_dim is not None else self.y_raw
194 |         data = data[:, np.newaxis] if len(data.shape) == 1 else data
195 | 
196 |         if pruned or cluster is not None:
197 |             data = data[tuple(self.pruned_idx)]
198 |             data = data[self.cluster_agency[cluster]] if cluster is not None else data
199 | 
200 |         return data
201 | 
202 |     @property
203 |     def y(self):
204 |         return self.get_y()
205 | 
206 |     def get_x_pred(self, pruned=True, raw=False):
207 | 
208 |         if raw:
209 |             return self.x_pred_raw[tuple(self.pruned_idx)] if pruned else self.x_pred_raw
210 | 
211 |         data = self.x_pred_raw[:, self.y_dim] if self.y_dim is not None else self.x_pred_raw
212 |         data = data[:, np.newaxis] if len(data.shape) == 1 else data
213 | 
214 |         if pruned:
215 |             data = data[tuple(self.pruned_idx)]
216 | 
217 |         return data
218 | 
219 |     @property
220 |     def x_pred(self):
221 |         return self.get_x_pred()
222 | 
223 |     def get_dt(self, pruned=True):
224 | 
225 |         return self.dt_raw[tuple(self.pruned_idx)] if pruned else self.dt_raw
226 | 
227 |     @property
228 |     def dt(self):
229 |         return self.get_dt()
230 | 
231 |     def cluster(self, n_clusters, load_clusters=False, save_dir="", visualize_data=False):
232 |         self.n_clusters = n_clusters
233 | 
234 |         x_pruned = self.x
235 |         y_pruned = self.y
236 | 
237 |         file_name = os.path.join(save_dir, 'gmm.pkl')
238 |         loaded = False
239 |         gmm = None
240 |         if os.path.exists(file_name) and load_clusters:
241 |             print("Loaded GP clusters from last GP training session. File: {}".format(file_name))
242 |             gmm = joblib.load(file_name)
243 |             if gmm.n_components != n_clusters:
244 |                 print("The loaded GP does not coincide with the number of set clusters: Found {} but requested"
245 |                       "is {}".format(gmm.n_components, n_clusters))
246 |             else:
247 |                 loaded = True
248 |         if not loaded:
249 |             gmm = GaussianMixture(n_clusters).fit(x_pruned)
250 |         clusters = gmm.predict_proba(x_pruned)
251 |         centroids = gmm.means_
252 | 
253 |         if not load_clusters and n_clusters > 1:
254 |             safe_mknode_recursive(save_dir, 'gmm.pkl', overwrite=True)
255 |             joblib.dump(gmm, file_name)
256 | 
257 |         index_aux = np.arange(0, clusters.shape[0])
258 |         cluster_agency = {}
259 | 
260 |         # Add to the points corresponding to each cluster the points that correspond to it with top 2 probability,
261 |         # if this probability is high
262 |         for cluster in range(n_clusters):
263 |             top_1 = np.argmax(clusters, 1)
264 |             clusters_ = copy.deepcopy(clusters)
265 |             clusters_[index_aux, top_1] *= 0
266 |             top_2 = np.argmax(clusters_, 1)
267 |             idx = np.where(top_1 == cluster)[0]
268 |             idx = np.append(idx, np.where((top_2 == cluster) * (clusters_[index_aux, top_2] > 0.2))[0])
269 |             cluster_agency[cluster] = idx
270 | 
271 |         # Only works in len(x_features) >= 3
272 |         if visualize_data:
273 |             x_aux = self.get_x(pruned=False)
274 |             y_aux = self.get_y(pruned=False)
275 |             visualize_data_distribution(x_aux, y_aux, cluster_agency, x_pruned, y_pruned)
276 | 
277 |         self.cluster_agency = cluster_agency
278 |         self.centroids = centroids
279 | 
280 | 
281 | def restore_gp_regressors(pre_trained_models):
282 |     """
283 |     :param pre_trained_models: A dictionary with all the GP models to be restored in 'models'.
284 |     :return: The GP ensemble restored from the models.
285 |     :rtype: GPEnsemble
286 |     """
287 | 
288 |     gp_reg_ensemble = GPEnsemble()
289 |     # TODO: Deprecate compatibility mode with old models.
290 |     if all(item in list(pre_trained_models.keys()) for item in ["x_features", "u_features"]):
291 |         x_features = pre_trained_models["x_features"]
292 |         u_features = pre_trained_models["u_features"]
293 |     else:
294 |         x_features = u_features = None
295 | 
296 |     if isinstance(pre_trained_models['models'][0], dict):
297 |         pre_trained_gp_reg = {}
298 |         for _, model_dict in enumerate(pre_trained_models['models']):
299 |             if x_features is not None:
300 |                 gp_reg = npGPRegression(x_features, u_features, model_dict["reg_dim"])
301 |             else:
302 |                 gp_reg = npGPRegression(model_dict["x_features"], model_dict["u_features"], model_dict["reg_dim"])
303 |             gp_reg.load(model_dict)
304 |             if model_dict["reg_dim"] not in pre_trained_gp_reg.keys():
305 |                 pre_trained_gp_reg[model_dict["reg_dim"]] = [gp_reg]
306 |             else:
307 |                 pre_trained_gp_reg[model_dict["reg_dim"]] += [gp_reg]
308 | 
309 |         # Add the GP's in a per-output-dim basis into the Ensemble
310 |         for key in np.sort(list(pre_trained_gp_reg.keys())):
311 |             gp_reg_ensemble.add_model(pre_trained_gp_reg[key])
312 |     else:
313 |         raise NotImplementedError("Cannot load this format of GP model.")
314 | 
315 |     return gp_reg_ensemble
316 | 
317 | 
318 | def read_dataset(name, train_split, sim_options):
319 |     """
320 |     Attempts to read a dataset given its name and its metadata.
321 |     :param name: Name of the dataset
322 |     :param train_split: Whether to load the training split (True) or the test split (False)
323 |     :param sim_options: Dictionary of metadata of the dataset to be read.
324 |     :return:
325 |     """
326 |     data_file = get_data_dir_and_file(name, training_split=train_split, params=sim_options, read_only=True)
327 |     if data_file is None:
328 |         raise FileNotFoundError
329 |     rec_file_dir, rec_file_name = data_file
330 |     if rec_file_name.startswith('npz'):
331 |         rec_file_name = rec_file_name[:-4] + '.npz'
332 |         rec_file = os.path.join(rec_file_dir, rec_file_name)
333 |         ds = np.load(rec_file)
334 |     else:
335 |         rec_file = os.path.join(rec_file_dir, rec_file_name)
336 |         ds = pd.read_csv(rec_file)
337 | 
338 |     return ds
339 | 
340 | 
341 | def world_to_body_velocity_mapping(state_sequence):
342 |     """
343 | 
344 |     :param state_sequence: N x 13 state array, where N is the number of states in the sequence.
345 |     :return: An N x 13 sequence of states, but where the velocities (assumed to be in positions 7, 8, 9) have been
346 |     rotated from world to body frame. The rotation is made using the quaternion in positions 3, 4, 5, 6.
347 |     """
348 | 
349 |     p, q, v_w, w = separate_variables(state_sequence)
350 |     v_b = []
351 |     for i in range(len(q)):
352 |         v_b.append(v_dot_q(v_w[i], quaternion_inverse(q[i])))
353 |     v_b = np.stack(v_b)
354 |     return np.concatenate((p, q, v_b, w), 1)


--------------------------------------------------------------------------------
/ros_dd_mpc/src/model_fitting/gp_visualization.py:
--------------------------------------------------------------------------------
  1 | """ Executable script for visual evaluation of the trained GP quality.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | import argparse
 15 | import numpy as np
 16 | import matplotlib.pyplot as plt
 17 | from src.model_fitting.gp_common import GPDataset, restore_gp_regressors, read_dataset
 18 | from config.configuration_parameters import ModelFitConfig as Conf
 19 | from src.utils.utils import load_pickled_models
 20 | from src.utils.visualization import visualize_gp_inference
 21 | 
 22 | 
 23 | def gp_visualization_experiment(quad_sim_options, dataset_name,
 24 |                                 x_cap, hist_bins, hist_thresh,
 25 |                                 x_vis_feats, u_vis_feats, y_vis_feats,
 26 |                                 grid_sampling_viz=False,
 27 |                                 load_model_version="", load_model_name="", pre_set_gp=None):
 28 | 
 29 |     # #### GP ENSEMBLE LOADING #### #
 30 |     if pre_set_gp is None:
 31 |         load_options = {"git": load_model_version, "model_name": load_model_name, "params": quad_sim_options}
 32 |         loaded_models = load_pickled_models(model_options=load_options)
 33 |         if loaded_models is None:
 34 |             raise FileNotFoundError("Model not found")
 35 |         gp_ensemble = restore_gp_regressors(loaded_models)
 36 |     else:
 37 |         gp_ensemble = pre_set_gp
 38 | 
 39 |     # #### DATASET LOADING #### #
 40 |     # Pre-set labels of the data:
 41 |     labels_x = [
 42 |         r'$p_x\:\left[m\right]$', r'$p_y\:\left[m\right]$', r'$p_z\:\left[m\right]$',
 43 |         r'$q_w\:\left[rad\right]$', r'$q_x\:\left[rad\right]$', r'$q_y\:\left[rad\right]$', r'$q_z\:\left[rad\right]$',
 44 |         r'$v_x\:\left[\frac{m}{s}\right]$', r'$v_y\:\left[\frac{m}{s}\right]$', r'$v_z\:\left[\frac{m}{s}\right]$',
 45 |         r'$w_x\:\left[\frac{rad}{s}\right]$', r'$w_y\:\left[\frac{rad}{s}\right]$', r'$w_z\:\left[\frac{rad}{s}\right]$'
 46 |     ]
 47 |     labels_u = [r'$u_1$', r'$u_2$', r'$u_3$', r'$u_4$']
 48 |     labels = [labels_x[feat] for feat in x_vis_feats]
 49 |     labels_ = [labels_u[feat] for feat in u_vis_feats]
 50 |     labels = labels + labels_
 51 | 
 52 |     if isinstance(dataset_name, str):
 53 |         test_ds = read_dataset(dataset_name, True, quad_sim_options)
 54 |         test_gp_ds = GPDataset(test_ds, x_features=x_vis_feats, u_features=u_vis_feats, y_dim=y_vis_feats,
 55 |                                cap=x_cap, n_bins=hist_bins, thresh=hist_thresh, visualize_data=False)
 56 |     else:
 57 |         test_gp_ds = dataset_name
 58 | 
 59 |     x_test = test_gp_ds.get_x(pruned=True, raw=True)
 60 |     u_test = test_gp_ds.get_u(pruned=True, raw=True)
 61 |     y_test = test_gp_ds.get_y(pruned=True, raw=False)
 62 |     dt_test = test_gp_ds.get_dt(pruned=True)
 63 |     x_pred = test_gp_ds.get_x_pred(pruned=True, raw=False)
 64 | 
 65 |     if isinstance(y_vis_feats, list):
 66 |         y_dims = [np.where(gp_ensemble.dim_idx == y_dim)[0][0] for y_dim in y_vis_feats]
 67 |     else:
 68 |         y_dims = np.where(gp_ensemble.dim_idx == y_vis_feats)[0]
 69 | 
 70 |     # #### VISUALIZE GRID SAMPLING RESULTS IN TRAINING SET RANGE #### #
 71 |     if grid_sampling_viz:
 72 |         visualize_gp_inference(x_test, u_test, y_test, gp_ensemble, x_vis_feats, y_dims, labels)
 73 | 
 74 |     # #### EVALUATE GP ON TEST SET #### #
 75 |     print("Test set prediction...")
 76 |     outs = gp_ensemble.predict(x_test.T, u_test.T, return_std=True, progress_bar=True)
 77 |     mean_estimate = np.atleast_2d(np.atleast_2d(outs["pred"])[y_dims])
 78 |     std_estimate = np.atleast_2d(np.atleast_2d(outs["cov_or_std"])[y_dims])
 79 |     mean_estimate = mean_estimate.T * dt_test[:, np.newaxis]
 80 |     std_estimate = std_estimate.T * dt_test[:, np.newaxis]
 81 | 
 82 |     # Undo dt normalization
 83 |     y_test *= dt_test[:, np.newaxis]
 84 | 
 85 |     # Error of nominal model
 86 |     nominal_diff = y_test
 87 | 
 88 |     # GP regresses model error, correct the predictions of the nominal model
 89 |     augmented_diff = nominal_diff - mean_estimate
 90 |     mean_estimate += x_pred
 91 | 
 92 |     nominal_rmse = np.mean(np.abs(nominal_diff), 0)
 93 |     augmented_rmse = np.mean(np.abs(augmented_diff), 0)
 94 | 
 95 |     labels = [r'$v_x$ error', r'$v_y$ error', r'$v_z$ error']
 96 |     t_vec = np.cumsum(dt_test)
 97 |     plt.figure()
 98 |     # mng = plt.get_current_fig_manager()
 99 |     # mng.resize(*mng.window.maxsize())
100 |     for i in range(std_estimate.shape[1]):
101 |         plt.subplot(std_estimate.shape[1], 1, i+1)
102 |         plt.plot(t_vec, np.zeros(augmented_diff[:, i].shape), 'k')
103 |         plt.plot(t_vec, augmented_diff[:, i], 'b', label='augmented_err')
104 |         plt.plot(t_vec, augmented_diff[:, i] - 3 * std_estimate[:, i], ':c')
105 |         plt.plot(t_vec, augmented_diff[:, i] + 3 * std_estimate[:, i], ':c', label='3 std')
106 |         if nominal_diff is not None:
107 |             plt.plot(t_vec, nominal_diff[:, i], 'r', label='nominal_err')
108 |             plt.title('Mean dt: %.2f s. Nominal RMSE: %.5f [m/s].  Augmented rmse: %.5f [m/s]' % (
109 |                 float(np.mean(dt_test)), nominal_rmse[i], augmented_rmse[i]))
110 |         else:
111 |             plt.title('Mean dt: %.2f s. Augmented RMSE: %.5f [m/s]' % (
112 |                 float(np.mean(dt_test)), float(augmented_rmse[i])))
113 | 
114 |         plt.ylabel(labels[i])
115 |         plt.legend()
116 | 
117 |         if i == std_estimate.shape[1] - 1:
118 |             plt.xlabel('time (s)')
119 | 
120 |     plt.show()
121 | 
122 | 
123 | if __name__ == '__main__':
124 | 
125 |     parser = argparse.ArgumentParser()
126 | 
127 |     parser.add_argument("--model_version", type=str, default="",
128 |                         help="Version to load for the regression models. By default it is an 8 digit git hash.")
129 | 
130 |     parser.add_argument("--model_name", type=str, default="",
131 |                         help="Name of the regression model within the specified <model_version> folder.")
132 | 
133 |     input_arguments = parser.parse_args()
134 | 
135 |     ds_name = Conf.ds_name
136 |     simulation_options = Conf.ds_metadata
137 | 
138 |     histogram_pruning_bins = Conf.histogram_bins
139 |     histogram_pruning_threshold = Conf.histogram_threshold
140 |     x_value_cap = Conf.velocity_cap
141 | 
142 |     x_viz = Conf.x_viz
143 |     u_viz = Conf.u_viz
144 |     y_viz = Conf.y_viz
145 |     gp_load_model_name = input_arguments.model_name
146 |     gp_load_model_version = input_arguments.model_version
147 |     gp_visualization_experiment(simulation_options, ds_name,
148 |                                 x_value_cap, histogram_pruning_bins, histogram_pruning_threshold,
149 |                                 x_vis_feats=x_viz, u_vis_feats=u_viz, y_vis_feats=y_viz,
150 |                                 grid_sampling_viz=True,
151 |                                 load_model_version=gp_load_model_version,
152 |                                 load_model_name=gp_load_model_name)


--------------------------------------------------------------------------------
/ros_dd_mpc/src/model_fitting/mlp_common.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | from torch.utils.data import Dataset
  3 | 
  4 | import ml_casadi.torch as mc
  5 | from config.configuration_parameters import GroundEffectMapConfig
  6 | from src.utils.ground_map import GroundMapWithBox
  7 | 
  8 | 
  9 | class GPToMLPDataset(Dataset):
 10 |     def __init__(self, gp_dataset, ground_effect=False):
 11 |         super().__init__()
 12 |         self.x = gp_dataset.x
 13 |         self.y = gp_dataset.y
 14 |         if ground_effect:
 15 |             self.x_raw = gp_dataset.x_raw[tuple(gp_dataset.pruned_idx)]
 16 |             assert self.x.shape[0] == self.x_raw.shape[0]
 17 |             map_conf = GroundEffectMapConfig
 18 |             ground_map = GroundMapWithBox(np.array(map_conf.box_min),
 19 |                                           np.array(map_conf.box_max),
 20 |                                           map_conf.box_height,
 21 |                                           horizon=map_conf.horizon,
 22 |                                           resolution=map_conf.resolution)
 23 |             self._map_res = map_conf.resolution
 24 |             self._static_ground_map, self._org_to_map_org = ground_map.at(np.array(map_conf.origin))
 25 |         else:
 26 |             self._static_ground_map = None
 27 | 
 28 |     def stats(self):
 29 |         if self._static_ground_map is None:
 30 |             return self.x.mean(axis=0), self.x.std(axis=0), self.y.mean(axis=0), self.y.std(axis=0)
 31 |         else:
 32 |             x_mean = np.hstack([self.x.mean(axis=0), np.zeros(9+4)])
 33 |             x_std = np.hstack([self.x.std(axis=0), np.ones(9+4)])
 34 |             return x_mean, x_std, self.y.mean(axis=0), self.y.std(axis=0)
 35 | 
 36 |     def __len__(self):
 37 |         return len(self.x)
 38 | 
 39 |     def __getitem__(self, item):
 40 |         if self._static_ground_map is None:
 41 |             return self.x[item], self.y[item]
 42 |         else:
 43 |             x, y, z = self.x_raw[item][:3]
 44 |             orientation = self.x_raw[item][3:7]
 45 |             x_idxs = np.floor((x - self._org_to_map_org[0]) / self._map_res).astype(int) - 1
 46 |             y_idxs = np.floor((y - self._org_to_map_org[1]) / self._map_res).astype(int) - 1
 47 |             ground_patch = self._static_ground_map[x_idxs:x_idxs+3, y_idxs:y_idxs+3]
 48 | 
 49 |             relative_ground_patch = 4 * (np.clip(z - ground_patch, 0, 0.5) - 0.25)
 50 | 
 51 |             flatten_relative_ground_patch = relative_ground_patch.flatten(order='F')
 52 | 
 53 |             ground_effect_in = np.hstack([flatten_relative_ground_patch, orientation*0])
 54 | 
 55 |             return np.hstack([self.x[item], ground_effect_in]), self.y[item]
 56 | 
 57 | 
 58 | class NormalizedMLP(mc.TorchMLCasadiModule):
 59 |     def __init__(self, model, x_mean, x_std, y_mean, y_std):
 60 |         super().__init__()
 61 |         self.model = model
 62 |         self.input_size = self.model.input_size
 63 |         self.output_size = self.model.output_size
 64 |         self.register_buffer('x_mean', x_mean)
 65 |         self.register_buffer('x_std', x_std)
 66 |         self.register_buffer('y_mean', y_mean)
 67 |         self.register_buffer('y_std', y_std)
 68 | 
 69 |     def forward(self, x):
 70 |         return (self.model((x - self.x_mean) / self.x_std) * self.y_std) + self.y_mean
 71 | 
 72 |     def cs_forward(self, x):
 73 |         return (self.model((x - self.x_mean.cpu().numpy()) / self.x_std.cpu().numpy()) * self.y_std.cpu().numpy()) + self.y_mean.cpu().numpy()
 74 | 
 75 | 
 76 | class QuadResidualModel(mc.TorchMLCasadiModule):
 77 |     def __init__(self, hidden_size, hidden_layers):
 78 |         super().__init__()
 79 |         self._input_size = 10
 80 |         self._output_size = 6
 81 |         self.force_model = mc.nn.MultiLayerPerceptron(7, hidden_size, 3, hidden_layers, 'Tanh')
 82 |         self.torque_model = mc.nn.MultiLayerPerceptron(10, hidden_size, 3, hidden_layers, 'Tanh')
 83 | 
 84 |     def forward(self, x):
 85 |         v = x[:, :3]
 86 |         w = x[:, 3:6]
 87 |         u = x[:, 6:]
 88 |         force_in = mc.hcat([v, u])
 89 |         torque_in = mc.hcat([v, w, u])
 90 |         force_out = self.force_model(force_in)
 91 |         torque_out = self.torque_model(torque_in)
 92 |         out = mc.hcat([force_out, torque_out])
 93 |         return out
 94 | 
 95 |     def cs_forward(self, x):
 96 |         v = x[:3]
 97 |         w = x[3:6]
 98 |         u = x[6:]
 99 |         force_in = mc.vcat([v, u])
100 |         torque_in = mc.vcat([v, w, u])
101 |         force_out = self.force_model(force_in)
102 |         torque_out = self.torque_model(torque_in)
103 |         out = mc.vcat([force_out, torque_out])
104 |         return out
105 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/model_fitting/mlp_fitting.py:
--------------------------------------------------------------------------------
  1 | """ Executable script to train a custom Gaussian Process on recorded flight data.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | 
 15 | import os
 16 | import time
 17 | import argparse
 18 | import subprocess
 19 | import numpy as np
 20 | 
 21 | import torch
 22 | from torch.utils.data import DataLoader
 23 | import ml_casadi.torch as mc
 24 | 
 25 | from tqdm import tqdm
 26 | 
 27 | from src.model_fitting.mlp_common import GPToMLPDataset, NormalizedMLP
 28 | from src.utils.utils import safe_mkdir_recursive, load_pickled_models
 29 | from src.utils.utils import get_model_dir_and_file
 30 | from src.model_fitting.gp_common import GPDataset, read_dataset
 31 | from config.configuration_parameters import ModelFitConfig as Conf
 32 | 
 33 | 
 34 | 
 35 | def main(x_features, u_features, reg_y_dims, model_ground_effect, quad_sim_options, dataset_name,
 36 |          x_cap, hist_bins, hist_thresh,
 37 |          model_name="model", epochs=100, batch_size=64, hidden_layers=1, hidden_size=32, plot=False):
 38 | 
 39 |     """
 40 |     Reads the dataset specified and trains a GP model or ensemble on it. The regressed variables is the time-derivative
 41 |     of one of the states.
 42 |     The feature and regressed variables are identified through the index number in the state vector. It is defined as:
 43 |     [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]
 44 |     The input vector is also indexed:
 45 |     [0: u_0, 1: u_1, 2: u_2, 3: u_3].
 46 | 
 47 |     :param x_features: List of n regression feature indices from the quadrotor state indexing.
 48 |     :type x_features: list
 49 |     :param u_features: List of n' regression feature indices from the input state.
 50 |     :type u_features: list
 51 |     :param reg_y_dims: Indices of output dimension being regressed as the time-derivative.
 52 |     :type reg_y_dims: List
 53 |     :param dataset_name: Name of the dataset
 54 |     :param quad_sim_options: Dictionary of metadata of the dataset to be read.
 55 |     :param x_cap: cap value (in absolute value) for dataset pruning. Any input feature that exceeds this number in any
 56 |     dimension will be removed.
 57 |     :param hist_bins: Number of bins used for histogram pruning
 58 |     :param hist_thresh: Any bin with less data percentage than this number will be removed.
 59 |     :param model_name: Given name to the currently trained GP.
 60 |     """
 61 | 
 62 |     # Get git commit hash for saving the model
 63 |     git_version = ''
 64 |     try:
 65 |         git_version = subprocess.check_output(['git', 'describe', '--always']).strip().decode("utf-8")
 66 |     except subprocess.CalledProcessError as e:
 67 |         print(e.returncode, e.output)
 68 |     print("The model will be saved using hash: %s" % git_version)
 69 | 
 70 |     gp_name_dict = {"git": git_version, "model_name": model_name, "params": quad_sim_options}
 71 |     save_file_path, save_file_name = get_model_dir_and_file(gp_name_dict)
 72 |     safe_mkdir_recursive(save_file_path)
 73 | 
 74 |     # #### DATASET LOADING #### #
 75 |     if isinstance(dataset_name, str):
 76 |         df_train = read_dataset(dataset_name, True, quad_sim_options)
 77 |         gp_dataset_train = GPDataset(df_train, x_features, u_features, reg_y_dims,
 78 |                                      cap=x_cap, n_bins=hist_bins, thresh=hist_thresh)
 79 |         gp_dataset_val = None
 80 |         try:
 81 |             df_val = read_dataset(dataset_name, False, quad_sim_options)
 82 |             gp_dataset_val = GPDataset(df_val, x_features, u_features, reg_y_dims,
 83 |                                        cap=x_cap, n_bins=hist_bins, thresh=hist_thresh)
 84 |         except:
 85 |             print('Could not find test dataset.')
 86 |     else:
 87 |         raise TypeError("dataset_name must be a string.")
 88 | 
 89 |     dataset_train = GPToMLPDataset(gp_dataset_train, ground_effect=model_ground_effect)
 90 |     x_mean, x_std, y_mean, y_std = dataset_train.stats()
 91 |     data_loader = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=0)
 92 |     input_dims = len(x_features) + len(u_features) + (9 + 4 if model_ground_effect else 0)
 93 |     mlp_model = mc.nn.MultiLayerPerceptron(input_dims, hidden_size, len(reg_y_dims), hidden_layers, 'Tanh')
 94 |     model = NormalizedMLP(mlp_model, torch.tensor(x_mean).float(), torch.tensor(x_std).float(),
 95 |                           torch.tensor(y_mean).float(), torch.tensor(y_std).float())
 96 |     print(f"Train Dataset has {len(dataset_train)} points.")
 97 | 
 98 |     if gp_dataset_val:
 99 |         dataset_val = GPToMLPDataset(gp_dataset_val, ground_effect=model_ground_effect)
100 |         data_loader_val = DataLoader(dataset_val, batch_size=10*batch_size, shuffle=False, num_workers=0)
101 |         print(f"Val Dataset has {len(dataset_val)} points.")
102 | 
103 |     print(f"Model has {sum(p.numel() for p in model.parameters() if p.requires_grad)} parameters.")
104 | 
105 |     if torch.cuda.is_available():
106 |         model = model.to('cuda:0')
107 | 
108 |     optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
109 | 
110 |     loss_infos = []
111 |     p_bar = tqdm(range(epochs))
112 |     for i in p_bar:
113 |         model.train()
114 |         losses = []
115 |         for x, y in data_loader:
116 |             if torch.cuda.is_available():
117 |                 x = x.to('cuda:0')
118 |                 y = y.to('cuda:0')
119 |             x = x.float()
120 |             y = y.float()
121 |             optimizer.zero_grad()
122 |             y_pred = model(x)
123 |             loss = torch.square(y - y_pred).mean()
124 |             loss.backward()
125 |             optimizer.step()
126 | 
127 |             losses.append(loss.item())
128 | 
129 |         losses_val = []
130 |         if gp_dataset_val:
131 |             with torch.no_grad():
132 |                 for x, y in data_loader_val:
133 |                     if torch.cuda.is_available():
134 |                         x = x.to('cuda:0')
135 |                         y = y.to('cuda:0')
136 |                     x = x.float()
137 |                     y = y.float()
138 |                     y_pred = model(x)
139 |                     loss = torch.square(y - y_pred)
140 | 
141 |                     losses_val.append(loss.numpy())
142 |         train_loss_info = np.mean(losses)
143 |         loss_info = np.mean(np.vstack(losses_val)) if losses_val else np.mean(losses)
144 |         loss_infos.append(loss_info)
145 |         p_bar.set_description(f'Train Loss: {train_loss_info}, Val Loss {loss_info:.6f}')
146 |         p_bar.refresh()
147 | 
148 |         save_dict = {
149 |             'state_dict': model.state_dict(),
150 |             'input_size': input_dims,
151 |             'hidden_size': hidden_size,
152 |             'output_size': len(reg_y_dims),
153 |             'hidden_layers': hidden_layers
154 |         }
155 |         torch.save(save_dict, os.path.join(save_file_path, f'{save_file_name}.pt'))
156 | 
157 |     if plot:
158 |         import matplotlib.pyplot as plt
159 |         plt.plot(loss_infos)
160 |         plt.show()
161 | 
162 | 
163 | if __name__ == '__main__':
164 | 
165 |     parser = argparse.ArgumentParser()
166 | 
167 |     parser.add_argument("--epochs", type=int, default="1000",
168 |                         help="Number of epochs to train the model.")
169 | 
170 |     parser.add_argument("--batch_size", type=int, default="64",
171 |                         help="Batch Size.")
172 | 
173 |     parser.add_argument("--hidden_layers", type=int, default="1",
174 |                         help="Number of hidden layers.")
175 | 
176 |     parser.add_argument("--hidden_size", type=int, default="32",
177 |                         help="Size of hidden layers.")
178 | 
179 |     parser.add_argument("--model_name", type=str, default="",
180 |                         help="Name assigned to the trained model.")
181 | 
182 |     parser.add_argument('--x', nargs='+', type=int, default=[7],
183 |                         help='Regression X variables. Must be a list of integers between 0 and 12. Velocities xyz '
184 |                              'correspond to indices 7, 8, 9.')
185 | 
186 |     parser.add_argument('--u', action="store_true",
187 |                         help='Use the control as input to the model.')
188 | 
189 |     parser.add_argument("--y", nargs='+', type=int, default=[7],
190 |                         help="Regression Y variable. Must be an integer between 0 and 12. Velocities xyz correspond to"
191 |                              "indices 7, 8, 9.")
192 | 
193 |     parser.add_argument('--ge', action="store_true",
194 |                         help='Use the ground map as input to the model.')
195 | 
196 |     parser.add_argument("--plot", dest="plot", action="store_true",
197 |                         help="Plot the loss after training.")
198 |     parser.set_defaults(plot=False)
199 | 
200 |     input_arguments = parser.parse_args()
201 | 
202 |     # Use vx, vy, vz as input features
203 |     x_feats = input_arguments.x
204 |     if input_arguments.u:
205 |         u_feats = [0, 1, 2, 3]
206 |     else:
207 |         u_feats = []
208 | 
209 |     model_ground_effect = input_arguments.ge
210 | 
211 |     # Regression dimension
212 |     y_regressed_dims = input_arguments.y
213 |     model_name = input_arguments.model_name
214 | 
215 |     epochs = input_arguments.epochs
216 |     batch_size = input_arguments.batch_size
217 |     hidden_layers = input_arguments.hidden_layers
218 |     hidden_size = input_arguments.hidden_size
219 |     plot = input_arguments.plot
220 | 
221 |     ds_name = Conf.ds_name
222 |     simulation_options = Conf.ds_metadata
223 | 
224 |     histogram_pruning_bins = Conf.histogram_bins
225 |     histogram_pruning_threshold = Conf.histogram_threshold
226 |     x_value_cap = Conf.velocity_cap
227 | 
228 |     main(x_feats, u_feats, y_regressed_dims, model_ground_effect, simulation_options, ds_name,
229 |          x_value_cap, histogram_pruning_bins, histogram_pruning_threshold,
230 |          model_name=model_name, epochs=epochs, batch_size=batch_size, hidden_layers=hidden_layers,
231 |          hidden_size=hidden_size, plot=plot)
232 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/model_fitting/mlp_quad_res_fitting.py:
--------------------------------------------------------------------------------
  1 | """ Executable script to train a custom Gaussian Process on recorded flight data.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | 
 15 | import os
 16 | import time
 17 | import argparse
 18 | import subprocess
 19 | import numpy as np
 20 | 
 21 | import torch
 22 | from torch.utils.data import DataLoader
 23 | import ml_casadi.torch as mc
 24 | 
 25 | from tqdm import tqdm
 26 | 
 27 | from src.model_fitting.mlp_common import GPToMLPDataset, NormalizedMLP, QuadResidualModel
 28 | from src.utils.utils import safe_mkdir_recursive, load_pickled_models
 29 | from src.utils.utils import get_model_dir_and_file
 30 | from src.model_fitting.gp_common import GPDataset, read_dataset
 31 | from config.configuration_parameters import ModelFitConfig as Conf
 32 | 
 33 | 
 34 | 
 35 | def main(x_features, u_features, reg_y_dims, quad_sim_options, dataset_name,
 36 |          x_cap, hist_bins, hist_thresh,
 37 |          model_name="model", epochs=100, batch_size=64, hidden_layers=1, hidden_size=32, plot=False):
 38 | 
 39 |     """
 40 |     Reads the dataset specified and trains a GP model or ensemble on it. The regressed variables is the time-derivative
 41 |     of one of the states.
 42 |     The feature and regressed variables are identified through the index number in the state vector. It is defined as:
 43 |     [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]
 44 |     The input vector is also indexed:
 45 |     [0: u_0, 1: u_1, 2: u_2, 3: u_3].
 46 | 
 47 |     :param x_features: List of n regression feature indices from the quadrotor state indexing.
 48 |     :type x_features: list
 49 |     :param u_features: List of n' regression feature indices from the input state.
 50 |     :type u_features: list
 51 |     :param reg_y_dims: Indices of output dimension being regressed as the time-derivative.
 52 |     :type reg_y_dims: List
 53 |     :param dataset_name: Name of the dataset
 54 |     :param quad_sim_options: Dictionary of metadata of the dataset to be read.
 55 |     :param x_cap: cap value (in absolute value) for dataset pruning. Any input feature that exceeds this number in any
 56 |     dimension will be removed.
 57 |     :param hist_bins: Number of bins used for histogram pruning
 58 |     :param hist_thresh: Any bin with less data percentage than this number will be removed.
 59 |     :param model_name: Given name to the currently trained GP.
 60 |     """
 61 | 
 62 |     # Get git commit hash for saving the model
 63 |     git_version = ''
 64 |     try:
 65 |         git_version = subprocess.check_output(['git', 'describe', '--always']).strip().decode("utf-8")
 66 |     except subprocess.CalledProcessError as e:
 67 |         print(e.returncode, e.output)
 68 |     print("The model will be saved using hash: %s" % git_version)
 69 | 
 70 |     gp_name_dict = {"git": git_version, "model_name": model_name, "params": quad_sim_options}
 71 |     save_file_path, save_file_name = get_model_dir_and_file(gp_name_dict)
 72 |     safe_mkdir_recursive(save_file_path)
 73 | 
 74 |     # #### DATASET LOADING #### #
 75 |     if isinstance(dataset_name, str):
 76 |         df_train = read_dataset(dataset_name, True, quad_sim_options)
 77 |         gp_dataset_train = GPDataset(df_train, x_features, u_features, reg_y_dims,
 78 |                                      cap=x_cap, n_bins=hist_bins, thresh=hist_thresh)
 79 |         gp_dataset_val = None
 80 |         try:
 81 |             df_val = read_dataset(dataset_name, False, quad_sim_options)
 82 |             gp_dataset_val = GPDataset(df_val, x_features, u_features, reg_y_dims,
 83 |                                        cap=x_cap, n_bins=hist_bins, thresh=hist_thresh)
 84 |         except:
 85 |             print('Could not find test dataset.')
 86 |     else:
 87 |         raise TypeError("dataset_name must be a string.")
 88 | 
 89 |     dataset_train = GPToMLPDataset(gp_dataset_train)
 90 |     x_mean, x_std, y_mean, y_std = dataset_train.stats()
 91 |     data_loader = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=0)
 92 |     input_dims = len(x_features) + len(u_features)
 93 |     mlp_model = QuadResidualModel(hidden_size, hidden_layers)
 94 |     model = NormalizedMLP(mlp_model, torch.tensor(x_mean).float(), torch.tensor(x_std).float(),
 95 |                           torch.tensor(y_mean).float(), torch.tensor(y_std).float())
 96 | 
 97 |     print(f"Train Dataset has {len(dataset_train)} points.")
 98 | 
 99 |     if gp_dataset_val:
100 |         dataset_val = GPToMLPDataset(gp_dataset_val)
101 |         data_loader_val = DataLoader(dataset_val, batch_size=10*batch_size, shuffle=False, num_workers=0)
102 |         print(f"Val Dataset has {len(dataset_val)} points.")
103 | 
104 |     print(f"Model has {sum(p.numel() for p in model.parameters() if p.requires_grad)} parameters.")
105 | 
106 |     if torch.cuda.is_available():
107 |         print('Pushing model to GPU.')
108 |         model = model.to('cuda:0')
109 | 
110 |     optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
111 | 
112 |     loss_infos = []
113 |     p_bar = tqdm(range(epochs))
114 |     for i in p_bar:
115 |         model.train()
116 |         losses = []
117 |         for x, y in data_loader:
118 |             if torch.cuda.is_available():
119 |                 x = x.to('cuda:0')
120 |                 y = y.to('cuda:0')
121 |             x = x.float()
122 |             y = y.float()
123 |             optimizer.zero_grad()
124 |             y_pred = model(x)
125 |             loss = torch.square(y - y_pred).mean()
126 |             loss.backward()
127 |             optimizer.step()
128 | 
129 |             losses.append(loss.item())
130 | 
131 |         losses_val = []
132 |         if gp_dataset_val:
133 |             with torch.no_grad():
134 |                 for x, y in data_loader_val:
135 |                     if torch.cuda.is_available():
136 |                         x = x.to('cuda:0')
137 |                         y = y.to('cuda:0')
138 |                     x = x.float()
139 |                     y = y.float()
140 |                     y_pred = model(x)
141 |                     loss = torch.square(y - y_pred)
142 | 
143 |                     losses_val.append(loss.cpu().numpy())
144 |         train_loss_info = np.mean(losses)
145 |         loss_info = np.mean(np.vstack(losses_val)) if losses_val else np.mean(losses)
146 |         loss_infos.append(loss_info)
147 |         p_bar.set_description(f'Train Loss: {train_loss_info}, Val Loss {loss_info:.6f}')
148 |         p_bar.refresh()
149 | 
150 |         save_dict = {
151 |             'state_dict': model.state_dict(),
152 |             'input_size': input_dims,
153 |             'hidden_size': hidden_size,
154 |             'output_size': len(reg_y_dims),
155 |             'hidden_layers': hidden_layers
156 |         }
157 |         torch.save(save_dict, os.path.join(save_file_path, f'{save_file_name}_{i}.pt'))
158 |         print(f'{i}: Val Loss {loss_info:.6f}')
159 | 
160 |     if plot:
161 |         import matplotlib.pyplot as plt
162 |         plt.plot(loss_infos)
163 |         plt.show()
164 | 
165 | 
166 | if __name__ == '__main__':
167 | 
168 |     parser = argparse.ArgumentParser()
169 | 
170 |     parser.add_argument("--epochs", type=int, default="1000",
171 |                         help="Number of epochs to train the model.")
172 | 
173 |     parser.add_argument("--batch_size", type=int, default="64",
174 |                         help="Batch Size.")
175 | 
176 |     parser.add_argument("--hidden_layers", type=int, default="1",
177 |                         help="Number of hidden layers.")
178 | 
179 |     parser.add_argument("--hidden_size", type=int, default="32",
180 |                         help="Size of hidden layers.")
181 | 
182 |     parser.add_argument("--model_name", type=str, default="",
183 |                         help="Name assigned to the trained model.")
184 | 
185 |     parser.add_argument("--plot", dest="plot", action="store_true",
186 |                         help="Plot the loss after training.")
187 |     parser.set_defaults(plot=False)
188 | 
189 |     input_arguments = parser.parse_args()
190 | 
191 |     # Use vx, vy, vz as input features
192 |     x_feats = [7, 8, 9, 10, 11, 12]
193 |     u_feats = [0, 1, 2, 3]
194 | 
195 |     # Regression dimension
196 |     y_regressed_dims = [7, 8, 9, 10, 11, 12]
197 |     model_name = input_arguments.model_name
198 | 
199 |     epochs = input_arguments.epochs
200 |     batch_size = input_arguments.batch_size
201 |     hidden_layers = input_arguments.hidden_layers
202 |     hidden_size = input_arguments.hidden_size
203 |     plot = input_arguments.plot
204 | 
205 |     ds_name = Conf.ds_name
206 |     simulation_options = Conf.ds_metadata
207 | 
208 |     histogram_pruning_bins = Conf.histogram_bins
209 |     histogram_pruning_threshold = Conf.histogram_threshold
210 |     x_value_cap = Conf.velocity_cap
211 | 
212 |     main(x_feats, u_feats, y_regressed_dims, simulation_options, ds_name,
213 |          x_value_cap, histogram_pruning_bins, histogram_pruning_threshold,
214 |          model_name=model_name, epochs=epochs, batch_size=batch_size, hidden_layers=hidden_layers,
215 |          hidden_size=hidden_size, plot=plot)
216 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/model_fitting/process_neurobem_dataset.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import os
  3 | import random
  4 | 
  5 | import numpy as np
  6 | import pandas as pd
  7 | from scipy.signal import savgol_filter
  8 | from tqdm import tqdm
  9 | 
 10 | from config.configuration_parameters import DirectoryConfig
 11 | from config.configuration_parameters import ModelFitConfig as Conf
 12 | from src.experiments.point_tracking_and_record import make_record_dict, jsonify
 13 | from src.quad_mpc.create_ros_dd_mpc import custom_quad_param_loader
 14 | from src.quad_mpc.quad_3d_mpc import Quad3DMPC
 15 | from src.utils.utils import safe_mkdir_recursive, v_dot_q
 16 | 
 17 | val_files = ['merged_2021-02-03-16-58-13_seg_2.csv']
 18 | 
 19 | def main(quad):
 20 |     full_path = os.path.join(DirectoryConfig.DATA_DIR, 'neurobem_processed_data_cmd')
 21 |     files = os.listdir(full_path)
 22 |     random.shuffle(files)
 23 | 
 24 |     rec_dicts = []
 25 |     for file in tqdm(files):
 26 |         if file in val_files:
 27 |             print('Skipping Validation File')
 28 |             continue
 29 |         try:
 30 |             rec_dict = make_record_dict(state_dim=13)
 31 |             process_file(os.path.join(full_path, file), quad, rec_dict)
 32 |             rec_dicts.append(rec_dict)
 33 |         except Exception as e:
 34 |             print(e)
 35 | 
 36 |     rec_dict = {}
 37 |     rec_dict['state_in'] = np.concatenate([data_dict['state_in'] for data_dict in rec_dicts], axis=0)
 38 |     rec_dict['input_in'] = np.concatenate([data_dict['input_in'] for data_dict in rec_dicts], axis=0)
 39 |     rec_dict['state_out'] = np.concatenate([data_dict['state_out'] for data_dict in rec_dicts], axis=0)
 40 |     rec_dict['state_ref'] = np.concatenate([data_dict['state_ref'] for data_dict in rec_dicts], axis=0)
 41 |     rec_dict['timestamp'] = np.concatenate([data_dict['timestamp'] for data_dict in rec_dicts], axis=0)
 42 |     rec_dict['dt'] = np.concatenate([data_dict['dt'] for data_dict in rec_dicts], axis=0)
 43 |     rec_dict['state_pred'] = np.concatenate([data_dict['state_pred'] for data_dict in rec_dicts], axis=0)
 44 |     rec_dict['error'] = np.concatenate([data_dict['error'] for data_dict in rec_dicts], axis=0)
 45 | 
 46 |     del rec_dicts
 47 | 
 48 |     # Save datasets
 49 |     save_file_folder = os.path.join(DirectoryConfig.DATA_DIR, 'neurobem_dataset', 'train')
 50 |     safe_mkdir_recursive(save_file_folder)
 51 |     save_file = os.path.join(save_file_folder, f'npz_dataset_001.npz')
 52 |     np.savez(save_file, **rec_dict)
 53 | 
 54 | 
 55 |     # # Validation
 56 |     # rec_dicts = []
 57 |     # for file in tqdm(files[-20:]):
 58 |     #     try:
 59 |     #         rec_dict = make_record_dict(state_dim=13)
 60 |     #         process_file(os.path.join(full_path, file), quad, rec_dict)
 61 |     #         rec_dicts.append(rec_dict)
 62 |     #     except Exception as e:
 63 |     #         print(e)
 64 |     #
 65 |     # rec_dict = {}
 66 |     # rec_dict['state_in'] = np.concatenate([data_dict['state_in'] for data_dict in rec_dicts], axis=0)
 67 |     # rec_dict['input_in'] = np.concatenate([data_dict['input_in'] for data_dict in rec_dicts], axis=0)
 68 |     # rec_dict['state_out'] = np.concatenate([data_dict['state_out'] for data_dict in rec_dicts], axis=0)
 69 |     # rec_dict['state_ref'] = np.concatenate([data_dict['state_ref'] for data_dict in rec_dicts], axis=0)
 70 |     # rec_dict['timestamp'] = np.concatenate([data_dict['timestamp'] for data_dict in rec_dicts], axis=0)
 71 |     # rec_dict['dt'] = np.concatenate([data_dict['dt'] for data_dict in rec_dicts], axis=0)
 72 |     # rec_dict['state_pred'] = np.concatenate([data_dict['state_pred'] for data_dict in rec_dicts], axis=0)
 73 |     # rec_dict['error'] = np.concatenate([data_dict['error'] for data_dict in rec_dicts], axis=0)
 74 |     #
 75 |     # del rec_dicts
 76 |     #
 77 |     # # Save datasets
 78 |     # save_file_folder = os.path.join(DirectoryConfig.DATA_DIR, 'neurobem_dataset', 'test')
 79 |     # safe_mkdir_recursive(save_file_folder)
 80 |     # save_file = os.path.join(save_file_folder, f'npz_dataset_001.npz')
 81 |     # np.savez(save_file, **rec_dict)
 82 | 
 83 | def quaternion_to_euler(q):
 84 |     yaw = np.arctan2(2 * (q[:, 0] * q[:, 3] - q[:, 1] * q[:, 2]),
 85 |         1 - 2 * (q[:, 2] ** 2 + q[:, 3] ** 2))
 86 |     pitch = np.arcsin(2 * (q[:, 0] * q[:, 2] + q[:, 3] * q[:, 1]))
 87 |     roll = np.arctan2(2 * (q[:, 0] * q[:, 1] - q[:, 2] * q[:, 3]),
 88 |         1 - 2 * (q[:, 1] ** 2 + q[:, 2] ** 2))
 89 | 
 90 |     return np.stack([yaw, pitch, roll], axis=-1)
 91 | 
 92 | def val(quad):
 93 |     full_path = os.path.join(DirectoryConfig.DATA_DIR, 'neurobem_processed_data_cmd')
 94 |     #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',
 95 | 
 96 |     rec_dicts = []
 97 |     for file in tqdm(val_files):
 98 |         try:
 99 |             rec_dict = make_record_dict(state_dim=13)
100 |             process_file(os.path.join(full_path, file), quad, rec_dict)
101 |             rec_dicts.append(rec_dict)
102 |         except Exception as e:
103 |             print(e)
104 | 
105 |     rec_dict = {}
106 |     rec_dict['state_in'] = np.concatenate([data_dict['state_in'] for data_dict in rec_dicts], axis=0)
107 |     rec_dict['input_in'] = np.concatenate([data_dict['input_in'] for data_dict in rec_dicts], axis=0)
108 |     rec_dict['state_out'] = np.concatenate([data_dict['state_out'] for data_dict in rec_dicts], axis=0)
109 |     rec_dict['state_ref'] = np.concatenate([data_dict['state_ref'] for data_dict in rec_dicts], axis=0)
110 |     rec_dict['timestamp'] = np.concatenate([data_dict['timestamp'] for data_dict in rec_dicts], axis=0)
111 |     rec_dict['dt'] = np.concatenate([data_dict['dt'] for data_dict in rec_dicts], axis=0)
112 |     rec_dict['state_pred'] = np.concatenate([data_dict['state_pred'] for data_dict in rec_dicts], axis=0)
113 |     rec_dict['error'] = np.concatenate([data_dict['error'] for data_dict in rec_dicts], axis=0)
114 | 
115 |     del rec_dicts
116 | 
117 |     # Save datasets
118 |     save_file_folder = os.path.join(DirectoryConfig.DATA_DIR, 'neurobem_dataset', 'test')
119 |     safe_mkdir_recursive(save_file_folder)
120 |     save_file = os.path.join(save_file_folder, f'npz_dataset_001.npz')
121 |     np.savez(save_file, **rec_dict)
122 | 
123 | def process_file(file_path, quad, rec_dict):
124 |     data = pd.read_csv(file_path, encoding='latin-1')
125 |     for x_0, x_f, u, dt in consecutive_data_points(data):
126 |         resimulate(x_0, x_f, u, dt, quad, rec_dict)
127 | 
128 | 
129 | def consecutive_data_points(data):
130 |     for i in range(25, len(data)-1, 1):
131 |         data_0 = data.iloc[i]
132 |         data_1 = data.iloc[i+1]
133 |         data_c = data.iloc[i]  # Communication Delay
134 |         t0 = data_0['t']
135 |         t1 = data_1['t']
136 |         dt = t1 - t0
137 | 
138 |         x_0 = np.hstack([
139 |             data_0['pos x'],
140 |             data_0['pos y'],
141 |             data_0['pos z'],
142 |             data_0['quat w'],
143 |             data_0['quat x'],
144 |             data_0['quat y'],
145 |             data_0['quat z'],
146 |             data_0['vel x'],
147 |             data_0['vel y'],
148 |             data_0['vel z'],
149 |             data_0['ang vel x'],
150 |             data_0['ang vel y'],
151 |             data_0['ang vel z'],
152 |         ])
153 | 
154 |         x_1 = np.hstack([
155 |             data_1['pos x'],
156 |             data_1['pos y'],
157 |             data_1['pos z'],
158 |             data_1['quat w'],
159 |             data_1['quat x'],
160 |             data_1['quat y'],
161 |             data_1['quat z'],
162 |             data_1['vel x'],
163 |             data_1['vel y'],
164 |             data_1['vel z'],
165 |             data_1['ang vel x'],
166 |             data_1['ang vel y'],
167 |             data_1['ang vel z'],
168 |         ])
169 | 
170 |         # Velocity to world coordinate frame
171 |         x_0[7:10] = v_dot_q(x_0[7:10], x_0[3:7])
172 |         x_1[7:10] = v_dot_q(x_1[7:10], x_1[3:7])
173 | 
174 |         # Solve single rotor thrusts
175 |         Jm1 = 1 / quad.J
176 | 
177 |         a1 = Jm1[0] * quad.y_f
178 |         b1 = data_0['ang acc x'] - Jm1[0] * (quad.J[1] - quad.J[2]) * x_0[11] * x_0[12]
179 |         a2 = -Jm1[1] * quad.x_f
180 |         b2 = data_0['ang acc y'] - Jm1[1] * (quad.J[2] - quad.J[0]) * x_0[12] * x_0[10]
181 |         a3 = Jm1[2] * quad.z_l_tau
182 |         b3 = data_0['ang acc z'] - Jm1[2] * (quad.J[0] - quad.J[1]) * x_0[10] * x_0[11]
183 |         a4 = np.ones(4,)
184 |         b4 = data_c['cmd_thrust'] * quad_mpc.quad.mass
185 | 
186 |         u = np.linalg.solve([a1, a2, a3, a4], [b1, b2, b3, b4]) / quad.max_thrust
187 | 
188 |         # ypr_0 = quaternion_to_euler(x_0[np.newaxis, 3:7])[0][::-1]
189 |         # 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]
190 |         # x_0[11] = np.cos(ypr_0[0]) * x_0[11] - np.sin(ypr_0[0]) * x_0[12]
191 |         # 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]
192 |         #
193 |         # ypr_1 = quaternion_to_euler(x_1[np.newaxis, 3:7])[0][::-1]
194 |         # 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]) * \
195 |         #           x_1[12]
196 |         # x_1[11] = np.cos(ypr_1[0]) * x_1[11] - np.sin(ypr_1[0]) * x_1[12]
197 |         # 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]
198 | 
199 |         # ypr_0 = quaternion_to_euler(x_0[np.newaxis, 3:7])[0][::-1]
200 |         # x_0[10] = 1.0 * x_0[10] - np.sin(ypr_0[1]) * x_0[12]
201 |         # x_0[11] = np.cos(ypr_0[0]) * x_0[11] + np.sin(ypr_0[0]) * np.cos(ypr_0[1]) * x_0[12]
202 |         # x_0[12] = -np.sin(ypr_0[0]) * x_0[11] + np.cos(ypr_0[0]) * np.cos(ypr_0[1]) * x_0[12]
203 |         #
204 |         # ypr_1 = quaternion_to_euler(x_1[np.newaxis, 3:7])[0][::-1]
205 |         # x_1[10] = 1.0 * x_1[10] - np.sin(ypr_1[1]) * x_1[12]
206 |         # x_1[11] = np.cos(ypr_1[0]) * x_1[11] + np.sin(ypr_1[0]) * np.cos(ypr_1[1]) * x_1[12]
207 |         # x_1[12] = -np.sin(ypr_1[0]) * x_1[11] + np.cos(ypr_1[0]) * np.cos(ypr_1[1]) * x_1[12]
208 | 
209 |         # u = np.hstack([
210 |         #     data_0['mot 2'],
211 |         #     data_0['mot 3'],
212 |         #     data_0['mot 1'],
213 |         #     data_0['mot 4'],
214 |         # ])
215 |         #
216 |         # u = u ** 2 * quad.thrust_map[0] / quad.max_thrust
217 | 
218 |         yield x_0, x_1, u, dt
219 | 
220 | def f_rate(x, u, quad):
221 |     """
222 |     Time-derivative of the angular rate
223 |     :param x: 4-length array of input state with components: 3D pos, quaternion angle, 3D vel, 3D rate
224 |     :param u: control input vector (4-dimensional): [trust_motor_1, ..., thrust_motor_4]
225 |     :param t_d: disturbance torque (3D)
226 |     :return: angular rate differential increment (scalar): dr/dt
227 |     """
228 | 
229 |     rate = x[10:]
230 |     return np.array([
231 |         1 / quad.J[0] * (u.dot(quad.y_f) + (quad.J[1] - quad.J[2]) * rate[1] * rate[2]),
232 |         1 / quad.J[1] * (-u.dot(quad.x_f) + (quad.J[2] - quad.J[0]) * rate[2] * rate[0]),
233 |         1 / quad.J[2] * (u.dot(quad.z_l_tau) + (quad.J[0] - quad.J[1]) * rate[0] * rate[1])
234 |     ]).squeeze()
235 | 
236 | def resimulate(x_0, x_f, u, dt, quad_mpc, rec_dict):
237 |     x_pred, _ = quad_mpc.forward_prop(x_0, u, t_horizon=dt, use_gp=False)
238 |     x_pred = x_pred[-1, np.newaxis, :]
239 | 
240 |     rec_dict['state_in'] = np.append(rec_dict['state_in'], x_0[np.newaxis, :], axis=0)
241 |     rec_dict['input_in'] = np.append(rec_dict['input_in'], u[np.newaxis, :], axis=0)
242 |     rec_dict['state_out'] = np.append(rec_dict['state_out'], x_f[np.newaxis, :], axis=0)
243 |     rec_dict['state_ref'] = np.append(rec_dict['state_ref'], np.zeros_like(x_f[np.newaxis, :]), axis=0)
244 |     rec_dict['timestamp'] = np.append(rec_dict['timestamp'], np.zeros_like(dt))
245 |     rec_dict['dt'] = np.append(rec_dict['dt'], dt)
246 |     rec_dict['state_pred'] = np.append(rec_dict['state_pred'], x_pred, axis=0)
247 |     rec_dict['error'] = np.append(rec_dict['error'], x_f - x_pred, axis=0)
248 | 
249 | 
250 | if __name__ == '__main__':
251 |     parser = argparse.ArgumentParser()
252 | 
253 |     parser.add_argument("--quad", type=str, default="kingfisher",
254 |                         help="Name of the quad.")
255 | 
256 |     input_arguments = parser.parse_args()
257 | 
258 |     ds_name = Conf.ds_name
259 |     simulation_options = Conf.ds_metadata
260 | 
261 |     quad = custom_quad_param_loader(input_arguments.quad)
262 |     quad_mpc = Quad3DMPC(quad)
263 | 
264 |     main(quad_mpc)
265 |     val(quad_mpc)
266 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/model_fitting/rdrv_fitting.py:
--------------------------------------------------------------------------------
  1 | """ Implementation and fitting of the RDRv linear regression model on flight data.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | 
 15 | import os
 16 | import joblib
 17 | import argparse
 18 | import subprocess
 19 | import numpy as np
 20 | import matplotlib.pyplot as plt
 21 | from sklearn.linear_model import LinearRegression
 22 | from src.model_fitting.gp_common import GPDataset, read_dataset
 23 | from src.utils.utils import v_dot_q, get_model_dir_and_file, safe_mknode_recursive
 24 | from config.configuration_parameters import ModelFitConfig as Conf
 25 | 
 26 | 
 27 | def linear_rdrv_fitting(x, y, feat_idx):
 28 |     drag_coeffs = np.zeros((3, 3))
 29 |     for i in range(x.shape[1]):
 30 |         reg = LinearRegression(fit_intercept=False).fit(x[:, i, np.newaxis], y[:, i])
 31 |         drag_coeffs[feat_idx[i], feat_idx[i]] = reg.coef_
 32 | 
 33 |     return drag_coeffs
 34 | 
 35 | 
 36 | def load_rdrv(model_options):
 37 |     directory, file_name = get_model_dir_and_file(model_options)
 38 |     rdrv_d = joblib.load(os.path.join(directory, file_name + '.pkl'))
 39 |     return rdrv_d
 40 | 
 41 | 
 42 | def main(model_name, features, quad_sim_options, dataset_name,
 43 |          x_cap, hist_bins, hist_thresh, plot=False):
 44 | 
 45 |     df_train = read_dataset(dataset_name, True, quad_sim_options)
 46 |     gp_dataset = GPDataset(df_train, features, [], features,
 47 |                            cap=x_cap, n_bins=hist_bins, thresh=hist_thresh, visualize_data=False)
 48 | 
 49 |     # Get X,Y datasets for the specified regression dimensions (subset of [7, 8, 9])
 50 |     a_err_b = gp_dataset.y
 51 |     v_b = gp_dataset.x
 52 | 
 53 |     drag_coeffs = linear_rdrv_fitting(v_b, a_err_b, np.array(features) - 7)
 54 | 
 55 |     # Get git commit hash for saving the model
 56 |     git_version = ''
 57 |     try:
 58 |         git_version = subprocess.check_output(['git', 'describe', '--always']).strip().decode("utf-8")
 59 |     except subprocess.CalledProcessError as e:
 60 |         print(e.returncode, e.output)
 61 |     gp_name_dict = {"git": git_version, "model_name": model_name, "params": quad_sim_options}
 62 |     save_file_path, save_file_name = get_model_dir_and_file(gp_name_dict)
 63 |     safe_mknode_recursive(save_file_path, save_file_name + '.pkl', overwrite=True)
 64 |     file_name = os.path.join(save_file_path, save_file_name + '.pkl')
 65 |     joblib.dump(drag_coeffs, file_name)
 66 | 
 67 |     if not plot:
 68 |         return drag_coeffs
 69 | 
 70 |     # Get X,Y datasets for velocity dimensions [7, 8, 9]
 71 |     a_err_b = gp_dataset.get_y(pruned=True, raw=True)[:, 7:10]
 72 |     v_b = gp_dataset.get_x(pruned=True, raw=True)[:, 7:10]
 73 | 
 74 |     # Compute predictions with RDRv model
 75 |     preds = []
 76 |     for i in range(len(a_err_b)):
 77 |         preds.append(np.matmul(drag_coeffs, v_b[i]))
 78 |     preds = np.stack(preds)
 79 | 
 80 |     ax_names = [r'$v_B^x$', r'$v_B^y$', r'$v_B^z$']
 81 |     y_dim_name = [r'drag $a_B^x$', 'drag $a_B^y$', 'drag $a_B^z$']
 82 |     fig, ax = plt.subplots(len(features), 1, sharex='all')
 83 |     for i in range(len(features)):
 84 |         ax[i].scatter(v_b[:, i], a_err_b[:, i], label='data')
 85 |         ax[i].scatter(v_b[:, i], preds[:, i], label='RDRv')
 86 |         ax[i].legend()
 87 |         ax[i].set_ylabel(y_dim_name[features[i] - 7])
 88 |         ax[i].set_xlabel(ax_names[features[i] - 7])
 89 |     ax[0].set_title('Body reference frame')
 90 | 
 91 |     # Remap error to world coordinates and predictions too
 92 |     q = gp_dataset.get_x(pruned=True, raw=True)[:, 3:7]
 93 |     for i in range(len(a_err_b)):
 94 |         a_err_b[i] = v_dot_q(a_err_b[i], q[i])
 95 |         preds[i] = v_dot_q(preds[i], q[i])
 96 | 
 97 |     ax_names = [r'$v_W^x$', r'$v_W^y$', r'$v_W^z$']
 98 |     y_dim_name = [r'drag $a_W^x$', 'drag $a_W^y$', 'drag $a_W^z$']
 99 |     fig, ax = plt.subplots(len(features), 1, sharex='all')
100 |     for i in range(len(features)):
101 |         ax[i].scatter(v_b[:, i], a_err_b[:, i], label='data')
102 |         ax[i].scatter(v_b[:, i], preds[:, i], label='RDRv')
103 |         ax[i].legend()
104 |         ax[i].set_ylabel(y_dim_name[features[i] - 7])
105 |         ax[i].set_xlabel(ax_names[features[i] - 7])
106 |     ax[0].set_title('World reference frame')
107 |     plt.show()
108 | 
109 |     return drag_coeffs
110 | 
111 | 
112 | if __name__ == "__main__":
113 |     parser = argparse.ArgumentParser()
114 | 
115 |     parser.add_argument("--model_name", type=str, default="",
116 |                         help="Name assigned to the trained model.")
117 | 
118 |     parser.add_argument('--x', nargs='+', type=int, default=[7, 8, 9],
119 |                         help='Regression X and Y variables. Must be a list of integers between 0 and 12.'
120 |                              'Velocities xyz correspond to indices 7, 8, 9.')
121 | 
122 |     input_arguments = parser.parse_args()
123 | 
124 |     reg_dimensions = input_arguments.x
125 |     rdrv_name = input_arguments.model_name
126 | 
127 |     histogram_bins = Conf.histogram_bins
128 |     histogram_threshold = Conf.histogram_threshold
129 |     histogram_cap = Conf.velocity_cap
130 | 
131 |     ds_name = Conf.ds_name
132 |     simulation_options = Conf.ds_metadata
133 | 
134 |     main(model_name=rdrv_name, features=reg_dimensions, quad_sim_options=simulation_options, dataset_name=ds_name,
135 |          x_cap=histogram_cap, hist_bins=histogram_bins, hist_thresh=histogram_threshold,
136 |          plot=Conf.visualize_training_result)
137 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/model_fitting/system_identification.py:
--------------------------------------------------------------------------------
  1 | import argparse
  2 | import os
  3 | 
  4 | import numpy as np
  5 | import pandas as pd
  6 | import rosbag
  7 | import rospy
  8 | 
  9 | from config.configuration_parameters import ModelFitConfig as Conf
 10 | from src.experiments.point_tracking_and_record import make_record_dict
 11 | from src.quad_mpc.create_ros_dd_mpc import custom_quad_param_loader
 12 | from src.quad_mpc.quad_3d_mpc import Quad3DMPC
 13 | from src.utils.utils import jsonify, get_data_dir_and_file
 14 | 
 15 | 
 16 | def odometry_parse(odom_msg):
 17 |     p = [odom_msg.pose.pose.position.x, odom_msg.pose.pose.position.y, odom_msg.pose.pose.position.z]
 18 |     q = [odom_msg.pose.pose.orientation.w, odom_msg.pose.pose.orientation.x, odom_msg.pose.pose.orientation.y,
 19 |          odom_msg.pose.pose.orientation.z]
 20 |     v = [odom_msg.twist.twist.linear.x, odom_msg.twist.twist.linear.y, odom_msg.twist.twist.linear.z]
 21 |     w = [odom_msg.twist.twist.angular.x, odom_msg.twist.twist.angular.y, odom_msg.twist.twist.angular.z]
 22 | 
 23 |     return np.array(p + q + v + w)
 24 | 
 25 | 
 26 | def thrust_motor_model(motor_tau, thrust, thrust_des, dt):
 27 |     if motor_tau < 1e-12:
 28 |         return thrust_des
 29 |     tau_inv = 1 / motor_tau
 30 |     thrust_out = (
 31 |             tau_inv ** 2 * (thrust_des - 2 * (thrust * thrust_des) ** 0.5 + thrust) * dt ** 2
 32 |             + 2 * tau_inv * ((thrust * thrust_des) ** 0.5 - thrust) * dt
 33 |             + thrust
 34 |     )
 35 |     return thrust_out
 36 | 
 37 | 
 38 | def system_identification(quad_mpc, ds_name, sim_options):
 39 |     """
 40 |     Processes a rosbag with recorded odometry and send commands to do a system identification of the residual error.
 41 |     :param quad_mpc: Quad model used for forward simulation
 42 |     :param ds_name: Dataset name
 43 |     :param sim_options: Simulation options.
 44 |     :return:
 45 |     """
 46 |     # Open bag
 47 |     data_file = get_data_dir_and_file(ds_name, training_split=True, params=sim_options, read_only=True)
 48 |     if data_file is None:
 49 |         raise FileNotFoundError
 50 |     rec_file_dir, rec_file_name = data_file
 51 |     rec_file_name_bag = rec_file_name.replace('.csv', '.bag')
 52 |     rec_file = os.path.join(rec_file_dir, rec_file_name_bag)
 53 |     save_file = os.path.join(rec_file_dir, rec_file_name)
 54 |     bag = rosbag.Bag(rec_file)
 55 | 
 56 |     # Re-simulate thrust with motor model and communication delay
 57 |     thrust = None
 58 |     control = {'t': [], 'thrust': []}
 59 |     for topic, msg, t in bag.read_messages(topics=['control']):
 60 |         t = msg.header.stamp
 61 |         desired_thrust = np.array(msg.thrusts)
 62 |         if thrust is None:
 63 |             thrust = desired_thrust
 64 |         t = t + rospy.Duration.from_sec(quad_mpc.quad.comm_delay) + rospy.Duration.from_sec(0.001)
 65 |         for dt in np.arange(0.001, 0.021, step=0.001):
 66 |             thrust = thrust_motor_model(quad_mpc.quad.motor_tau, thrust, desired_thrust, 0.001)
 67 |             t_at = t + rospy.Duration.from_sec(dt)
 68 |             control['t'].append(t_at.to_sec())
 69 |             control['thrust'].append(thrust)
 70 | 
 71 |     control = pd.DataFrame(control)
 72 |     control = control.set_index('t')
 73 |     control = control[~control.index.duplicated(keep='last')]
 74 |     control = control.sort_index()
 75 | 
 76 |     state = {'t': [], 'state': []}
 77 |     states_list = []
 78 |     recording = False
 79 |     for topic, msg, t in bag.read_messages(topics=['recording_ctrl', 'state']):
 80 |         if topic == 'recording_ctrl':
 81 |             recording = msg.data
 82 |         if topic == 'state' and recording:
 83 |             t = msg.header.stamp
 84 |             state['t'].append(t.to_sec())
 85 |             x = odometry_parse(msg)
 86 |             state['state'].append(x)
 87 |             states_list.append(np.array(x))
 88 | 
 89 |     states_np = np.array(states_list)
 90 | 
 91 |     filtered_states_np = states_np.copy()
 92 |     from scipy.signal import savgol_filter
 93 |     window_size_xy = 31
 94 |     window_size_z = 31
 95 |     window_size_q = 31
 96 |     window_size_v = 31
 97 |     window_size_w = 121
 98 |     poly_order = 4
 99 |     filtered_states_np[:, :2] = savgol_filter(filtered_states_np[:, :2], window_size_xy, poly_order, axis=0)
100 |     filtered_states_np[:, 2] = savgol_filter(filtered_states_np[:, 2], window_size_z, poly_order, axis=0)
101 |     filtered_states_np[:, 3:7] = savgol_filter(filtered_states_np[:, 3:7], window_size_q, 2, axis=0)
102 |     filtered_states_np[:, 3:7] = filtered_states_np[:, 3:7] / np.linalg.norm(filtered_states_np[:, 3:7], axis=1, keepdims=True)
103 |     filtered_states_np[:, 7:10] = savgol_filter(filtered_states_np[:, 7:10], window_size_v, 2, axis=0)
104 |     filtered_states_np[:, 10:] = savgol_filter(filtered_states_np[:, 10:], window_size_w, poly_order, axis=0)
105 | 
106 |     # def extract_windows(array, sub_window_size):
107 |     #     examples = []
108 |     #
109 |     #     for i in range(0, array.shape[0]-sub_window_size, sub_window_size):
110 |     #         example = array[i:sub_window_size + i]
111 |     #         examples.append(np.expand_dims(example, 0))
112 |     #
113 |     #     return np.vstack(examples)
114 |     #
115 |     # from scipy.spatial.transform import Rotation as R
116 |     # from scipy.spatial.transform import RotationSpline
117 |     # window_rot = extract_windows(filtered_states_np[:, 3:7], 20)
118 |     # mean_rot = []
119 |     # for i in range(window_rot.shape[0]):
120 |     #     mean_rot.append(R.from_quat(window_rot[i]).mean().as_quat())
121 |     # mean_rot = np.vstack(mean_rot)
122 |     # rotations = R.from_quat(mean_rot)
123 |     # spline = RotationSpline(state['t'][9::20][:mean_rot.shape[0]], rotations)
124 |     # filtered_states_np[:, 3:7] = spline(state['t']).as_quat()
125 | 
126 |     import matplotlib.pyplot as plt
127 |     # plt.plot(states_np[:, 9])
128 |     #plt.plot(filtered_states_np[3975:4025, 0])
129 |     #plt.plot(states_np[3975:4025, 0])
130 |     #plt.plot(np.gradient(filtered_states_np[3975:4025, 1])/1e-2)
131 |     #plt.plot(filtered_states_np[3975:4025, 7])
132 |     #plt.plot(filtered_states_np[1800:2000, 3])
133 |     #plt.plot(states_np[1800:2000, 3])
134 |     #plt.plot(filtered_states_np[3775:4025, 3:7] * np.sign(filtered_states_np[3775:4025, 3]))
135 |     #plt.plot(filtered_states_np[3775:4025, 10])
136 |     #plt.plot(states_np[3775:4025, 10])
137 |     #plt.show()
138 | 
139 |     state = {'t': state['t'], 'state': []}
140 |     for i in range(filtered_states_np.shape[0]):
141 |         state['state'].append(filtered_states_np[i])
142 | 
143 |     rec_dict = make_record_dict(state_dim=13)
144 |     for i in range(1, len(state['t'])):
145 |         last_state_idx = control.index.get_loc(state['t'][i-1], method='nearest')
146 |         curr_state_idx = control.index.get_loc(state['t'][i], method='nearest')
147 |         u0 = control['thrust'].iloc[last_state_idx] / quad_mpc.quad.max_thrust
148 |         u1 = control['thrust'].iloc[curr_state_idx] / quad_mpc.quad.max_thrust
149 | 
150 |         # Only use subsequent states with similar thrust state for system identification (within 1%)
151 |         if np.all(np.abs(u0 - u1) < 0.01):
152 |             x_0 = state['state'][i-1]
153 |             x_f = state['state'][i]
154 |             u = np.vstack([u0, u1]).mean(0)
155 |             dt = state['t'][i] - state['t'][i-1]
156 |             if dt > 0.015:
157 |                 continue
158 |             x_pred, _ = quad_mpc.forward_prop(x_0, u, t_horizon=dt, use_gp=False)
159 |             x_pred = x_pred[-1, np.newaxis, :]
160 | 
161 |             rec_dict['state_in'] = np.append(rec_dict['state_in'], x_0[np.newaxis, :], axis=0)
162 |             rec_dict['input_in'] = np.append(rec_dict['input_in'], u[np.newaxis, :], axis=0)
163 |             rec_dict['state_out'] = np.append(rec_dict['state_out'], x_f[np.newaxis, :], axis=0)
164 |             rec_dict['state_ref'] = np.append(rec_dict['state_ref'], np.zeros_like(x_f[np.newaxis, :]), axis=0)
165 |             rec_dict['timestamp'] = np.append(rec_dict['timestamp'], state['t'][i])
166 |             rec_dict['dt'] = np.append(rec_dict['dt'], dt)
167 |             rec_dict['state_pred'] = np.append(rec_dict['state_pred'], x_pred, axis=0)
168 |             rec_dict['error'] = np.append(rec_dict['error'], x_f - x_pred, axis=0)
169 | 
170 |     # Save datasets
171 |     for key in rec_dict.keys():
172 |         print(key, " ", rec_dict[key].shape)
173 |         rec_dict[key] = jsonify(rec_dict[key])
174 |     df = pd.DataFrame(rec_dict)
175 |     df.to_csv(save_file, index=True, header=True)
176 |     return
177 |     import matplotlib.pyplot as plt
178 |     #plt.plot(states_np[:, 9])
179 |     plt.plot(np.gradient(out[:, 8]))
180 |     plt.show()
181 |     return
182 | 
183 |     # Match each state to the corresponding thrust state of the quad
184 |     rec_dict = make_record_dict(state_dim=13)
185 |     recording = False
186 |     last_state_msg = None
187 |     for topic, msg, t in bag.read_messages(topics=['recording_ctrl', 'state']):
188 |         if topic == 'recording_ctrl':
189 |             recording = msg.data
190 |             if not recording:
191 |                 last_state_msg = None
192 |         if topic == 'state' and recording:
193 |             if last_state_msg is not None:
194 |                 last_state_idx = control.index.get_loc(last_state_msg.header.stamp.to_sec(), method='nearest')
195 |                 curr_state_idx = control.index.get_loc(msg.header.stamp.to_sec(), method='nearest')
196 |                 u0 = control['thrust'].iloc[last_state_idx] / quad_mpc.quad.max_thrust
197 |                 u1 = control['thrust'].iloc[curr_state_idx] / quad_mpc.quad.max_thrust
198 | 
199 |                 # Only use subsequent states with similar thrust state for system identification (within 1%)
200 |                 if np.all(np.abs(u0 - u1) < 0.01):
201 |                     x_0 = odometry_parse(last_state_msg)
202 |                     x_f = odometry_parse(msg)
203 |                     u = np.vstack([u0, u1]).mean(0)
204 |                     dt = msg.header.stamp.to_sec() - last_state_msg.header.stamp.to_sec()
205 |                     x_pred, _ = quad_mpc.forward_prop(x_0, u, t_horizon=dt, use_gp=False)
206 |                     x_pred = x_pred[-1, np.newaxis, :]
207 | 
208 |                     rec_dict['state_in'] = np.append(rec_dict['state_in'], x_0[np.newaxis, :], axis=0)
209 |                     rec_dict['input_in'] = np.append(rec_dict['input_in'], u[np.newaxis, :], axis=0)
210 |                     rec_dict['state_out'] = np.append(rec_dict['state_out'], x_f[np.newaxis, :], axis=0)
211 |                     rec_dict['state_ref'] = np.append(rec_dict['state_ref'], np.zeros_like(x_f[np.newaxis, :]), axis=0)
212 |                     rec_dict['timestamp'] = np.append(rec_dict['timestamp'], msg.header.stamp.to_sec())
213 |                     rec_dict['dt'] = np.append(rec_dict['dt'], dt)
214 |                     rec_dict['state_pred'] = np.append(rec_dict['state_pred'], x_pred, axis=0)
215 |                     rec_dict['error'] = np.append(rec_dict['error'], x_f - x_pred, axis=0)
216 | 
217 |             last_state_msg = msg
218 | 
219 |     # Save datasets
220 |     for key in rec_dict.keys():
221 |         print(key, " ", rec_dict[key].shape)
222 |         rec_dict[key] = jsonify(rec_dict[key])
223 |     df = pd.DataFrame(rec_dict)
224 |     df.to_csv(save_file, index=True, header=True)
225 | 
226 | 
227 | if __name__ == '__main__':
228 | 
229 |     parser = argparse.ArgumentParser()
230 | 
231 |     parser.add_argument("--quad", type=str, default="kingfisher",
232 |                         help="Name of the quad.")
233 | 
234 |     input_arguments = parser.parse_args()
235 | 
236 |     ds_name = Conf.ds_name
237 |     simulation_options = Conf.ds_metadata
238 | 
239 |     quad = custom_quad_param_loader(input_arguments.quad)
240 |     quad_mpc = Quad3DMPC(quad)
241 | 
242 |     system_identification(quad_mpc, ds_name, simulation_options)


--------------------------------------------------------------------------------
/ros_dd_mpc/src/quad_mpc/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/TUM-AAS/neural-mpc/5f15661c2d1a3b6072911d8993214c50a68e11eb/ros_dd_mpc/src/quad_mpc/__init__.py


--------------------------------------------------------------------------------
/ros_dd_mpc/src/quad_mpc/create_ros_dd_mpc.py:
--------------------------------------------------------------------------------
  1 | """ Class for interfacing the data-augmented MPC with ROS.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | from src.quad_mpc.quad_3d import Quadrotor3D
 15 | from src.quad_mpc.quad_3d_mpc import Quad3DMPC
 16 | import numpy as np
 17 | import os
 18 | import yaml
 19 | 
 20 | 
 21 | def custom_quad_param_loader(quad_name):
 22 | 
 23 |     this_path = os.path.dirname(os.path.realpath(__file__))
 24 |     params_file = os.path.join(this_path, '..', '..', 'config', quad_name + '.yaml')
 25 | 
 26 |     # Get parameters for drone
 27 |     with open(params_file, "r") as stream:
 28 |         attrib = yaml.safe_load(stream)
 29 | 
 30 |     quad = Quadrotor3D(noisy=False, drag=False, payload=False, motor_noise=False)
 31 |     quad.mass = float(attrib['mass']) + (float(attrib['mass_rotor']) if 'mass_rotor' in attrib else 0) * 4
 32 |     quad.J = np.array(attrib['inertia'])
 33 | 
 34 |     quad.thrust_map = attrib["thrust_map"]
 35 | 
 36 |     if 'thrust_max' in attrib:
 37 |         quad.max_thrust = attrib['thrust_max']
 38 |     else:
 39 |         float(attrib["motor_omega_max"]) ** 2 * quad.thrust_map[0]
 40 |     quad.c = float(attrib['kappa'])
 41 | 
 42 |     if 'arm_length' in attrib and attrib['rotors_config'] == 'cross':
 43 |         quad.length = float(attrib['arm_length'])
 44 |         h = np.cos(np.pi / 4) * quad.length
 45 |         quad.x_f = np.array([h, -h, -h, h])
 46 |         quad.y_f = np.array([-h, h, -h, h])
 47 |     elif 'arm_length' in attrib and attrib['rotors_config'] == 'plus':
 48 |         quad.length = float(attrib['arm_length'])
 49 |         quad.x_f = np.array([0, 0, -quad.length, quad.length])
 50 |         quad.y_f = np.array([-quad.length, quad.length, 0, 0])
 51 |     else:
 52 |         tbm_fr = np.array(attrib['tbm_fr'])
 53 |         tbm_bl = np.array(attrib['tbm_bl'])
 54 |         tbm_br = np.array(attrib['tbm_br'])
 55 |         tbm_fl = np.array(attrib['tbm_fl'])
 56 |         quad.length = np.linalg.norm(tbm_fr)
 57 |         quad.x_f = np.array([tbm_fr[0], tbm_bl[0], tbm_br[0], tbm_fl[0]])
 58 |         quad.y_f = np.array([tbm_fr[1], tbm_bl[1], tbm_br[1], tbm_fl[1]])
 59 |     quad.z_l_tau = np.array([-quad.c, -quad.c, quad.c, quad.c])
 60 | 
 61 |     if 'motor_tau' in attrib:
 62 |         quad.motor_tau = attrib['motor_tau']
 63 | 
 64 |     if 'comm_delay' in attrib:
 65 |         quad.comm_delay = attrib['comm_delay']
 66 | 
 67 |     return quad
 68 | 
 69 | 
 70 | class ROSDDMPC:
 71 |     def __init__(self, t_horizon, n_mpc_nodes, opt_dt, quad_name, point_reference=False, models=None,
 72 |                  model_conf=None, rdrv=None):
 73 | 
 74 |         quad = custom_quad_param_loader(quad_name)
 75 | 
 76 |         # Initialize quad MPC
 77 |         if point_reference:
 78 |             acados_config = {
 79 |                 "solver_type": "SQP",
 80 |                 "terminal_cost": True
 81 |             }
 82 |         else:
 83 |             acados_config = {
 84 |                 "solver_type": "SQP_RTI",
 85 |                 "terminal_cost": False
 86 |             }
 87 | 
 88 |         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])
 89 |         r_diagonal = np.array([1.0, 1.0, 1.0, 1.0])
 90 | 
 91 |         q_mask = np.array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]).T
 92 | 
 93 |         quad_mpc = Quad3DMPC(quad, t_horizon=t_horizon, optimization_dt=opt_dt, n_nodes=n_mpc_nodes,
 94 |                              pre_trained_models=models, model_conf=model_conf, model_name=quad_name,
 95 |                              solver_options=acados_config,
 96 |                              q_mask=q_mask, q_cost=q_diagonal, r_cost=r_diagonal, rdrv_d_mat=rdrv)
 97 | 
 98 |         self.quad_name = quad_name
 99 |         self.quad = quad
100 |         self.quad_mpc = quad_mpc
101 | 
102 |         # Last optimization
103 |         self.last_w = None
104 | 
105 |     def set_state(self, x):
106 |         """
107 |         Set quadrotor state estimate from an odometry measurement
108 |         :param x: measured state from odometry. List with 13 components with the format: [p_xyz, q_wxyz, v_xyz, w_xyz]
109 |         """
110 | 
111 |         self.quad.set_state(x)
112 | 
113 |     def set_gp_state(self, x):
114 |         """
115 |         Set a quadrotor state estimate from an odometry measurement. While the input state in the function set_state()
116 |         will be used as initial state for the MPC optimization, the input state of this function will be used to
117 |         evaluate the GP's. If this function is never called, then the GP's will be evaluated with the state from
118 |         set_state() too.
119 |         :param x: measured state from odometry. List with 13 components with the format: [p_xyz, q_wxyz, v_xyz, w_xyz]
120 |         """
121 | 
122 |         self.quad.set_gp_state(x)
123 | 
124 |     def set_reference(self, x_ref, u_ref):
125 |         """
126 |         Set a reference state for the optimizer.
127 |         :param x_ref: list with 4 sub-components (position, angle quaternion, velocity, body rate). If these four
128 |         are lists, then this means a single target point is used. If they are Nx3 and Nx4 (for quaternion) numpy arrays,
129 |         then they are interpreted as a sequence of N tracking points.
130 |         :param u_ref: Optional target for the optimized control inputs
131 |         """
132 |         return self.quad_mpc.set_reference(x_reference=x_ref, u_reference=u_ref)
133 | 
134 |     def optimize(self, model_data):
135 |         w_opt, x_opt = self.quad_mpc.optimize(use_model=model_data, return_x=True)
136 | 
137 |         return w_opt, x_opt
138 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/quad_mpc/quad_3d.py:
--------------------------------------------------------------------------------
  1 | """ Implementation of the Simplified Simulator and its quadrotor dynamics.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | 
 15 | from math import sqrt
 16 | import numpy as np
 17 | from src.utils.utils import quaternion_to_euler, skew_symmetric, v_dot_q, unit_quat, quaternion_inverse
 18 | 
 19 | 
 20 | class Quadrotor3D:
 21 | 
 22 |     def __init__(self, noisy=False, drag=False, payload=False, motor_noise=False):
 23 |         """
 24 |         Initialization of the 3D quadrotor class
 25 |         :param noisy: Whether noise is used in the simulation
 26 |         :type noisy: bool
 27 |         :param drag: Whether to simulate drag or not.
 28 |         :type drag: bool
 29 |         :param payload: Whether to simulate a payload force in the simulation
 30 |         :type payload: bool
 31 |         :param motor_noise: Whether non-gaussian noise is considered in the motor inputs
 32 |         :type motor_noise: bool
 33 |         """
 34 | 
 35 |         # Either 'x' or '+'. The xy configuration of the thrusters in the body plane.
 36 |         configuration = 'x'
 37 | 
 38 |         # Maximum thrust in Newtons of a thruster when rotating at maximum speed.
 39 |         self.max_thrust = 20
 40 |         self.thrust_map = 1e-6
 41 | 
 42 |         # Motor time constant
 43 |         self.motor_tau = 0.
 44 | 
 45 |         # Communication delay
 46 |         self.comm_delay = 0.
 47 | 
 48 |         # System state space
 49 |         self.pos = np.zeros((3,))
 50 |         self.vel = np.zeros((3,))
 51 |         self.angle = np.array([1., 0., 0., 0.])  # Quaternion format: qw, qx, qy, qz
 52 |         self.a_rate = np.zeros((3,))
 53 | 
 54 |         # System state space for GP evaluation.
 55 |         self.gp_pos = None
 56 |         self.gp_vel = None
 57 |         self.gp_angle = None
 58 |         self.gp_a_rate = None
 59 | 
 60 |         # Input constraints
 61 |         self.max_input_value = 1  # Motors at full thrust
 62 |         self.min_input_value = 0  # Motors turned off
 63 | 
 64 |         # Quadrotor intrinsic parameters
 65 |         self.J = np.array([.03, .03, .06])  # N m s^2 = kg m^2
 66 |         self.mass = 1.0  # kg
 67 | 
 68 |         # Length of motor to CoG segment
 69 |         self.length = 0.47 / 2  # m
 70 | 
 71 |         # Positions of thrusters
 72 |         if configuration == '+':
 73 |             self.x_f = np.array([self.length, 0, -self.length, 0])
 74 |             self.y_f = np.array([0, self.length, 0, -self.length])
 75 |         elif configuration == 'x':
 76 |             h = np.cos(np.pi / 4) * self.length
 77 |             self.x_f = np.array([h, -h, -h, h])
 78 |             self.y_f = np.array([-h, -h, h, h])
 79 | 
 80 |         # For z thrust torque calculation
 81 |         self.c = 0.013  # m   (z torque generated by each motor)
 82 |         self.z_l_tau = np.array([-self.c, self.c, -self.c, self.c])
 83 | 
 84 |         # Gravity vector
 85 |         self.g = np.array([[0], [0], [9.8066]])  # m s^-2
 86 | 
 87 |         # Actuation thrusts
 88 |         self.u_noiseless = np.array([0.0, 0.0, 0.0, 0.0])
 89 |         self.u = np.array([0.0, 0.0, 0.0, 0.0])  # N
 90 | 
 91 |         # Drag coefficients [kg / m]
 92 |         self.rotor_drag_xy = 0.3
 93 |         self.rotor_drag_z = 0.0  # No rotor drag in the z dimension
 94 |         self.rotor_drag = np.array([self.rotor_drag_xy, self.rotor_drag_xy, self.rotor_drag_z])[:, np.newaxis]
 95 |         self.aero_drag = 0.08
 96 | 
 97 |         self.drag = drag
 98 |         self.noisy = noisy
 99 |         self.motor_noise = motor_noise
100 | 
101 |         self.payload_mass = 0.3  # kg
102 |         self.payload_mass = self.payload_mass * payload
103 | 
104 |     def set_state(self, *args, **kwargs):
105 |         if len(args) != 0:
106 |             assert len(args) == 1 and len(args[0]) == 13
107 |             self.pos[0], self.pos[1], self.pos[2], \
108 |             self.angle[0], self.angle[1], self.angle[2], self.angle[3], \
109 |             self.vel[0], self.vel[1], self.vel[2], \
110 |             self.a_rate[0], self.a_rate[1], self.a_rate[2] \
111 |                 = args[0]
112 | 
113 |         else:
114 |             self.pos = kwargs["pos"]
115 |             self.angle = kwargs["angle"]
116 |             self.vel = kwargs["vel"]
117 |             self.a_rate = kwargs["rate"]
118 | 
119 |     def set_gp_state(self, *args, **kwargs):
120 |         if self.gp_pos is None:
121 |             # Instantiate these variables for the first time
122 |             self.gp_pos = np.zeros((3,))
123 |             self.gp_vel = np.zeros((3,))
124 |             self.gp_angle = np.array([1., 0., 0., 0.])  # Quaternion format: qw, qx, qy, qz
125 |             self.gp_a_rate = np.zeros((3,))
126 | 
127 |         if len(args) != 0:
128 |             assert len(args) == 1 and len(args[0]) == 13
129 |             self.gp_pos[0], self.gp_pos[1], self.gp_pos[2], \
130 |             self.gp_angle[0], self.gp_angle[1], self.gp_angle[2], self.gp_angle[3], \
131 |             self.gp_vel[0], self.gp_vel[1], self.gp_vel[2], \
132 |             self.gp_a_rate[0], self.gp_a_rate[1], self.gp_a_rate[2] \
133 |                 = args[0]
134 | 
135 |         else:
136 |             self.gp_pos = kwargs["pos"]
137 |             self.gp_angle = kwargs["angle"]
138 |             self.gp_vel = kwargs["vel"]
139 |             self.gp_a_rate = kwargs["rate"]
140 | 
141 |     def get_state(self, quaternion=False, stacked=False):
142 | 
143 |         if quaternion and not stacked:
144 |             return [self.pos, self.angle, self.vel, self.a_rate]
145 |         if quaternion and stacked:
146 |             return [self.pos[0], self.pos[1], self.pos[2], self.angle[0], self.angle[1], self.angle[2], self.angle[3],
147 |                     self.vel[0], self.vel[1], self.vel[2], self.a_rate[0], self.a_rate[1], self.a_rate[2]]
148 | 
149 |         angle = quaternion_to_euler(self.angle)
150 |         if not quaternion and stacked:
151 |             return [self.pos[0], self.pos[1], self.pos[2], angle[0], angle[1], angle[2],
152 |                     self.vel[0], self.vel[1], self.vel[2], self.a_rate[0], self.a_rate[1], self.a_rate[2]]
153 |         return [self.pos, angle, self.vel, self.a_rate]
154 | 
155 |     def get_gp_state(self, quaternion=False, stacked=False):
156 | 
157 |         if self.gp_pos is None:
158 |             return None
159 | 
160 |         if quaternion and not stacked:
161 |             return [self.gp_pos, self.gp_angle, self.gp_vel, self.gp_a_rate]
162 |         if quaternion and stacked:
163 |             return [self.gp_pos[0], self.gp_pos[1], self.gp_pos[2],
164 |                     self.gp_angle[0], self.gp_angle[1], self.gp_angle[2], self.gp_angle[3],
165 |                     self.gp_vel[0], self.gp_vel[1], self.gp_vel[2],
166 |                     self.gp_a_rate[0], self.gp_a_rate[1], self.gp_a_rate[2]]
167 | 
168 |         gp_angle = quaternion_to_euler(self.gp_angle)
169 |         if not quaternion and stacked:
170 |             return [self.gp_pos[0], self.gp_pos[1], self.gp_pos[2],
171 |                     gp_angle[0], gp_angle[1], gp_angle[2],
172 |                     self.gp_vel[0], self.gp_vel[1], self.gp_vel[2],
173 |                     self.gp_a_rate[0], self.gp_a_rate[1], self.gp_a_rate[2]]
174 |         return [self.gp_pos, gp_angle, self.gp_vel, self.gp_a_rate]
175 | 
176 |     def get_control(self, noisy=False):
177 |         if not noisy:
178 |             return self.u_noiseless
179 |         else:
180 |             return self.u
181 | 
182 |     def update(self, u, dt):
183 |         """
184 |         Runge-Kutta 4th order dynamics integration
185 | 
186 |         :param u: 4-dimensional vector with components between [0.0, 1.0] that represent the activation of each motor.
187 |         :param dt: time differential
188 |         """
189 | 
190 |         # Clip inputs
191 |         for i, u_i in enumerate(u):
192 |             self.u_noiseless[i] = max(min(u_i, self.max_input_value), self.min_input_value)
193 | 
194 |         # Apply noise to inputs (uniformly distributed noise with standard deviation proportional to input magnitude)
195 |         if self.motor_noise:
196 |             for i, u_i in enumerate(self.u_noiseless):
197 |                 std = 0.02 * sqrt(u_i)
198 |                 noise_u = np.random.normal(loc=0.1 * (u_i / 1.3) ** 2, scale=std)
199 |                 self.u[i] = max(min(u_i - noise_u, self.max_input_value), self.min_input_value) * self.max_thrust
200 |         else:
201 |             self.u = self.u_noiseless * self.max_thrust
202 | 
203 |         # Generate disturbance forces / torques
204 |         if self.noisy:
205 |             f_d = np.random.normal(size=(3, 1), scale=10 * dt)
206 |             t_d = np.random.normal(size=(3, 1), scale=10 * dt)
207 |         else:
208 |             f_d = np.zeros((3, 1))
209 |             t_d = np.zeros((3, 1))
210 | 
211 |         x = self.get_state(quaternion=True, stacked=False)
212 | 
213 |         # RK4 integration
214 |         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)]
215 |         x_aux = [x[i] + dt / 2 * k1[i] for i in range(4)]
216 |         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)]
217 |         x_aux = [x[i] + dt / 2 * k2[i] for i in range(4)]
218 |         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)]
219 |         x_aux = [x[i] + dt * k3[i] for i in range(4)]
220 |         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)]
221 |         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
222 |              range(4)]
223 | 
224 |         # Ensure unit quaternion
225 |         x[1] = unit_quat(x[1])
226 | 
227 |         self.pos, self.angle, self.vel, self.a_rate = x
228 | 
229 |     def f_pos(self, x):
230 |         """
231 |         Time-derivative of the position vector
232 |         :param x: 4-length array of input state with components: 3D pos, quaternion angle, 3D vel, 3D rate
233 |         :return: position differential increment (vector): d[pos_x; pos_y]/dt
234 |         """
235 | 
236 |         vel = x[2]
237 |         return vel
238 | 
239 |     def f_att(self, x):
240 |         """
241 |         Time-derivative of the attitude in quaternion form
242 |         :param x: 4-length array of input state with components: 3D pos, quaternion angle, 3D vel, 3D rate
243 |         :return: attitude differential increment (quaternion qw, qx, qy, qz): da/dt
244 |         """
245 | 
246 |         rate = x[3]
247 |         angle_quaternion = x[1]
248 | 
249 |         return 1 / 2 * skew_symmetric(rate).dot(angle_quaternion)
250 | 
251 |     def f_vel(self, x, u, f_d):
252 |         """
253 |         Time-derivative of the velocity vector
254 |         :param x: 4-length array of input state with components: 3D pos, quaternion angle, 3D vel, 3D rate
255 |         :param u: control input vector (4-dimensional): [trust_motor_1, ..., thrust_motor_4]
256 |         :param f_d: disturbance force vector (3-dimensional)
257 |         :return: 3D velocity differential increment (vector): d[vel_x; vel_y; vel_z]/dt
258 |         """
259 | 
260 |         a_thrust = np.array([[0], [0], [np.sum(u)]]) / self.mass
261 | 
262 |         if self.drag:
263 |             # Transform velocity to body frame
264 |             v_b = v_dot_q(x[2], quaternion_inverse(x[1]))[:, np.newaxis]
265 |             # Compute aerodynamic drag acceleration in world frame
266 |             a_drag = -self.aero_drag * v_b ** 2 * np.sign(v_b) / self.mass
267 |             # Add rotor drag
268 |             a_drag -= self.rotor_drag * v_b / self.mass
269 |             # Transform drag acceleration to world frame
270 |             a_drag = v_dot_q(a_drag, x[1])
271 |         else:
272 |             a_drag = np.zeros((3, 1))
273 | 
274 |         angle_quaternion = x[1]
275 | 
276 |         a_payload = -self.payload_mass * self.g / self.mass
277 | 
278 |         return np.squeeze(-self.g + a_payload + a_drag + v_dot_q(a_thrust + f_d / self.mass, angle_quaternion))
279 | 
280 |     def f_rate(self, x, u, t_d):
281 |         """
282 |         Time-derivative of the angular rate
283 |         :param x: 4-length array of input state with components: 3D pos, quaternion angle, 3D vel, 3D rate
284 |         :param u: control input vector (4-dimensional): [trust_motor_1, ..., thrust_motor_4]
285 |         :param t_d: disturbance torque (3D)
286 |         :return: angular rate differential increment (scalar): dr/dt
287 |         """
288 | 
289 |         rate = x[3]
290 |         return np.array([
291 |             1 / self.J[0] * (u.dot(self.y_f) + t_d[0] + (self.J[1] - self.J[2]) * rate[1] * rate[2]),
292 |             1 / self.J[1] * (-u.dot(self.x_f) + t_d[1] + (self.J[2] - self.J[0]) * rate[2] * rate[0]),
293 |             1 / self.J[2] * (u.dot(self.z_l_tau) + t_d[2] + (self.J[0] - self.J[1]) * rate[0] * rate[1])
294 |         ]).squeeze()
295 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/quad_mpc/quad_3d_mpc.py:
--------------------------------------------------------------------------------
  1 | """ Implementation of the data-augmented MPC for quadrotors.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | 
 15 | import numpy as np
 16 | import torch
 17 | 
 18 | from src.quad_mpc.quad_3d_optimizer import Quad3DOptimizer
 19 | from src.model_fitting.gp_common import restore_gp_regressors
 20 | from src.utils.quad_3d_opt_utils import simulate_plant, uncertainty_forward_propagation
 21 | from src.utils.utils import make_bx_matrix
 22 | 
 23 | 
 24 | class Quad3DMPC:
 25 |     def __init__(self, my_quad, t_horizon=1.0, n_nodes=5, q_cost=None, r_cost=None,
 26 |                  optimization_dt=5e-2, simulation_dt=5e-4, pre_trained_models=None, model_name="my_quad", q_mask=None,
 27 |                  solver_options=None, rdrv_d_mat=None, model_conf=None):
 28 |         """
 29 |         :param my_quad: Quadrotor3D simulator object
 30 |         :type my_quad: Quadrotor3D
 31 |         :param t_horizon: time horizon for optimization loop MPC controller
 32 |         :param n_nodes: number of MPC nodes
 33 |         :param optimization_dt: time step between two successive optimizations intended to be used.
 34 |         :param simulation_dt: discretized time-step for the quadrotor simulation
 35 |         :param pre_trained_models: additional pre-trained GP regressors to be combined with nominal model in the MPC
 36 |         :param q_cost: diagonal of Q matrix for LQR cost of MPC cost function. Must be a numpy array of length 13.
 37 |         :param r_cost: diagonal of R matrix for LQR cost of MPC cost function. Must be a numpy array of length 4.
 38 |         :param q_mask: Optional boolean mask that determines which variables from the state compute towards the
 39 |         cost function. In case no argument is passed, all variables are weighted.
 40 |         :param solver_options: Optional set of extra options dictionary for acados solver.
 41 |         :param rdrv_d_mat: 3x3 matrix that corrects the drag with a linear model according to Faessler et al. 2018. None
 42 |         if not used
 43 |         """
 44 | 
 45 |         if rdrv_d_mat is not None:
 46 |             # rdrv is currently not compatible with covariance mode or with GP-MPC.
 47 |             print("RDRv mode")
 48 |             self.rdrv = rdrv_d_mat
 49 |             assert pre_trained_models is None
 50 |         else:
 51 |             self.rdrv = None
 52 | 
 53 |         self.quad = my_quad
 54 |         self.simulation_dt = simulation_dt
 55 |         self.optimization_dt = optimization_dt
 56 | 
 57 |         # motor commands from last step
 58 |         self.motor_u = np.array([0., 0., 0., 0.])
 59 | 
 60 |         self.n_nodes = n_nodes
 61 |         self.t_horizon = t_horizon
 62 | 
 63 |         self.mlp = None
 64 |         self.mlp_approx = None
 65 | 
 66 |         # Load augmented dynamics model with GP regressor
 67 |         if isinstance(pre_trained_models, torch.nn.Module):
 68 |             self.gp_ensemble = None
 69 |             self.B_x = {}
 70 |             x_dims = len(my_quad.get_state(quaternion=True, stacked=True))
 71 |             pred_dims = [7, 8, 9] + ([10, 11, 12] if model_conf['torque_output'] else [])
 72 |             for y_dim in pred_dims:
 73 |                 self.B_x[y_dim] = make_bx_matrix(x_dims, [y_dim])
 74 |                 self.mlp = pre_trained_models
 75 | 
 76 |         elif pre_trained_models is not None:
 77 |             self.gp_ensemble = restore_gp_regressors(pre_trained_models)
 78 |             x_dims = len(my_quad.get_state(quaternion=True, stacked=True))
 79 |             self.B_x = {}
 80 |             for y_dim in self.gp_ensemble.gp.keys():
 81 |                 self.B_x[y_dim] = make_bx_matrix(x_dims, [y_dim])
 82 | 
 83 |         else:
 84 |             self.gp_ensemble = None
 85 |             self.B_x = {}  # Selection matrix of the GP regressor-modified system states
 86 | 
 87 |         # For MPC optimization use
 88 |         self.quad_opt = Quad3DOptimizer(my_quad, t_horizon=t_horizon, n_nodes=n_nodes,
 89 |                                         q_cost=q_cost, r_cost=r_cost,
 90 |                                         B_x=self.B_x, gp_regressors=self.gp_ensemble,
 91 |                                         mlp_regressor=self.mlp, mlp_conf=model_conf,
 92 |                                         model_name=model_name, q_mask=q_mask,
 93 |                                         solver_options=solver_options, rdrv_d_mat=rdrv_d_mat)
 94 | 
 95 |     def clear(self):
 96 |         self.quad_opt.clear_acados_model()
 97 | 
 98 |     def get_state(self):
 99 |         """
100 |         Returns the state of the drone, with the angle described as a wxyz quaternion
101 |         :return: 13x1 array with the drone state: [p_xyz, a_wxyz, v_xyz, r_xyz]
102 |         """
103 | 
104 |         x = np.expand_dims(self.quad.get_state(quaternion=True, stacked=True), 1)
105 |         return x
106 | 
107 |     def set_reference(self, x_reference, u_reference=None):
108 |         """
109 |         Sets a target state for the MPC optimizer
110 |         :param x_reference: list with 4 sub-components (position, angle quaternion, velocity, body rate). If these four
111 |         are lists, then this means a single target point is used. If they are Nx3 and Nx4 (for quaternion) numpy arrays,
112 |         then they are interpreted as a sequence of N tracking points.
113 |         :param u_reference: Optional target for the optimized control inputs
114 |         """
115 | 
116 |         if isinstance(x_reference[0], list):
117 |             # Target state is just a point
118 |             return self.quad_opt.set_reference_state(x_reference, u_reference)
119 |         else:
120 |             # Target state is a sequence
121 |             return self.quad_opt.set_reference_trajectory(x_reference, u_reference)
122 | 
123 |     def optimize(self, use_model=0, return_x=False):
124 |         """
125 |         Runs MPC optimization to reach the pre-set target.
126 |         :param use_model: Integer. Select which dynamics model to use from the available options.
127 |         :param return_x: bool, whether to also return the optimized sequence of states alongside with the controls.
128 | 
129 |         :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), ...,
130 |         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
131 |         state prediction.
132 |         """
133 | 
134 |         quad_current_state = self.quad.get_state(quaternion=True, stacked=True)
135 |         quad_gp_state = self.quad.get_gp_state(quaternion=True, stacked=True)
136 | 
137 |         # Remove rate state for simplified model NLP
138 |         out_out = self.quad_opt.run_optimization(quad_current_state, use_model=use_model, return_x=return_x,
139 |                                                  gp_regression_state=quad_gp_state)
140 |         return out_out
141 | 
142 |     def simulate(self, ref_u):
143 |         """
144 |         Runs the simulation step for the dynamics model of the quadrotor 3D.
145 | 
146 |         :param ref_u: 4-length reference vector of control inputs
147 |         """
148 | 
149 |         # Simulate step
150 |         self.quad.update(ref_u, self.simulation_dt)
151 | 
152 |     def simulate_plant(self, w_opt, t_horizon=None, dt_vec=None, progress_bar=False):
153 |         """
154 |         Given a sequence of n inputs, evaluates the simulated discrete-time plant model n steps into the future. The
155 |         current drone state will not be changed by calling this method.
156 |         :param w_opt: sequence of control n x m control inputs, where n is the number of steps and m is the
157 |         dimensionality of a control input.
158 |         :param t_horizon: time corresponding to the duration of the n control inputs. In the case that the w_opt comes
159 |         from an MPC optimization, this parameter should be the MPC time horizon.
160 |         :param dt_vec: a vector of timestamps, the same length as w_opt, corresponding to the total time each input is
161 |         applied.
162 |         :param progress_bar: boolean - whether to show a progress bar on the console or not.
163 |         :return: the sequence of simulated quadrotor states.
164 |         """
165 | 
166 |         if t_horizon is None and dt_vec is None:
167 |             t_horizon = self.t_horizon
168 | 
169 |         return simulate_plant(self.quad, w_opt, simulation_dt=self.simulation_dt, simulate_func=self.simulate,
170 |                               t_horizon=t_horizon, dt_vec=dt_vec, progress_bar=progress_bar)
171 | 
172 |     def forward_prop(self, x_0, w_opt, cov_0=None, t_horizon=None, dt=None, use_gp=False, use_model=0):
173 |         """
174 |         Computes the forward propagation of the state using an MPC-optimized control input sequence.
175 |         :param x_0: Initial n-length state vector
176 |         :param w_opt: Optimized m*4-length sequence of control inputs from MPC, with the vector format:
177 |         [u_1(1), u_2(1), u_3(1), u_4(1), ..., u_3(m), u_4(m)]
178 |         :param cov_0: Initial covariance estimate (default 0). Can be either a positive semi-definite matrix or a
179 |         1D vector, which will be the diagonal of the covariance matrix. In both cases, the resulting covariance matrix
180 |         must be nxn shape, where n is the length of x_0.
181 |         :param t_horizon: time span of the control inputs (default is the time horizon of the MPC)
182 |         :param dt: Optional. Vector of length m, with the corresponding integration time for every control input in
183 |         w_opt. If none is provided, the default integration step is used.
184 |         :param use_gp: Boolean, whether to use GP regressors when performing the integration or not.
185 |         :param use_model: Integer. Select which dynamics model to use from the available options.
186 |         :return: An m x n array of propagated (expected) state estimates, and an m x n x n array with the m propagated
187 |         covariance matrices.
188 |         """
189 | 
190 |         # Default parameters
191 |         if dt is not None:
192 |             assert len(dt) == len(w_opt) / 4
193 |             t_horizon = np.sum(dt)
194 |         if t_horizon is None:
195 |             t_horizon = self.t_horizon
196 |         if cov_0 is None:
197 |             cov_0 = np.diag(np.zeros_like(np.squeeze(x_0)))
198 |         elif len(cov_0.shape) == 1:
199 |             cov_0 = np.diag(cov_0)
200 |         elif len(cov_0.shape) > 2:
201 |             TypeError("The initial covariance value must be either a 1D vector of a 2D matrix")
202 | 
203 |         gp_ensemble = self.gp_ensemble if use_gp else None
204 | 
205 |         # Compute forward propagation of state pdf
206 |         return uncertainty_forward_propagation(x_0, w_opt, t_horizon=t_horizon, covar=cov_0, dt=dt,
207 |                                                discrete_dynamics_f=self.quad_opt.discretize_f_and_q,
208 |                                                dynamics_jac_f=self.quad_opt.quad_xdot_jac,
209 |                                                B_x=self.B_x, gp_regressors=gp_ensemble,
210 |                                                use_model=use_model, m_integration_steps=1)
211 | 
212 |     @staticmethod
213 |     def reshape_input_sequence(u_seq):
214 |         """
215 |         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)]
216 |         to a 2D n x m array.
217 |         :param u_seq: 1D input sequence
218 |         :return: 2D input sequence, were n is the number of control inputs and m is the dimension of a single input.
219 |         """
220 | 
221 |         k = np.arange(u_seq.shape[0] / 4, dtype=int)
222 |         u_seq = np.atleast_2d(u_seq).T if len(u_seq.shape) == 1 else u_seq
223 |         u_seq = np.concatenate((u_seq[4 * k], u_seq[4 * k + 1], u_seq[4 * k + 2], u_seq[4 * k + 3]), 1)
224 |         return u_seq
225 | 
226 |     def reset(self):
227 |         return
228 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/utils/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/TUM-AAS/neural-mpc/5f15661c2d1a3b6072911d8993214c50a68e11eb/ros_dd_mpc/src/utils/__init__.py


--------------------------------------------------------------------------------
/ros_dd_mpc/src/utils/animator.py:
--------------------------------------------------------------------------------
  1 | """ Class for generating a comprehensive post-processed visualization of experimental flight results.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | import matplotlib.pyplot as plt
 15 | import numpy as np
 16 | from mpl_toolkits.mplot3d import axes3d
 17 | import matplotlib.animation as animation
 18 | from config.configuration_parameters import DirectoryConfig
 19 | 
 20 | 
 21 | class Dynamic3DTrajectory:
 22 |     def __init__(self, pos_data, vel_data, pos_ref, vel_ref, t_vec_ref, legends, sparsing_factor=1):
 23 | 
 24 |         # Add reference to data so it also moves
 25 |         pos_data = pos_ref + pos_data
 26 |         self.data = np.array(pos_data)
 27 | 
 28 |         self.reference = pos_ref[0]
 29 |         self.data_len = len(pos_data[0])
 30 |         self.n_lines = len(pos_data)
 31 | 
 32 |         self.t_vec = t_vec_ref - t_vec_ref[0]
 33 | 
 34 |         self.legends = legends
 35 | 
 36 |         vel_data = vel_ref + vel_data
 37 |         self.vel_data = np.array(vel_data)
 38 | 
 39 |         if sparsing_factor == 0:
 40 |             sparse_data = np.arange(0, self.data_len)
 41 |         else:
 42 |             sparse_data = np.arange(0, self.data_len, sparsing_factor)
 43 | 
 44 |         self.sparsed_data = self.data[:, sparse_data, :]
 45 |         self.vel_data = self.vel_data[:, sparse_data, :]
 46 |         self.t_vec = self.t_vec[sparse_data]
 47 | 
 48 |         self.max_buffer_size = 50
 49 |         self.colors = ["tab:blue", "tab:orange", "tab:green", "tab:red"]
 50 | 
 51 |         self.data_len = len(sparse_data)
 52 | 
 53 |         x_data = np.concatenate(tuple([dat[sparse_data, 0] for dat in pos_data]))
 54 |         y_data = np.concatenate(tuple([dat[sparse_data, 1] for dat in pos_data]))
 55 |         z_data = np.concatenate(tuple([dat[sparse_data, 2] for dat in pos_data]))
 56 | 
 57 |         self.max_x = np.max(x_data)
 58 |         self.min_x = np.min(x_data)
 59 |         self.max_y = np.max(y_data)
 60 |         self.min_y = np.min(y_data)
 61 |         self.max_z = np.max(z_data)
 62 |         self.min_z = np.min(z_data)
 63 | 
 64 |         # Make dimensions more similar and add a bit of padding
 65 |         range_x = self.max_x - self.min_x
 66 |         range_y = self.max_y - self.min_y
 67 |         range_z = self.max_z - self.min_z
 68 |         max_range = max(range_x, range_y, range_z)
 69 |         self.max_x = self.max_x + 0.25 * (max_range - range_x) + (self.max_x - self.min_x) * 0.2
 70 |         self.min_x = self.min_x - 0.25 * (max_range - range_x) - (self.max_x - self.min_x) * 0.2
 71 |         self.max_y = self.max_y + 0.25 * (max_range - range_y) + (self.max_y - self.min_y) * 0.2
 72 |         self.min_y = self.min_y - 0.25 * (max_range - range_y) - (self.max_y - self.min_y) * 0.2
 73 |         self.max_z = self.max_z + 0.25 * (max_range - range_z) + (self.max_z - self.min_z) * 0.2
 74 |         self.min_z = self.min_z - 0.25 * (max_range - range_z) - (self.max_z - self.min_z) * 0.2
 75 | 
 76 |         self.figure = None
 77 |         self.ax = None
 78 |         self.pos_err_ax = None
 79 |         self.speed_ax = None
 80 |         self.top_down_ax = None
 81 |         self.x_time_ax = None
 82 |         self.y_time_ax = None
 83 | 
 84 |         self.lines = None
 85 | 
 86 |         self.n_3d_lines = None
 87 |         self.vel_bars_0_idx = None
 88 |         self.pos_err_0_idx = None
 89 |         self.top_down_0_idx = None
 90 |         self.x_time_0_idx = None
 91 |         self.y_time_0_idx = None
 92 | 
 93 |     def on_launch(self):
 94 |         self.figure = plt.figure(figsize=(14, 10))
 95 | 
 96 |         self.ax = axes3d.Axes3D(self.figure, rect=(-0.02, 0.3, 0.65, 0.7))
 97 | 
 98 |         self.ax.set_zlim3d([self.min_z, self.max_z])
 99 |         self.ax.set_ylim3d([self.min_y, self.max_y])
100 |         self.ax.set_xlim3d([self.min_x, self.max_x])
101 |         self.ax.set_xlabel(r'$p_x\: [m]$')
102 |         self.ax.set_ylabel(r'$p_y\: [m]$')
103 |         self.ax.set_zlabel(r'$p_z\: [m]$')
104 | 
105 |         self.ax.plot(self.reference[:, 0], self.reference[:, 1], self.reference[:, 2], '-', alpha=0.2)
106 | 
107 |         self.lines = sum([self.ax.plot([], [], [], '-', c=self.colors[i], label=self.legends[i])
108 |                           for i in range(self.n_lines)], [])
109 |         projection_lines = sum([self.ax.plot([], [], [], '-', c=self.colors[i], alpha=0.2)
110 |                                 for i in range(self.n_lines)], [])
111 |         projection_lines += sum([self.ax.plot([], [], [], '-', c=self.colors[i], alpha=0.2)
112 |                                  for i in range(self.n_lines)], [])
113 |         projection_lines += sum([self.ax.plot([], [], [], '-', c=self.colors[i], alpha=0.2)
114 |                                  for i in range(self.n_lines)], [])
115 | 
116 |         self.lines += projection_lines
117 | 
118 |         pos_balls = sum([self.ax.plot([], [], [], 'o', c=self.colors[i]) for i in range(self.n_lines)], [])
119 |         self.lines += pos_balls
120 |         self.n_3d_lines = len(self.lines)
121 | 
122 |         self.ax.legend()
123 |         self.ax.set_title('3D Visualization')
124 | 
125 |         self.vel_bars_0_idx = len(self.lines)
126 |         bar_height = 0.1
127 | 
128 |         self.speed_ax = plt.axes((0.65, 0.8, 0.3, 0.05))
129 |         vel_bar = [self.speed_ax.barh([0], [0], color=self.colors[0], height=bar_height)[0]]
130 |         self.speed_ax.set_xlim([0, np.max(np.sqrt(np.sum(self.vel_data[0] ** 2, -1))) * 1.05])
131 |         self.speed_ax.set_ylim([-bar_height, bar_height * 1.05])
132 | 
133 |         self.speed_ax.set_xlabel(r'$\|\mathbf{v}\|\:[m/s]$')
134 |         self.speed_ax.set_yticks([0])
135 |         self.speed_ax.set_yticklabels([self.legends[0]])
136 |         self.speed_ax.grid()
137 |         self.speed_ax.set_title('Quadrotor speed')
138 | 
139 |         self.lines += vel_bar
140 | 
141 |         self.pos_err_ax = plt.axes((0.65, 0.6, 0.3, 0.1))
142 |         self.pos_err_ax.grid()
143 |         pos_err = []
144 |         p_errors = np.stack([self.data[0] - self.data[i+1] for i in range(self.n_lines - 1)])
145 | 
146 |         pos_err += [self.pos_err_ax.barh([(i - 1) * bar_height * 1.05], [i], color=self.colors[i], height=bar_height)[0]
147 |                     for i in range(1, self.n_lines)]
148 |         self.pos_err_ax.set_xlim([np.min(p_errors) * 1.05, np.max(p_errors) * 1.05])
149 |         self.pos_err_ax.set_ylim([-bar_height, ((self.n_lines - 1) * 1.05) * bar_height])
150 |         self.pos_err_ax.set_title('XY position error')
151 |         self.pos_err_ax.set_xlabel(r'$\|\mathbf{p}^* - \mathbf{p}\|\: [m]$')
152 |         self.pos_err_ax.set_yticks([i * bar_height * 1.05 for i in range(self.n_lines - 1)])
153 |         self.pos_err_ax.set_yticklabels(self.legends[1:])
154 | 
155 |         self.pos_err_0_idx = len(self.lines)
156 |         self.lines += pos_err
157 | 
158 |         self.top_down_ax = plt.axes((0.65, 0.08, 0.3, 0.4))
159 |         top_down_lines = sum([self.top_down_ax.plot([], [], color=self.colors[i], label=self.legends[i])
160 |                               for i in range(self.n_lines)], [])
161 |         top_down_lines += sum([self.top_down_ax.plot([], [], 'o', color=self.colors[i])
162 |                                for i in range(self.n_lines)], [])
163 |         self.top_down_ax.plot(self.reference[:, 0], self.reference[:, 1], '-', alpha=0.2)
164 |         self.top_down_0_idx = len(self.lines)
165 |         self.top_down_ax.grid()
166 |         self.top_down_ax.set_xlim([self.min_x, self.max_x])
167 |         self.top_down_ax.set_ylim([self.min_y, self.max_y])
168 |         self.top_down_ax.set_title('Top down view')
169 |         self.top_down_ax.set_xlabel(r'$p_x\:[m]$')
170 |         self.top_down_ax.set_ylabel(r'$p_y\:[m]$')
171 |         self.top_down_ax.legend()
172 |         self.lines += top_down_lines
173 | 
174 |         self.x_time_ax = plt.axes((0.05, 0.08, 0.26, 0.2))
175 |         self.x_time_ax.axhline(y=0, linestyle='--', alpha=0.5)
176 |         x_time_lines = sum([self.x_time_ax.plot([], [], color=self.colors[i], label=self.legends[i])
177 |                             for i in range(1, self.n_lines)], [])
178 |         self.x_time_0_idx = len(self.lines)
179 |         self.lines += x_time_lines
180 |         self.x_time_ax.set_xlim([self.t_vec[0], self.t_vec[-1]])
181 |         self.x_time_ax.set_ylim([0, np.max(np.abs(p_errors[:, :, 0])) * 1.05])
182 |         self.x_time_ax.grid()
183 |         self.x_time_ax.set_title(r'$p_x$ error')
184 |         self.x_time_ax.set_ylabel('pos [m]')
185 |         self.x_time_ax.legend(loc='upper right')
186 | 
187 |         self.y_time_ax = plt.axes((0.34, 0.08, 0.26, 0.2))
188 |         self.y_time_ax.axhline(y=0, linestyle='--', alpha=0.5)
189 |         y_time_lines = sum([self.y_time_ax.plot([], [], color=self.colors[i], label=self.legends[i])
190 |                             for i in range(1, self.n_lines)], [])
191 |         self.y_time_0_idx = len(self.lines)
192 |         self.lines += y_time_lines
193 |         self.y_time_ax.set_xlim([self.t_vec[0], self.t_vec[-1]])
194 |         self.y_time_ax.set_ylim([0, np.max(np.abs(p_errors[:, :, 1])) * 1.05])
195 |         self.y_time_ax.set_xlabel('Time [s]')
196 |         self.y_time_ax.set_title(r'$p_y$ error')
197 |         self.y_time_ax.grid()
198 |         self.y_time_ax.legend(loc='upper right')
199 | 
200 |     def on_init(self):
201 |         for i in range(self.n_3d_lines):
202 |             self.lines[i].set_data([], [])
203 |             self.lines[i].set_3d_properties([])
204 | 
205 |         self.lines[self.vel_bars_0_idx].set_width(0)
206 | 
207 |         for i in range(self.n_lines - 1):
208 |             self.lines[self.pos_err_0_idx + i].set_width(0)
209 | 
210 |         for i in range(self.n_lines * 2):
211 |             self.lines[self.top_down_0_idx + i].set_data([], [])
212 | 
213 |         for i in range(self.n_lines - 1):
214 |             self.lines[self.x_time_0_idx + i].set_data([], [])
215 |             self.lines[self.y_time_0_idx + i].set_data([], [])
216 | 
217 |         return self.lines
218 | 
219 |     def animate(self, i):
220 |         #i = (2 * i) % self.data.shape[1]
221 |         for j, (line, xi) in enumerate(zip(self.lines[:self.n_lines], self.sparsed_data)):
222 |             x, y, z = xi[:i].T
223 | 
224 |             if len(x) > self.max_buffer_size:
225 |                 x = x[len(x) - self.max_buffer_size:]
226 |                 y = y[len(y) - self.max_buffer_size:]
227 |                 z = z[len(z) - self.max_buffer_size:]
228 | 
229 |             line.set_data(x, y)
230 |             line.set_3d_properties(z)
231 | 
232 |             self.lines[j + self.n_lines].set_data(x, self.max_y)
233 |             self.lines[j + self.n_lines].set_3d_properties(z)
234 | 
235 |             self.lines[j + 2 * self.n_lines].set_data(np.ones(len(y)) * self.min_x, y)
236 |             self.lines[j + 2 * self.n_lines].set_3d_properties(z)
237 | 
238 |             self.lines[j + 3 * self.n_lines].set_data(x, y)
239 |             self.lines[j + 3 * self.n_lines].set_3d_properties(self.min_z)
240 | 
241 |             if len(x) > 0:
242 |                 self.lines[j + 4 * self.n_lines].set_data([x[-1]], [y[-1]])
243 |                 self.lines[j + 4 * self.n_lines].set_3d_properties([z[-1]])
244 | 
245 |                 self.lines[self.top_down_0_idx + j + self.n_lines].set_data(x[-1], y[-1])
246 | 
247 |             self.lines[self.top_down_0_idx + j].set_data(x, y)
248 | 
249 |         self.lines[self.vel_bars_0_idx].set_width(np.sqrt(np.sum(self.vel_data[0, i, :] ** 2)))
250 | 
251 |         for j in range(self.n_lines - 1):
252 |             self.lines[self.pos_err_0_idx + j].set_width(
253 |                 np.sqrt(np.sum((self.data[0][i, :2] - self.data[j+1][i, :2]) ** 2)))
254 | 
255 |         for j in reversed(range(self.n_lines - 1)):
256 |             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]))
257 |             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]))
258 | 
259 |         return self.lines
260 | 
261 |     def __call__(self, save=False):
262 |         self.on_launch()
263 | 
264 |         ani = animation.FuncAnimation(self.figure, self.animate, init_func=self.on_init, frames=self.data_len,
265 |                                       interval=20, blit=True, repeat=False)
266 | 
267 |         if save:
268 |             results_dir = DirectoryConfig.RESULTS_DIR
269 |             plt.rcParams['animation.ffmpeg_path'] = '/usr/local/bin/ffmpeg'
270 |             writer = animation.FFMpegWriter(fps=50)
271 |             ani.save(results_dir + '/animation.mp4', writer=writer)
272 | 
273 |         else:
274 |             plt.show()
275 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/utils/ground_map.py:
--------------------------------------------------------------------------------
 1 | import numpy as np
 2 | 
 3 | 
 4 | class GroundMap:
 5 |     def __init__(self, horizon=((-1, 1), (-1, 1)), resolution=0.1):
 6 |         assert ((horizon[0][1] - horizon[0][0]) / resolution).is_integer()
 7 |         assert ((horizon[1][1] - horizon[1][0]) / resolution).is_integer()
 8 |         self._horizon = horizon
 9 |         self._resolution = resolution
10 | 
11 |     def at(self, center: np.array):
12 |         map_center_to_map_origin = np.array([self._horizon[0][0], self._horizon[1][0]])
13 |         map_origin = center + map_center_to_map_origin
14 |         return self.draw(map_origin - 0.5 * self._resolution, self.empty_map), map_origin
15 | 
16 |     def draw(self, pos, map):
17 |         raise NotImplementedError
18 | 
19 |     @property
20 |     def empty_map(self):
21 |         x_len = int((self._horizon[0][1] - self._horizon[0][0]) / self._resolution) + 1
22 |         y_len = int((self._horizon[1][1] - self._horizon[1][0]) / self._resolution) + 1
23 |         empty_map = np.zeros((x_len, y_len))
24 |         return empty_map
25 | 
26 | 
27 | class GroundMapWithBox(GroundMap):
28 |     def __init__(self, box_min, box_max, height, *args, **kwargs):
29 |         super().__init__(*args, **kwargs)
30 |         self._box_min = box_min
31 |         self._box_max = box_max
32 |         self._height = height
33 | 
34 |     def draw(self, pos, map):
35 |         x_from = max(int((self._box_min[0] - pos[0]) // self._resolution), 0)
36 |         x_to = min(int((self._box_max[0] - pos[0]) // self._resolution), map.shape[0] - 1)
37 |         y_from = max(int((self._box_min[1] - pos[1]) // self._resolution), 0)
38 |         y_to = min(int((self._box_max[1] - pos[1]) // self._resolution), map.shape[1] - 1)
39 |         map[x_from:x_to, y_from:y_to] = self._height
40 |         return map
41 | 
42 | if __name__ == '__main__':
43 |     horizon = ((-7, 7), (-7, 7))
44 |     box_min_ = (-4.25, 9.37)
45 |     box_max_ = (-2.76, 10.13)
46 |     box_height_ = 1.0
47 | 
48 |     map = GroundMapWithBox(box_min_, box_max_, box_height_, horizon=horizon)
49 | 
50 |     map_at, o = map.at(np.array([3, 9]))
51 | 
52 |     print(map_at)
53 | 
54 |     m_x = int((1.01 + o[0]) / 0.1)
55 | 
56 |     print(m_x)


--------------------------------------------------------------------------------
/ros_dd_mpc/src/utils/keyframe_3d_gen.py:
--------------------------------------------------------------------------------
  1 | """ Generates a set of keypoints to generate a piece-wise polynomial trajectory between each pair.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | 
 15 | from sklearn.gaussian_process import GaussianProcessRegressor
 16 | from sklearn.gaussian_process.kernels import ExpSineSquared
 17 | import numpy as np
 18 | import matplotlib.pyplot as plt
 19 | from mpl_toolkits.mplot3d import Axes3D
 20 | 
 21 | 
 22 | def apply_map_limits(x, y, z, limits):
 23 | 
 24 |     # Find out which axis is the most constrained
 25 |     x_max_range = limits["x"][1] - limits["x"][0]
 26 |     y_max_range = limits["y"][1] - limits["y"][0]
 27 |     z_max_range = limits["z"][1] - limits["z"][0]
 28 | 
 29 |     x_actual_range = np.max(x) - np.min(x)
 30 |     y_actual_range = np.max(y) - np.min(y)
 31 |     z_actual_range = np.max(z) - np.min(z)
 32 | 
 33 |     # One or more of the ranges violates the constraints.
 34 |     if x_actual_range > x_max_range or y_actual_range > y_max_range or z_actual_range > z_max_range:
 35 |         shrink_ratio = max(x_actual_range / x_max_range, y_actual_range / y_max_range, z_actual_range / z_max_range)
 36 |         x = (x - np.mean(x)) / shrink_ratio
 37 |         y = (y - np.mean(y)) / shrink_ratio
 38 |         z = (z - np.mean(z)) / shrink_ratio
 39 | 
 40 |         x += (limits["x"][0] - np.min(x))
 41 |         y += (limits["y"][0] - np.min(y))
 42 |         z += (limits["z"][0] - np.min(z))
 43 | 
 44 |     return x, y, z
 45 | 
 46 | 
 47 | def center_and_scale(x, y, z, x_max, x_min, y_max, y_min, z_max, z_min):
 48 | 
 49 |     x -= (x_min + (x_max - x_min) / 2)
 50 |     y -= (y_min + (y_max - y_min) / 2)
 51 |     z -= (z_min + (z_max - z_min) / 2)
 52 | 
 53 |     scaling = np.mean([x_max - x_min, y_max - y_min, z_max - z_min])
 54 |     x = x * 6 / scaling
 55 |     y = y * 6 / scaling
 56 |     z = z * 6 / scaling
 57 | 
 58 |     return x, y, z
 59 | 
 60 | 
 61 | def random_periodical_trajectory(plot=False, random_state=None, map_limits=None):
 62 | 
 63 |     if random_state is None:
 64 |         random_state = np.random.randint(0, 9999)
 65 | 
 66 |     kernel_z = ExpSineSquared(length_scale=5.5, periodicity=60)
 67 |     kernel_y = ExpSineSquared(length_scale=4.5, periodicity=30) + ExpSineSquared(length_scale=4.0, periodicity=15)
 68 |     kernel_x = ExpSineSquared(length_scale=4.5, periodicity=30) + ExpSineSquared(length_scale=4.5, periodicity=60)
 69 | 
 70 |     gp_x = GaussianProcessRegressor(kernel=kernel_x)
 71 |     gp_y = GaussianProcessRegressor(kernel=kernel_y)
 72 |     gp_z = GaussianProcessRegressor(kernel=kernel_z)
 73 | 
 74 |     # High resolution sampling for track boundaries
 75 |     inputs_x = np.linspace(0, 60, 100)
 76 |     inputs_y = np.linspace(0, 30, 100)
 77 |     inputs_z = np.linspace(0, 60, 100)
 78 | 
 79 |     x_sample_hr = gp_x.sample_y(inputs_x[:, np.newaxis], 1, random_state=random_state)
 80 |     y_sample_hr = gp_y.sample_y(inputs_y[:, np.newaxis], 1, random_state=random_state)
 81 |     z_sample_hr = gp_z.sample_y(inputs_z[:, np.newaxis], 1, random_state=random_state)
 82 | 
 83 |     max_x_coords = np.max(x_sample_hr, 0)
 84 |     max_y_coords = np.max(y_sample_hr, 0)
 85 |     max_z_coords = np.max(z_sample_hr, 0)
 86 | 
 87 |     min_x_coords = np.min(x_sample_hr, 0)
 88 |     min_y_coords = np.min(y_sample_hr, 0)
 89 |     min_z_coords = np.min(z_sample_hr, 0)
 90 | 
 91 |     x_sample_hr, y_sample_hr, z_sample_hr = center_and_scale(
 92 |         x_sample_hr, y_sample_hr, z_sample_hr,
 93 |         max_x_coords, min_x_coords, max_y_coords, min_y_coords, max_z_coords, min_z_coords)
 94 | 
 95 |     # Additional constraint on map limits
 96 |     if map_limits is not None:
 97 |         x_sample_hr, y_sample_hr, z_sample_hr = apply_map_limits(x_sample_hr, y_sample_hr, z_sample_hr, map_limits)
 98 | 
 99 |     # Low resolution for control points
100 |     lr_ind = np.linspace(0, len(x_sample_hr) - 1, 10, dtype=int)
101 |     lr_ind[-1] = 0
102 |     x_sample_lr = x_sample_hr[lr_ind, :]
103 |     y_sample_lr = y_sample_hr[lr_ind, :]
104 |     z_sample_lr = z_sample_hr[lr_ind, :]
105 | 
106 |     x_sample_diff = x_sample_hr[lr_ind + 1, :] - x_sample_lr
107 |     y_sample_diff = y_sample_hr[lr_ind + 1, :] - y_sample_lr
108 |     z_sample_diff = z_sample_hr[lr_ind + 1, :] - z_sample_lr
109 | 
110 |     # Get angles at keypoints
111 |     a_x = np.arctan2(z_sample_diff, y_sample_diff) * 0
112 |     a_y = np.arctan2(x_sample_diff, z_sample_diff) * 0
113 |     a_z = (np.arctan2(y_sample_diff, x_sample_diff) - np.pi/4) * 0
114 | 
115 |     if plot:
116 |         # Plot checking
117 |         # Build rotation matrices
118 |         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]]
119 |         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]]
120 |         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]]
121 | 
122 |         main_axes = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
123 |         quiver_axes = np.zeros((len(lr_ind), 3, 3))
124 | 
125 |         for i in range(len(lr_ind)):
126 |             r_mat = rz[i].dot(ry[i].dot(rx[i]))
127 |             rot_body = r_mat.dot(main_axes)
128 |             quiver_axes[i, :, :] = rot_body
129 | 
130 |         shortest_axis = min(np.max(x_sample_hr) - np.min(x_sample_hr),
131 |                             np.max(y_sample_hr) - np.min(y_sample_hr),
132 |                             np.max(z_sample_hr) - np.min(z_sample_hr))
133 |         fig = plt.figure()
134 |         ax = fig.add_subplot(111, projection='3d')
135 |         ax.scatter(x_sample_lr, y_sample_lr, z_sample_lr)
136 |         ax.plot(x_sample_hr[:, 0], y_sample_hr[:, 0], z_sample_hr[:, 0], '-', alpha=0.5)
137 |         ax.quiver(x_sample_lr[:, 0], y_sample_lr[:, 0], z_sample_lr[:, 0],
138 |                   x_sample_diff[:, 0], y_sample_diff[:, 0], z_sample_diff[:, 0], color='g',
139 |                   length=shortest_axis / 6, normalize=True, label='traj. norm')
140 |         ax.quiver(x_sample_lr, y_sample_lr, z_sample_lr,
141 |                   quiver_axes[:, 0, :], quiver_axes[:, 1, :], quiver_axes[:, 2, :], color='b',
142 |                   length=shortest_axis / 6, normalize=True, label='quad. att.')
143 |         ax.tick_params(labelsize=16)
144 |         ax.legend(fontsize=16)
145 |         ax.set_xlabel('x [m]', size=16, labelpad=10)
146 |         ax.set_ylabel('y [m]', size=16, labelpad=10)
147 |         ax.set_zlabel('z [m]', size=16, labelpad=10)
148 |         ax.set_xlim([np.min(x_sample_hr), np.max(x_sample_hr)])
149 |         ax.set_ylim([np.min(y_sample_hr), np.max(y_sample_hr)])
150 |         ax.set_zlim([np.min(z_sample_hr), np.max(z_sample_hr)])
151 |         ax.set_title('Source keypoints', size=18)
152 |         plt.show()
153 | 
154 |     curve = np.concatenate((x_sample_lr, y_sample_lr, z_sample_lr), 1)
155 |     attitude = np.concatenate((a_x, a_y, a_z), 1)
156 | 
157 |     return curve, attitude
158 | 
159 | 
160 | if __name__ == "__main__":
161 |     limits = {
162 |         "x": [-0.6, 4],
163 |         "y": [-2, 2],
164 |         "z": [0.1, 2]
165 |     }
166 |     random_periodical_trajectory(plot=True, map_limits=limits)
167 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/utils/quad_3d_opt_utils.py:
--------------------------------------------------------------------------------
  1 | """ Set of utility functions for the quadrotor optimizer and simplified simulator.
  2 | 
  3 | This program is free software: you can redistribute it and/or modify it under
  4 | the terms of the GNU General Public License as published by the Free Software
  5 | Foundation, either version 3 of the License, or (at your option) any later
  6 | version.
  7 | This program is distributed in the hope that it will be useful, but WITHOUT
  8 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  9 | FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 10 | You should have received a copy of the GNU General Public License along with
 11 | this program. If not, see <http://www.gnu.org/licenses/>.
 12 | """
 13 | 
 14 | import casadi as cs
 15 | import numpy as np
 16 | from src.quad_mpc.quad_3d import Quadrotor3D
 17 | from tqdm import tqdm
 18 | 
 19 | 
 20 | def discretize_dynamics_and_cost(t_horizon, n_points, m_steps_per_point, x, u, dynamics_f, cost_f, ind):
 21 |     """
 22 |     Integrates the symbolic dynamics and cost equations until the time horizon using a RK4 method.
 23 |     :param t_horizon: time horizon in seconds
 24 |     :param n_points: number of control input points until time horizon
 25 |     :param m_steps_per_point: number of integrations steps per control input
 26 |     :param x: 4-element list with symbolic vectors for position (3D), angle (4D), velocity (3D) and rate (3D)
 27 |     :param u: 4-element symbolic vector for control input
 28 |     :param dynamics_f: symbolic dynamics function written in CasADi symbolic syntax.
 29 |     :param cost_f: symbolic cost function written in CasADi symbolic syntax. If None, then cost 0 is returned.
 30 |     :param ind: Only used for trajectory tracking. Index of cost function to use.
 31 |     :return: a symbolic function that computes the dynamics integration and the cost function at n_control_inputs
 32 |     points until the time horizon given an initial state and
 33 |     """
 34 | 
 35 |     if isinstance(cost_f, list):
 36 |         # Select the list of cost functions
 37 |         cost_f = cost_f[ind * m_steps_per_point:(ind + 1) * m_steps_per_point]
 38 |     else:
 39 |         cost_f = [cost_f] * m_steps_per_point
 40 | 
 41 |     # Fixed step Runge-Kutta 4 integrator
 42 |     dt = t_horizon / n_points / m_steps_per_point
 43 |     x0 = x
 44 |     q = 0
 45 | 
 46 |     for j in range(m_steps_per_point):
 47 |         k1 = dynamics_f(x=x, u=u)['x_dot']
 48 |         k2 = dynamics_f(x=x + dt / 2 * k1, u=u)['x_dot']
 49 |         k3 = dynamics_f(x=x + dt / 2 * k2, u=u)['x_dot']
 50 |         k4 = dynamics_f(x=x + dt * k3, u=u)['x_dot']
 51 |         x_out = x + dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
 52 | 
 53 |         x = x_out
 54 | 
 55 |         if cost_f and cost_f[j] is not None:
 56 |             q = q + cost_f[j](x=x, u=u)['q']
 57 | 
 58 |     return cs.Function('F', [x0, u], [x, q], ['x0', 'p'], ['xf', 'qf'])
 59 | 
 60 | 
 61 | def _forward_prop_core(x_0, u_seq, t_horizon, discrete_dynamics_f, dynamics_jac_f, B_x, gp_ensemble, covar, dt,
 62 |                        m_int_steps, use_model):
 63 |     """
 64 |     Propagates forward the state estimate described by the mean vector x_0 and the covariance matrix covar, and a
 65 |     sequence of inputs for the system u_seq. These inputs can either be numerical or symbolic.
 66 | 
 67 |     :param x_0: initial mean state of the state probability density function. Vector of length m
 68 |     :param u_seq: sequence of flattened N control inputs. I.e. vector of size N*4
 69 |     :param t_horizon: time horizon corresponding to sequence of inputs
 70 |     :param discrete_dynamics_f: symbolic function to compute the discrete dynamics of the system.
 71 |     :param dynamics_jac_f: symbolic function to compute the  dynamics jacobian of the system.
 72 |     :param B_x: Matrix to convert map from regressor output to state.
 73 |     :param gp_ensemble: The ensemble of GP's. Can be None if no GP's are used.
 74 |     :param covar: Initial covariance matrix of shape m x m, of the state probability density function.
 75 |     :param dt: Vector of N timestamps, the same length as w_opt / 2, corresponding to the time each input is applied.
 76 |     :param m_int_steps: number of intermediate integration steps per control node.
 77 |     :param use_model: The number (index) of the dynamics model to use from the available options.
 78 |     :return: The sequence of mean and covariance estimates for every corresponding input, as well as the computed
 79 |     cost for each stage.
 80 |     """
 81 | 
 82 |     # Reshape input sequence to a N x 4 (1D) vector. Assume control input dim = 4
 83 |     k = np.arange(int(u_seq.shape[0] / 4))
 84 |     input_sequence = cs.horzcat(u_seq[4 * k], u_seq[4 * k + 1], u_seq[4 * k + 2], u_seq[4 * k + 3])
 85 | 
 86 |     N = int(u_seq.shape[0] / 4)
 87 | 
 88 |     if dt is None:
 89 |         dt = t_horizon / N * np.ones((N, 1))
 90 | 
 91 |     if len(dt.shape) == 1:
 92 |         dt = np.expand_dims(dt, 1)
 93 | 
 94 |     # Initialize sequence of propagated states
 95 |     mu_x = [x_0]
 96 |     cov_x = [covar]
 97 |     cost_x = [0]
 98 | 
 99 |     for k in range(N):
100 | 
101 |         # Get current control input and current state mean and covariance
102 |         u_k = input_sequence[k, :]
103 |         mu_k = mu_x[k]
104 |         sig_k = cov_x[k]
105 | 
106 |         # mu(k+1) vector from propagation equations. Pass state through nominal dynamics with GP mean augmentation if GP
107 |         # is available. Otherwise use nominal dynamics only.
108 |         f_func = discrete_dynamics_f(dt[k], 1, m_int_steps, k, use_gp=gp_ensemble is not None, use_model=use_model)
109 | 
110 |         fk = f_func(x0=mu_k, p=u_k)
111 |         mu_next = fk['xf']
112 |         stage_cost = fk['qf']
113 | 
114 |         # K(k+1) matrix from propagation equations
115 |         K_mat = sig_k
116 | 
117 |         # Evaluate linearized dynamics at current state
118 |         l_mat = dynamics_jac_f(mu_k, u_k) * dt[k] + np.eye(mu_k.shape[0])
119 | 
120 |         # Adjust matrices if GP model available.
121 |         if gp_ensemble is not None:
122 | 
123 |             raise NotImplementedError  # TODO: re-implement covariance propagation with GP
124 |             # Get subset of features for GP regressor if GP regressor available
125 |             z_k = cs.mtimes(B_z, cs.vertcat(mu_k, u_k.T))
126 | 
127 |             # GP prediction. Stack predictions vertically along prediction dimension (e.g. vx, vy, ...)
128 |             _, gp_covar_preds, gp_noise_prior = gp_ensemble.predict(z_k, return_cov=True, gp_idx=use_model)
129 | 
130 |             # Covariance forward propagation.
131 |             l_mat = cs.horzcat(l_mat, B_x * dt[k])  # left-multiplying matrix
132 |             sig_ld_comp = cs.mtimes(gp_prediction_jac(z_k, B_x, B_z, gp_ensemble, use_model), sig_k)
133 |             K_mat = cs.vertcat(K_mat, sig_ld_comp)
134 |             K_mat = cs.horzcat(K_mat, cs.vertcat(sig_ld_comp.T, cs.diag(gp_covar_preds + gp_noise_prior)))
135 | 
136 |         # Add next state estimate to lists
137 |         mu_x += [mu_next]
138 |         cov_x += [cs.mtimes(cs.mtimes(l_mat, K_mat), l_mat.T)]
139 |         cost_x += [stage_cost]
140 | 
141 |     cov_x = [cov for cov in cov_x]
142 |     return mu_x, cov_x, cost_x
143 | 
144 | 
145 | def uncertainty_forward_propagation(x_0, u_seq, t_horizon, discrete_dynamics_f, dynamics_jac_f, B_x=None,
146 |                                     gp_regressors=None, covar=None, dt=None, m_integration_steps=1, use_model=0):
147 |     if covar is None:
148 |         covar = np.zeros((len(x_0), len(x_0)))
149 |     else:
150 |         assert covar.shape == (len(x_0), len(x_0))
151 | 
152 |     x_0 = np.array(x_0)
153 | 
154 |     mu_x, cov_x, _ = _forward_prop_core(x_0, u_seq, t_horizon, discrete_dynamics_f, dynamics_jac_f, B_x,
155 |                                         gp_regressors, covar, dt, m_integration_steps, use_model)
156 | 
157 |     mu_x = cs.horzcat(*mu_x)
158 |     cov_x = cs.horzcat(*cov_x)
159 | 
160 |     mu_prop = np.array(mu_x).T
161 |     cov_prop = np.array(cov_x).reshape((mu_prop.shape[1], mu_prop.shape[1], -1), order='F').transpose(2, 0, 1)
162 |     return mu_prop, cov_prop
163 | 
164 | 
165 | def gp_prediction_jac(z, Bx, Bz, gp_ensemble, gp_idx):
166 |     """
167 |     Computes the symbolic function for the Jacobian of the expected values of the joint GP regression model,
168 |     evaluated at point z.
169 | 
170 |     :param z: A dictionary of symbolic vector at which the Jacobian must be evaluated. Each entry in the dictionary is
171 |     indexed by the output dimension index. The dimension m_z on any given value must be the expected dimension of the
172 |     feature inputs for the regressor.
173 |     :param Bx: Matrix to convert map from regressor output to state.
174 |     :param Bz: Matrix to map from (stacked) state vector and input vector to feature vector. If the gp_ensemble is not
175 |     homogeneous, this parameter must be a dictionary specifying in each dimension which Bz matrix to use.
176 |     :param gp_ensemble: GPEnsemble object with all the gp regressors
177 |     :param gp_idx: which gp to compute the jacobian to from the ensemble
178 |     :return: A Jacobian matrix of size n x m_x, where n is the number of variables regressed by the joint GP
179 |     regressor and m_x is the dimension of the state vector.
180 |     """
181 | 
182 |     # Figure out which variables are necessary to compute the Jacobian wrt.
183 |     out_dims = np.matmul(Bx, np.ones((Bx.shape[1], 1)))
184 |     if isinstance(Bz, dict):
185 |         z_dims = {}
186 |         for dim in Bz.keys():
187 |             z_dims[dim] = np.matmul(Bz[dim][:, :len(out_dims)], out_dims)
188 |     else:
189 |         bz = np.matmul(Bz[:, :len(out_dims)], out_dims)
190 |         z_dims = {k: v for k, v in zip(z.keys(), [bz] * len(z.keys()))}
191 |         Bz = {k: v for k, v in zip(z.keys(), [Bz] * len(z.keys()))}
192 | 
193 |     jac = []
194 |     for dim in gp_idx.keys():
195 |         # Mapping from z to x
196 |         inv_Bz = Bz[dim][:, :len(out_dims)].T
197 | 
198 |         gp = gp_ensemble.gp[dim][gp_idx[dim][0]]
199 |         jac += [cs.mtimes(inv_Bz, gp.eval_gp_jac(z[dim]) * z_dims[dim])]
200 | 
201 |     return cs.horzcat(*jac).T
202 | 
203 | 
204 | def simulate_plant(quad, w_opt, simulation_dt, simulate_func, t_horizon=None, dt_vec=None, progress_bar=False):
205 |     """
206 |     Given a sequence of n inputs, evaluates the simulated discrete-time plant model n steps into the future. The
207 |     current drone state will not be changed by calling this method.
208 |     :param quad: Quadrotor3D simulator object
209 |     :type quad: Quadrotor3D
210 |     :param w_opt: sequence of control n x m control inputs, where n is the number of steps and m is the
211 |     dimensionality of a control input.
212 |     :param simulation_dt: simulation step
213 |     :param simulate_func: simulation function (inner loop) from the quadrotor MPC.
214 |     :param t_horizon: time corresponding to the duration of the n control inputs. In the case that the w_opt comes
215 |     from an MPC optimization, this parameter should be the MPC time horizon.
216 |     :param dt_vec: a vector of timestamps, the same length as w_opt, corresponding to the total time each input is
217 |     applied.
218 |     :param progress_bar: boolean - whether to show a progress bar on the console or not.
219 |     :return: the sequence of simulated quadrotor states.
220 |     """
221 | 
222 |     # Default_parameters
223 |     if t_horizon is None and dt_vec is None:
224 |         raise ValueError("At least the t_horizon or dt should be provided")
225 | 
226 |     if t_horizon is None:
227 |         t_horizon = np.sum(dt_vec)
228 | 
229 |     state_safe = quad.get_state(quaternion=True, stacked=True)
230 | 
231 |     # Compute true simulated trajectory
232 |     total_sim_time = t_horizon
233 |     sim_traj = []
234 | 
235 |     if dt_vec is None:
236 |         change_control_input = np.arange(0, t_horizon, t_horizon / (w_opt.shape[0]))[1:]
237 |         first_dt = t_horizon / w_opt.shape[0]
238 |         sim_traj.append(quad.get_state(quaternion=True, stacked=True))
239 |     else:
240 |         change_control_input = np.cumsum(dt_vec)
241 |         first_dt = dt_vec[0]
242 | 
243 |     t_start_ep = 1e-6
244 |     int_range = tqdm(np.arange(t_start_ep, total_sim_time, simulation_dt)) if progress_bar else \
245 |         np.arange(t_start_ep, total_sim_time, simulation_dt)
246 | 
247 |     current_ind = 0
248 |     past_ind = 0
249 |     for t_elapsed in int_range:
250 |         ref_u = w_opt[current_ind, :].T
251 |         simulate_func(ref_u)
252 |         if t_elapsed + simulation_dt >= first_dt:
253 |             current_ind = np.argwhere(change_control_input <= t_elapsed + simulation_dt)[-1, 0] + 1
254 |             if past_ind != current_ind:
255 |                 sim_traj.append(quad.get_state(quaternion=True, stacked=True))
256 |                 past_ind = current_ind
257 | 
258 |     if dt_vec is None:
259 |         sim_traj.append(quad.get_state(quaternion=True, stacked=True))
260 |     sim_traj = np.squeeze(sim_traj)
261 | 
262 |     quad.set_state(state_safe)
263 | 
264 |     return sim_traj
265 | 
266 | 
267 | def get_reference_chunk(reference_traj, reference_u, current_idx, n_mpc_nodes, reference_over_sampling):
268 |     """
269 |     Extracts the reference states and controls for the current MPC optimization given the over-sampled counterparts.
270 | 
271 |     :param reference_traj: The reference trajectory, which has been finely over-sampled by a factor of
272 |     reference_over_sampling. It should be a vector of shape (Nx13), where N is the length of the trajectory in samples.
273 |     :param reference_u: The reference controls, following the same requirements as reference_traj. Should be a vector
274 |     of shape (Nx4).
275 |     :param current_idx: Current index of the trajectory tracking. Should be an integer number between 0 and N-1.
276 |     :param n_mpc_nodes: Number of MPC nodes considered in the optimization.
277 |     :param reference_over_sampling: The over-sampling factor of the reference trajectories. Should be a positive
278 |     integer.
279 |     :return: Returns the chunks of reference selected for the current MPC iteration. Two numpy arrays will be returned:
280 |         - An ((N+1)x13) array, corresponding to the reference trajectory. The first row is the state of current_idx.
281 |         - An (Nx4) array, corresponding to the reference controls.
282 |     """
283 | 
284 |     # Dense references
285 |     ref_traj_chunk = reference_traj[current_idx:current_idx + (n_mpc_nodes + 1) * reference_over_sampling, :]
286 |     ref_u_chunk = reference_u[current_idx:current_idx + n_mpc_nodes * reference_over_sampling, :]
287 | 
288 |     # Indices for down-sampling the reference to number of MPC nodes
289 |     downsample_ref_ind = np.arange(0, min(reference_over_sampling * (n_mpc_nodes + 1), ref_traj_chunk.shape[0]),
290 |                                    reference_over_sampling, dtype=int)
291 | 
292 |     # Sparser references (same dt as node separation)
293 |     ref_traj_chunk = ref_traj_chunk[downsample_ref_ind, :]
294 |     ref_u_chunk = ref_u_chunk[downsample_ref_ind[:max(len(downsample_ref_ind) - 1, 1)], :]
295 | 
296 |     return ref_traj_chunk, ref_u_chunk
297 | 


--------------------------------------------------------------------------------
/ros_dd_mpc/src/utils/trajectory_generator.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | import matplotlib.pyplot as plt
  3 | 
  4 | 
  5 | def draw_poly(traj, u_traj, t, target_points=None, target_t=None):
  6 |     """
  7 |     Plots the generated trajectory of length n with the used keypoints.
  8 |     :param traj: Full generated reference trajectory. Numpy array of shape nx13
  9 |     :param u_traj: Generated reference inputs. Numpy array of shape nx4
 10 |     :param t: Timestamps of the references. Numpy array of length n
 11 |     :param target_points: m position keypoints used for trajectory generation. Numpy array of shape 3 x m.
 12 |     :param target_t: Timestamps of the reference position keypoints. If not passed, then they are extracted from the
 13 |     t vector, assuming constant time separation.
 14 |     """
 15 | 
 16 |     ders = 2
 17 |     dims = 3
 18 | 
 19 |     y_labels = [r'pos $[m]$', r'vel $[m/s]$', r'acc $[m/s^2]$', r'jer $[m/s^3]$']
 20 |     dim_legends = ['x', 'y', 'z']
 21 | 
 22 |     if target_t is None and target_points is not None:
 23 |         target_t = np.linspace(0, t[-1], target_points.shape[1])
 24 | 
 25 |     p_traj = traj[:, :3]
 26 |     a_traj = traj[:, 3:7]
 27 |     v_traj = traj[:, 7:10]
 28 |     r_traj = traj[:, 10:]
 29 | 
 30 |     plt_traj = [p_traj, v_traj]
 31 | 
 32 |     fig = plt.figure()
 33 |     for d_ord in range(ders):
 34 | 
 35 |         plt.subplot(ders + 2, 2, d_ord * 2 + 1)
 36 | 
 37 |         for dim in range(dims):
 38 | 
 39 |             plt.plot(t, plt_traj[d_ord][:, dim], label=dim_legends[dim])
 40 | 
 41 |             if d_ord == 0 and target_points is not None:
 42 |                 plt.plot(target_t, target_points[dim, :], 'bo')
 43 | 
 44 |         plt.gca().set_xticklabels([])
 45 |         plt.legend()
 46 |         plt.grid()
 47 |         plt.ylabel(y_labels[d_ord])
 48 | 
 49 |     dim_legends = [['w', 'x', 'y', 'z'], ['x', 'y', 'z']]
 50 |     y_labels = [r'att $[quat]$', r'rate $[rad/s]$']
 51 |     plt_traj = [a_traj, r_traj]
 52 |     for d_ord in range(ders):
 53 | 
 54 |         plt.subplot(ders + 2, 2, d_ord * 2 + 1 + ders * 2)
 55 |         for dim in range(plt_traj[d_ord].shape[1]):
 56 |             plt.plot(t, plt_traj[d_ord][:, dim], label=dim_legends[d_ord][dim])
 57 | 
 58 |         plt.legend()
 59 |         plt.grid()
 60 |         plt.ylabel(y_labels[d_ord])
 61 |         if d_ord == ders - 1:
 62 |             plt.xlabel(r'time $[s]$')
 63 |         else:
 64 |             plt.gca().set_xticklabels([])
 65 | 
 66 |     ax = fig.add_subplot(2, 2, 2, projection='3d')
 67 |     plt.plot(p_traj[:, 0], p_traj[:, 1], p_traj[:, 2])
 68 |     if target_points is not None:
 69 |         plt.plot(target_points[0, :], target_points[1, :], target_points[2, :], 'bo')
 70 |     plt.title('Target position trajectory')
 71 |     ax.set_xlabel(r'$p_x [m]$')
 72 |     ax.set_ylabel(r'$p_y [m]$')
 73 |     ax.set_zlabel(r'$p_z [m]$')
 74 | 
 75 |     plt.subplot(ders + 1, 2, (ders + 1) * 2)
 76 |     for i in range(u_traj.shape[1]):
 77 |         plt.plot(t, u_traj[:, i], label=r'$u_{}$'.format(i))
 78 |     plt.grid()
 79 |     plt.legend()
 80 |     plt.gca().yaxis.set_label_position("right")
 81 |     plt.gca().yaxis.tick_right()
 82 |     plt.xlabel(r'time $[s]$')
 83 |     plt.ylabel(r'single thrusts $[N]$')
 84 |     plt.title('Control inputs')
 85 | 
 86 |     plt.suptitle('Generated polynomial trajectory')
 87 | 
 88 |     plt.show()
 89 | 
 90 | 
 91 | def get_full_traj(poly_coeffs, target_dt, int_dt):
 92 | 
 93 |     dims = poly_coeffs.shape[-1]
 94 |     full_traj = np.zeros((4, dims, 0))
 95 |     t_total = np.zeros((0,))
 96 | 
 97 |     if isinstance(target_dt, float):
 98 |         # Adjust target_dt to make it divisible by int_dt
 99 |         target_dt = round(target_dt / int_dt) * int_dt
100 | 
101 |         # Assign target time for each keypoint using homogeneous spacing
102 |         t_vec = np.arange(0, target_dt * (poly_coeffs.shape[0] + 1) - 1e-5, target_dt)
103 | 
104 |     else:
105 |         # The time between each pair of points is assigned independently
106 |         # First, also adjust each value of the target_dt vector to make it divisible by int_dt
107 |         for i, dt in enumerate(target_dt):
108 |             target_dt[i] = round(dt / int_dt) * int_dt
109 | 
110 |         t_vec = np.append(np.zeros(1), np.cumsum(target_dt[:-1]))
111 | 
112 |     for seg in range(len(t_vec) - 1):
113 | 
114 |         # Select time sampling (linear or quadratic) mode
115 |         tau_dt = np.arange(t_vec[seg], t_vec[seg + 1] + 1e-5, int_dt)
116 | 
117 |         # Re-normalize time sampling vector between -1 and 1
118 |         t1 = (tau_dt - t_vec[seg]) / (t_vec[seg + 1] - t_vec[seg]) * 2 - 1
119 | 
120 |         # Compression ratio
121 |         compress = 2 / np.diff(t_vec)[seg]
122 | 
123 |         # Integrate current segment of trajectory
124 |         traj = np.zeros((4, dims, len(t1)))
125 | 
126 |         for der_order in range(4):
127 |             for i in range(dims):
128 |                 traj[der_order, i, :] = np.polyval(np.polyder(poly_coeffs[seg, :, i], der_order), t1) * (compress ** der_order)
129 | 
130 |         if seg < len(t_vec) - 2:
131 |             # Remove last sample (will be the initial point of next segment)
132 |             traj = traj[:, :, :-1]
133 |             t_seg = tau_dt[:-1]
134 |         else:
135 |             t_seg = tau_dt
136 | 
137 |         full_traj = np.concatenate((full_traj, traj), axis=-1)
138 |         t_total = np.concatenate((t_total, t_seg))
139 | 
140 |     # Separate into p_xyz and yaw trajectories
141 |     yaw_traj = full_traj[:, -1, :]
142 |     full_traj = full_traj[:, :-1, :]
143 | 
144 |     return full_traj, yaw_traj, t_total
145 | 
146 | 
147 | def fit_multi_segment_polynomial_trajectory(p_targets, yaw_targets):
148 | 
149 |     p_targets = np.concatenate((p_targets, yaw_targets[np.newaxis, :]), 0)
150 |     m = multiple_waypoints(p_targets.shape[1] - 1)
151 | 
152 |     dims = p_targets.shape[0]
153 |     n_segments = p_targets.shape[1]
154 | 
155 |     poly_coefficients = np.zeros((n_segments - 1, 8, dims))
156 |     for dim in range(dims):
157 |         b = rhs_generation(p_targets[dim, :])
158 |         poly_coefficients[:, :, dim] = np.fliplr(np.linalg.solve(m, b).reshape(n_segments - 1, 8))
159 | 
160 |     return poly_coefficients
161 | 
162 | 
163 | def matrix_generation(ts):
164 |     b = np.array([[1, ts,  ts**2, ts**3,    ts**4,    ts**5,     ts**6,     ts**7],
165 |                   [0, 1, 2*ts,  3*ts**2,  4*ts**3,  5*ts**4,   6*ts**5,   7*ts**6],
166 |                   [0, 0, 2,     6*ts,    12*ts**2, 20*ts**3,  30*ts**4,  42*ts**5],
167 |                   [0, 0, 0,     6,       24*ts,    60*ts**2, 120*ts**3, 210*ts**4],
168 |                   [0, 0, 0,     0,       24,      120*ts,    360*ts**2, 840*ts**3],
169 |                   [0, 0, 0,     0,       0,       120,       720*ts,   2520*ts**2],
170 |                   [0, 0, 0,     0,       0,       0,         720,      5040*ts],
171 |                   [0, 0, 0,     0,       0,       0,         0,        5040]])
172 | 
173 |     return b
174 | 
175 | 
176 | def multiple_waypoints(n_segments):
177 | 
178 |     m = np.zeros((8 * n_segments, 8 * n_segments))
179 | 
180 |     for i in range(n_segments):
181 | 
182 |         if i == 0:
183 | 
184 |             # initial condition of the first curve
185 |             b = matrix_generation(-1.0)
186 |             m[8 * i:8 * i + 4, 8 * i:8 * i + 8] = b[:4, :]
187 | 
188 |             # intermediary condition of the first curve
189 |             b = matrix_generation(1.0)
190 |             m[8 * i + 4:8 * i + 7 + 4, 8 * i:8 * i + 8] = b[:-1, :]
191 | 
192 |             # starting condition of the second curve position and derivatives
193 |             b = matrix_generation(-1.0)
194 |             m[8 * i + 4 + 1:8 * i + 4 + 7, 8 * (i + 1):8 * (i + 1) + 8] = -b[1:-1, :]
195 |             m[8 * i + 4 + 7:8 * i + 4 + 8, 8 * (i + 1):8 * (i + 1) + 8] = b[0, :]
196 | 
197 |         elif i != n_segments - 1:
198 | 
199 |             # starting condition of the ith curve position and derivatives
200 |             b = matrix_generation(1.0)
201 |             m[8 * i + 4:8 * i + 7 + 4, 8 * i:8 * i + 8] = b[:-1, :]
202 | 
203 |             # end condition of the ith curve position and derivatives
204 |             b = matrix_generation(-1.0)
205 |             m[8 * i + 4 + 1:8 * i + 4 + 7, 8 * (i + 1):8 * (i + 1) + 8] = -b[1:-1, :]
206 |             m[8 * i + 4 + 7:8 * i + 4 + 8, 8 * (i + 1):8 * (i + 1) + 8] = b[0, :]
207 | 
208 |         if i == n_segments - 1:
209 |             # end condition of the final curve position and derivatives (4 boundary conditions)
210 |             b = matrix_generation(1.0)
211 |             m[8 * i + 4:8 * i + 4 + 4, 8 * i:8 * i + 8] = b[:4, :]
212 | 
213 |     return m
214 | 
215 | 
216 | 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):
217 | 
218 |     if v_start is None:
219 |         v_start = np.array([0, 0])
220 |     if v_end is None:
221 |         v_end = np.array([0, 0])
222 |     if a_start is None:
223 |         a_start = np.array([0, 0])
224 |     if a_end is None:
225 |         a_end = np.array([0, 0])
226 |     if j_start is None:
227 |         j_start = np.array([0, 0])
228 |     if j_end is None:
229 |         j_end = np.array([0, 0])
230 | 
231 |     poly_coefficients = np.zeros((8, len(p_start)))
232 | 
233 |     tf = 1
234 |     ti = -1
235 |     A = np.array(([
236 |         [1 * tf ** 7,   1 * tf ** 6,   1 * tf ** 5,   1 * tf ** 4,   1 * tf ** 3,  1 * tf ** 2,  1 * tf ** 1,  1],
237 |         [7 * tf ** 6,   6 * tf ** 5,   5 * tf ** 4,   4 * tf ** 3,   3 * tf ** 2,  2 * tf ** 1,  1,            0],
238 |         [42 * tf ** 5,  30 * tf ** 4,  20 * tf ** 3,  12 * tf ** 2,  6 * tf ** 1,  2,            0,            0],
239 |         [210 * tf ** 4, 120 * tf ** 3, 60 * tf ** 2,  24 * tf ** 1,  6,            0,            0,            0],
240 |         [1 * ti ** 7,   1 * ti ** 6,   1 * ti ** 5,   1 * ti ** 4,   1 * ti ** 3,  1 * ti ** 2,  1 * ti ** 1,  1],
241 |         [7 * ti ** 6,   6 * ti ** 5,   5 * ti ** 4,   4 * ti ** 3,   3 * ti ** 2,  2 * ti ** 1,  1,            0],
242 |         [42 * ti ** 5,  30 * ti ** 4,  20 * ti ** 3,  12 * ti ** 2,  6 * ti ** 1,  2,            0,            0],
243 |         [210 * ti ** 4, 120 * ti ** 3, 60 * ti ** 2,  24 * ti ** 1,  6,            0,            0,            0]]))
244 | 
245 |     A = np.tile(A[:, :, np.newaxis], (1, 1, len(p_start)))
246 | 
247 |     b = np.concatenate((p_end, v_end, a_end, j_end, p_start, v_start, a_start, j_start)).reshape(8, -1)
248 | 
249 |     for i in range(len(p_start)):
250 |         poly_coefficients[:, i] = np.linalg.inv(A[:, :, i]).dot(np.array(b[:, i]))
251 | 
252 |     return np.expand_dims(poly_coefficients, 0)
253 | 
254 | 
255 | def rhs_generation(x):
256 |     n = x.shape[0] - 1
257 | 
258 |     big_x = np.zeros((8 * n))
259 |     big_x[:4] = np.array([x[0], 0, 0, 0]).T
260 |     big_x[-4:] = np.array([x[-1], 0, 0, 0]).T
261 | 
262 |     for i in range(1, n):
263 |         big_x[8 * (i - 1) + 4:8 * (i - 1) + 8 + 4] = np.array([x[i], 0, 0, 0, 0, 0, 0, x[i]]).T
264 | 
265 |     return big_x
266 | 


--------------------------------------------------------------------------------