├── .gitignore ├── LICENSE ├── README.md ├── __init__.py ├── mpc_kuka_avoid.py ├── mpc_kuka_contact.py ├── mpc_kuka_reaching.py ├── mpc_utils.py ├── ocp_kuka_avoid.py ├── ocp_kuka_contact.py ├── ocp_kuka_obstacle.py ├── ocp_kuka_reaching.py ├── ocp_utils.py ├── pin_utils.py └── test_riccati_gains_mpc.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2023, Machines in Motion Laboratory 4 | 5 | Redistribution and use in source and binary forms, with or without 6 | modification, are permitted provided that the following conditions are met: 7 | 8 | 1. Redistributions of source code must retain the above copyright notice, this 9 | list of conditions and the following disclaimer. 10 | 11 | 2. Redistributions in binary form must reproduce the above copyright notice, 12 | this list of conditions and the following disclaimer in the documentation 13 | and/or other materials provided with the distribution. 14 | 15 | 3. Neither the name of the copyright holder nor the names of its 16 | contributors may be used to endorse or promote products derived from 17 | this software without specific prior written permission. 18 | 19 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 20 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 21 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 22 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 23 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 24 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 25 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 26 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 27 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 28 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # minimal_examples_crocoddyl 2 | Demo scripts to quickly get started with Crocoddyl and MPC. 3 | 4 | An easy way to integrate this package is to clone this repository into your project and import ocp_utils and pin_utils for your use. 5 | 6 | 7 | # Dependencies 8 | For OCP and MPC scripts 9 | - [mim_robots](https://github.com/machines-in-motion/mim_robots) 10 | - [Crocoddyl](https://github.com/loco-3d/crocoddyl) 11 | - [matplotlib](https://matplotlib.org/) 12 | 13 | For visualization (optional) 14 | - [meshcat](https://github.com/meshcat-dev/meshcat) 15 | - [gepetto-viewer](https://github.com/Gepetto/gepetto-viewer) 16 | 17 | # Usage 18 | For the reaching task, run `python ocp_kuka_reaching.py` to solve the OCP and visualize / plot the solution. Run `python mpc_kuka_reaching.py` to simulate it in MPC in PyBullet. Same for contact task. 19 | 20 | The scripts are minimal and self-explanatory. The machinery for data extraction and plotting is hidden in the utils. 21 | 22 | # Copyrights 23 | Copyright(c) 2019-2023 New York University 24 | 25 | # License 26 | BSD 3-Clause License 27 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | import pin_utils 2 | import ocp_utils -------------------------------------------------------------------------------- /mpc_kuka_avoid.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example script : MPC simulation with KUKA arm 3 | static target reaching task while avoiding an obstacle. 4 | Based from the mpc_kuka_reaching.py code. 5 | 6 | Written by Arthur Haffemayer. 7 | """ 8 | 9 | import crocoddyl 10 | import mim_solvers 11 | import numpy as np 12 | import pinocchio as pin 13 | 14 | np.set_printoptions(precision=4, linewidth=180) 15 | 16 | import pin_utils, mpc_utils 17 | 18 | from mim_robots.pybullet.env import BulletEnvWithGround 19 | from mim_robots.robot_loader import load_bullet_wrapper 20 | import hppfcl 21 | 22 | try: 23 | from colmpc import ResidualDistanceCollision 24 | except: 25 | print( 26 | """It needs to be built with COLMPC https://github.com/agimus-project/colmpc 27 | and HPPFCL the devel branch for it to work.""" 28 | ) 29 | 30 | import pybullet as p 31 | 32 | 33 | # # # # # # # # # # # # # # # # # # # 34 | ### LOAD ROBOT MODEL and SIMU ENV ### 35 | # # # # # # # # # # # # # # # # # # # 36 | # Simulation environment 37 | env = BulletEnvWithGround(p.GUI, dt=1e-3) 38 | # Robot simulator 39 | robot_simulator = load_bullet_wrapper("iiwa_convex") 40 | env.add_robot(robot_simulator) 41 | 42 | # Extract robot model 43 | nq = robot_simulator.pin_robot.model.nq 44 | nv = robot_simulator.pin_robot.model.nv 45 | nu = nq 46 | nx = nq + nv 47 | q0 = np.array([0.1, 0.7, 0.0, 0.7, -0.5, 1.5, 0.0]) 48 | v0 = np.zeros(nv) 49 | x0 = np.concatenate([q0, v0]) 50 | # Add robot to simulation and initialize 51 | env.add_robot(robot_simulator) 52 | robot_simulator.reset_state(q0, v0) 53 | robot_simulator.forward_robot(q0, v0) 54 | print("[PyBullet] Created robot (id = " + str(robot_simulator.robotId) + ")") 55 | 56 | robot_simulator.pin_robot.collision_model = pin_utils.transform_model_into_capsules( 57 | robot_simulator.pin_robot.collision_model 58 | ) 59 | # # # # # # # # # # # # 60 | ### ADDING OBSTACLE ### 61 | # # # # # # # # # # # # 62 | 63 | OBSTACLE_POSE = pin.SE3(pin.utils.rotate("x", np.pi), np.array([-0.2, 0.15, 0.7])) 64 | OBSTACLE = hppfcl.Sphere(1e-1) 65 | OBSTACLE_GEOM_OBJECT = pin.GeometryObject( 66 | "obstacle", 67 | 0, 68 | 0, 69 | OBSTACLE, 70 | OBSTACLE_POSE, 71 | ) 72 | ID_OBSTACLE = robot_simulator.pin_robot.collision_model.addGeometryObject( 73 | OBSTACLE_GEOM_OBJECT 74 | ) 75 | 76 | shapes_avoiding_collision_with_obstacles = [ 77 | "L3_capsule_0", 78 | "L4_capsule_0", 79 | "L4_capsule_1", 80 | "L5_capsule_0", 81 | "L6_capsule_0", 82 | "L6_capsule_1", 83 | "L7_capsule_0", 84 | ] 85 | 86 | # Adding the collisions pairs to the collision model 87 | for shape in shapes_avoiding_collision_with_obstacles: 88 | robot_simulator.pin_robot.collision_model.addCollisionPair( 89 | pin.CollisionPair( 90 | robot_simulator.pin_robot.collision_model.getGeometryId(shape), 91 | robot_simulator.pin_robot.collision_model.getGeometryId("obstacle"), 92 | ) 93 | ) 94 | 95 | 96 | # # # # # # # # # # # # # # # 97 | ### SETUP CROCODDYL OCP ### 98 | # # # # # # # # # # # # # # # 99 | # State and actuation model 100 | state = crocoddyl.StateMultibody(robot_simulator.pin_robot.model) 101 | actuation = crocoddyl.ActuationModelFull(state) 102 | # Running and terminal cost models 103 | runningCostModel = crocoddyl.CostModelSum(state) 104 | terminalCostModel = crocoddyl.CostModelSum(state) 105 | # Constraint model managers 106 | runningConstraintModelManager = crocoddyl.ConstraintModelManager(state, actuation.nu) 107 | terminalConstraintModelManager = crocoddyl.ConstraintModelManager(state, actuation.nu) 108 | 109 | # Create collision avoidance constraints 110 | 111 | # Creating the residual 112 | if len(robot_simulator.pin_robot.collision_model.collisionPairs) != 0: 113 | for col_idx in range(len(robot_simulator.pin_robot.collision_model.collisionPairs)): 114 | # obstacleDistanceResidual = ResidualCollision( 115 | # state, robot_simulator.pin_robot.collision_model, cdata, col_idx 116 | # ) 117 | obstacleDistanceResidual = ResidualDistanceCollision(state, 7, robot_simulator.pin_robot.collision_model, col_idx) 118 | 119 | # Creating the inequality constraint 120 | constraint = crocoddyl.ConstraintModelResidual( 121 | state, 122 | obstacleDistanceResidual, 123 | np.array([2e-2]), 124 | np.array([np.inf]), 125 | ) 126 | 127 | # Adding the constraint to the constraint manager 128 | runningConstraintModelManager.addConstraint("col_" + str(col_idx), constraint) 129 | terminalConstraintModelManager.addConstraint("col_term_" + str(col_idx), constraint) 130 | 131 | # Create cost terms 132 | # Control regularization cost 133 | uResidual = crocoddyl.ResidualModelControlGrav(state) 134 | uRegCost = crocoddyl.CostModelResidual(state, uResidual) 135 | # State regularization cost 136 | xResidual = crocoddyl.ResidualModelState(state, x0) 137 | xRegCost = crocoddyl.CostModelResidual(state, xResidual) 138 | # endeff frame translation cost 139 | endeff_frame_id = robot_simulator.pin_robot.model.getFrameId("contact") 140 | # endeff_translation = robot.data.oMf[endeff_frame_id].translation.copy() 141 | endeff_translation = np.array( 142 | [-0.4, 0.3, 0.7] 143 | ) # move endeff +10 cm along x in WORLD frame 144 | frameTranslationResidual = crocoddyl.ResidualModelFrameTranslation( 145 | state, endeff_frame_id, endeff_translation 146 | ) 147 | frameTranslationCost = crocoddyl.CostModelResidual(state, frameTranslationResidual) 148 | # Add costs 149 | runningCostModel.addCost("stateReg", xRegCost, 1e-1) 150 | runningCostModel.addCost("ctrlRegGrav", uRegCost, 1e-4) 151 | runningCostModel.addCost("translation", frameTranslationCost, 100) 152 | terminalCostModel.addCost("stateReg", xRegCost, 1e-1) 153 | terminalCostModel.addCost("translation", frameTranslationCost, 100) 154 | # Create Differential Action Model (DAM), i.e. continuous dynamics and cost functions 155 | running_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics( 156 | state, actuation, runningCostModel, runningConstraintModelManager 157 | ) 158 | terminal_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics( 159 | state, actuation, terminalCostModel, terminalConstraintModelManager 160 | ) 161 | # Create Integrated Action Model (IAM), i.e. Euler integration of continuous dynamics and cost 162 | dt = 1e-2 163 | runningModel = crocoddyl.IntegratedActionModelEuler(running_DAM, dt) 164 | terminalModel = crocoddyl.IntegratedActionModelEuler(terminal_DAM, 0.0) 165 | # Create the shooting problem 166 | T = 100 167 | problem = crocoddyl.ShootingProblem(x0, [runningModel] * T, terminalModel) 168 | # Create solver + callbacks 169 | solver = mim_solvers.SolverCSQP(problem) 170 | # solver.setCallbacks([crocoddyl.CallbackLogger(), 171 | # crocoddyl.CallbackVerbose()]) 172 | # Warm start : initial state + gravity compensation 173 | xs_init = [x0 for i in range(T + 1)] 174 | us_init = solver.problem.quasiStatic(xs_init[:-1]) 175 | # Solve 176 | solver.termination_tolerance = 1e-4 177 | solver.with_callbacks = True 178 | solver.solve(xs_init, us_init, 100) 179 | solver.with_callbacks = False 180 | solver.max_qp_iters = 50 181 | # # # # # # # # # # # # 182 | ### MPC SIMULATION ### 183 | # # # # # # # # # # # # 184 | # OCP parameters 185 | ocp_params = {} 186 | ocp_params["N_h"] = T 187 | ocp_params["dt"] = dt 188 | ocp_params["maxiter"] = 3 189 | ocp_params["pin_model"] = robot_simulator.pin_robot.model 190 | ocp_params["armature"] = runningModel.differential.armature 191 | ocp_params["id_endeff"] = endeff_frame_id 192 | ocp_params["active_costs"] = solver.problem.runningModels[ 193 | 0 194 | ].differential.costs.active.tolist() 195 | # Simu parameters 196 | sim_params = {} 197 | sim_params["sim_freq"] = int(1.0 / env.dt) 198 | sim_params["mpc_freq"] = 1000 199 | sim_params["T_sim"] = 2.0 200 | log_rate = 100 201 | # Initialize simulation data 202 | sim_data = mpc_utils.init_sim_data(sim_params, ocp_params, x0) 203 | # Display target 204 | mpc_utils.display_ball(endeff_translation, RADIUS=0.025, COLOR=[1.0, 0.0, 0.0, 0.6]) 205 | mpc_utils.display_ball( 206 | OBSTACLE_POSE.translation, RADIUS=OBSTACLE.radius, COLOR=[1.0, 0.0, 0.0, 0.6] 207 | ) 208 | 209 | # Simulate 210 | mpc_cycle = 0 211 | for i in range(sim_data["N_sim"]): 212 | 213 | if i % log_rate == 0: 214 | print("\n SIMU step " + str(i) + "/" + str(sim_data["N_sim"]) + "\n") 215 | 216 | # Solve OCP if we are in a planning cycle (MPC/planning frequency) 217 | if i % int(sim_params["sim_freq"] / sim_params["mpc_freq"]) == 0: 218 | # Set x0 to measured state 219 | solver.problem.x0 = sim_data["state_mea_SIM_RATE"][i, :] 220 | # Warm start using previous solution 221 | xs_init = list(solver.xs[1:]) + [solver.xs[-1]] 222 | xs_init[0] = sim_data["state_mea_SIM_RATE"][i, :] 223 | us_init = list(solver.us[1:]) + [solver.us[-1]] 224 | 225 | # Solve OCP & record MPC predictions 226 | solver.solve(xs_init, us_init, ocp_params["maxiter"]) 227 | sim_data["state_pred"][mpc_cycle, :, :] = np.array(solver.xs) 228 | sim_data["ctrl_pred"][mpc_cycle, :, :] = np.array(solver.us) 229 | # Extract relevant predictions for interpolations 230 | x_curr = sim_data["state_pred"][ 231 | mpc_cycle, 0, : 232 | ] # x0* = measured state (q^, v^ ) 233 | x_pred = sim_data["state_pred"][ 234 | mpc_cycle, 1, : 235 | ] # x1* = predicted state (q1*, v1*) 236 | u_curr = sim_data["ctrl_pred"][ 237 | mpc_cycle, 0, : 238 | ] # u0* = optimal control (tau0*) 239 | # Record costs references 240 | q = sim_data["state_pred"][mpc_cycle, 0, : sim_data["nq"]] 241 | sim_data["ctrl_ref"][mpc_cycle, :] = pin_utils.get_u_grav( 242 | q, 243 | solver.problem.runningModels[0].differential.pinocchio, 244 | ocp_params["armature"], 245 | ) 246 | sim_data["state_ref"][mpc_cycle, :] = ( 247 | solver.problem.runningModels[0] 248 | .differential.costs.costs["stateReg"] 249 | .cost.residual.reference 250 | ) 251 | sim_data["lin_pos_ee_ref"][mpc_cycle, :] = ( 252 | solver.problem.runningModels[0] 253 | .differential.costs.costs["translation"] 254 | .cost.residual.reference 255 | ) 256 | 257 | # Select reference control and state for the current MPC cycle 258 | x_ref_MPC_RATE = x_curr + sim_data["ocp_to_mpc_ratio"] * (x_pred - x_curr) 259 | u_ref_MPC_RATE = u_curr 260 | if mpc_cycle == 0: 261 | sim_data["state_des_MPC_RATE"][mpc_cycle, :] = x_curr 262 | sim_data["ctrl_des_MPC_RATE"][mpc_cycle, :] = u_ref_MPC_RATE 263 | sim_data["state_des_MPC_RATE"][mpc_cycle + 1, :] = x_ref_MPC_RATE 264 | 265 | # Increment planning counter 266 | mpc_cycle += 1 267 | 268 | # Select reference control and state for the current SIMU cycle 269 | x_ref_SIM_RATE = x_curr + sim_data["ocp_to_mpc_ratio"] * (x_pred - x_curr) 270 | u_ref_SIM_RATE = u_curr 271 | 272 | # First prediction = measurement = initialization of MPC 273 | if i == 0: 274 | sim_data["state_des_SIM_RATE"][i, :] = x_curr 275 | sim_data["ctrl_des_SIM_RATE"][i, :] = u_ref_SIM_RATE 276 | sim_data["state_des_SIM_RATE"][i + 1, :] = x_ref_SIM_RATE 277 | 278 | # Send torque to simulator & step simulator 279 | robot_simulator.send_joint_command(u_ref_SIM_RATE) 280 | env.step() 281 | # Measure new state from simulator 282 | q_mea_SIM_RATE, v_mea_SIM_RATE = robot_simulator.get_state() 283 | # Update pinocchio model 284 | robot_simulator.forward_robot(q_mea_SIM_RATE, v_mea_SIM_RATE) 285 | # Record data 286 | x_mea_SIM_RATE = np.concatenate([q_mea_SIM_RATE, v_mea_SIM_RATE]).T 287 | sim_data["state_mea_SIM_RATE"][i + 1, :] = x_mea_SIM_RATE 288 | 289 | 290 | plot_data = mpc_utils.extract_plot_data_from_sim_data(sim_data) 291 | 292 | mpc_utils.plot_mpc_results( 293 | plot_data, 294 | which_plots=["all"], 295 | PLOT_PREDICTIONS=True, 296 | pred_plot_sampling=int(sim_params["mpc_freq"] / 10), 297 | ) 298 | -------------------------------------------------------------------------------- /mpc_kuka_contact.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Example script : MPC simulation with KUKA arm 3 | contact force task 4 | ''' 5 | 6 | import crocoddyl 7 | import mim_solvers 8 | import numpy as np 9 | import pinocchio as pin 10 | np.set_printoptions(precision=4, linewidth=180) 11 | import pin_utils, mpc_utils 12 | 13 | from mim_robots.pybullet.env import BulletEnvWithGround 14 | from mim_robots.robot_loader import load_bullet_wrapper 15 | import pybullet as p 16 | import pinocchio as pin 17 | 18 | # # # # # # # # # # # # # # # # # # # 19 | ### LOAD ROBOT MODEL and SIMU ENV ###  20 | # # # # # # # # # # # # # # # # # # #  21 | # Simulation environment 22 | env = BulletEnvWithGround(p.GUI, dt=1e-3) 23 | # Robot simulator 24 | robot_simulator = load_bullet_wrapper("iiwa") 25 | # Extract robot model 26 | nq = robot_simulator.pin_robot.model.nq 27 | nv = robot_simulator.pin_robot.model.nv 28 | nu = nq; nx = nq+nv 29 | q0 = np.array([0.1, 0.7, 0., 0.7, -0.5, 1.5, 0.]) 30 | v0 = np.zeros(nv) 31 | x0 = np.concatenate([q0, v0]) 32 | # Add robot to simulation and initialize 33 | env.add_robot(robot_simulator) 34 | robot_simulator.reset_state(q0, v0) 35 | robot_simulator.forward_robot(q0, v0) 36 | print("[PyBullet] Created robot (id = "+str(robot_simulator.robotId)+")") 37 | 38 | # Display contact surface 39 | contact_frame_id = robot_simulator.pin_robot.model.getFrameId("contact") 40 | contact_frame_placement = robot_simulator.pin_robot.data.oMf[contact_frame_id] 41 | offset = 0.03348 42 | contact_frame_placement.translation = contact_frame_placement.act(np.array([0., 0., offset])) 43 | mpc_utils.display_contact_surface(contact_frame_placement, with_collision=True) 44 | 45 | # # # # # # # # # # # # # # # 46 | ### SETUP CROCODDYL OCP ### 47 | # # # # # # # # # # # # # # # 48 | # State and actuation model 49 | state = crocoddyl.StateMultibody(robot_simulator.pin_robot.model) 50 | actuation = crocoddyl.ActuationModelFull(state) 51 | # Running and terminal cost models 52 | runningCostModel = crocoddyl.CostModelSum(state) 53 | terminalCostModel = crocoddyl.CostModelSum(state) 54 | # Contact model 55 | contactModel = crocoddyl.ContactModelMultiple(state, actuation.nu) 56 | # Create 3D contact on the en-effector frame 57 | contact_position = robot_simulator.pin_robot.data.oMf[contact_frame_id].copy() 58 | baumgarte_gains = np.array([0., 50.]) 59 | pinRef = pin.LOCAL_WORLD_ALIGNED 60 | contact3d = crocoddyl.ContactModel6D(state, contact_frame_id, contact_position, pinRef, baumgarte_gains) 61 | # Populate contact model with contacts 62 | contactModel.addContact("contact", contact3d, active=True) 63 | # Create cost terms 64 | # Control regularization cost 65 | uResidual = crocoddyl.ResidualModelContactControlGrav(state) 66 | uRegCost = crocoddyl.CostModelResidual(state, uResidual) 67 | # State regularization cost 68 | xResidual = crocoddyl.ResidualModelState(state, x0) 69 | xRegCost = crocoddyl.CostModelResidual(state, xResidual) 70 | # End-effector frame force cost 71 | desired_wrench = np.array([0., 0., -20., 0., 0., 0.]) 72 | frameForceResidual = crocoddyl.ResidualModelContactForce(state, contact_frame_id, pin.Force(desired_wrench), 6, actuation.nu) 73 | contactForceCost = crocoddyl.CostModelResidual(state, frameForceResidual) 74 | # Populate cost models with cost terms 75 | runningCostModel.addCost("stateReg", xRegCost, 1e-2) 76 | runningCostModel.addCost("ctrlRegGrav", uRegCost, 1e-4) 77 | runningCostModel.addCost("force", contactForceCost, 10.) 78 | terminalCostModel.addCost("stateReg", xRegCost, 1e-2) 79 | # Create Differential Action Model (DAM), i.e. continuous dynamics and cost functions 80 | running_DAM = crocoddyl.DifferentialActionModelContactFwdDynamics(state, actuation, contactModel, runningCostModel, inv_damping=0., enable_force=True) 81 | terminal_DAM = crocoddyl.DifferentialActionModelContactFwdDynamics(state, actuation, contactModel, terminalCostModel, inv_damping=0., enable_force=True) 82 | # Create Integrated Action Model (IAM), i.e. Euler integration of continuous dynamics and cost 83 | dt = 1e-2 84 | runningModel = crocoddyl.IntegratedActionModelEuler(running_DAM, dt) 85 | terminalModel = crocoddyl.IntegratedActionModelEuler(terminal_DAM, 0.) 86 | # Create the shooting problem 87 | T = 100 88 | problem = crocoddyl.ShootingProblem(x0, [runningModel] * T, terminalModel) 89 | # Create solver + callbacks 90 | solver = mim_solvers.SolverSQP(problem) 91 | # solver.setCallbacks([crocoddyl.CallbackLogger(), 92 | # crocoddyl.CallbackVerbose()]) 93 | # Warm start : initial state + gravity compensation 94 | xs_init = [x0 for i in range(T+1)] 95 | us_init = solver.problem.quasiStatic(xs_init[:-1]) 96 | # Solve 97 | solver.termination_tolerance = 1e-4 98 | solver.with_callbacks = True 99 | solver.solve(xs_init, us_init, 100) 100 | solver.with_callbacks = False 101 | 102 | 103 | # # # # # # # # # # # # 104 | ### MPC SIMULATION ### 105 | # # # # # # # # # # # # 106 | # OCP parameters 107 | ocp_params = {} 108 | ocp_params['N_h'] = T 109 | ocp_params['dt'] = dt 110 | ocp_params['maxiter'] = 10 111 | ocp_params['pin_model'] = robot_simulator.pin_robot.model 112 | ocp_params['armature'] = runningModel.differential.armature 113 | ocp_params['id_endeff'] = contact_frame_id 114 | ocp_params['active_costs'] = solver.problem.runningModels[0].differential.costs.active.tolist() 115 | 116 | # Simu parameters 117 | sim_params = {} 118 | sim_params['sim_freq'] = int(1./env.dt) 119 | sim_params['mpc_freq'] = 1000 120 | sim_params['T_sim'] = 2. 121 | log_rate = 100 122 | # Initialize simulation data 123 | sim_data = mpc_utils.init_sim_data(sim_params, ocp_params, x0) 124 | # Simulate 125 | mpc_cycle = 0 126 | for i in range(sim_data['N_sim']): 127 | 128 | if(i%log_rate==0): 129 | print("\n SIMU step "+str(i)+"/"+str(sim_data['N_sim'])+"\n") 130 | 131 | # Solve OCP if we are in a planning cycle (MPC/planning frequency) 132 | if(i%int(sim_params['sim_freq']/sim_params['mpc_freq']) == 0): 133 | # Set x0 to measured state 134 | solver.problem.x0 = sim_data['state_mea_SIM_RATE'][i, :] 135 | # Warm start using previous solution 136 | xs_init = list(solver.xs[1:]) + [solver.xs[-1]] 137 | xs_init[0] = sim_data['state_mea_SIM_RATE'][i, :] 138 | us_init = list(solver.us[1:]) + [solver.us[-1]] 139 | 140 | # Solve OCP & record MPC predictions 141 | solver.solve(xs_init, us_init, ocp_params['maxiter']) 142 | sim_data['state_pred'][mpc_cycle, :, :] = np.array(solver.xs) 143 | sim_data['ctrl_pred'][mpc_cycle, :, :] = np.array(solver.us) 144 | sim_data ['force_pred'][mpc_cycle, :, :] = np.array([solver.problem.runningDatas[i].differential.multibody.contacts.contacts['contact'].f.vector for i in range(ocp_params['N_h'])]) 145 | # Extract relevant predictions for interpolations 146 | x_curr = sim_data['state_pred'][mpc_cycle, 0, :] # x0* = measured state (q^, v^ ) 147 | x_pred = sim_data['state_pred'][mpc_cycle, 1, :] # x1* = predicted state (q1*, v1*) 148 | u_curr = sim_data['ctrl_pred'][mpc_cycle, 0, :] # u0* = optimal control (tau0*) 149 | f_curr = sim_data['force_pred'][mpc_cycle, 0, :] 150 | f_pred = sim_data['force_pred'][mpc_cycle, 1, :] 151 | # Record costs references 152 | q = sim_data['state_pred'][mpc_cycle, 0, :sim_data['nq']] 153 | sim_data['ctrl_ref'][mpc_cycle, :] = pin_utils.get_u_grav(q, solver.problem.runningModels[0].differential.pinocchio, ocp_params['armature']) 154 | sim_data['f_ee_ref'][mpc_cycle, :] = solver.problem.runningModels[0].differential.costs.costs['force'].cost.residual.reference.vector 155 | sim_data['state_ref'][mpc_cycle, :] = solver.problem.runningModels[0].differential.costs.costs['stateReg'].cost.residual.reference 156 | 157 | 158 | # Select reference control and state for the current MPC cycle 159 | x_ref_MPC_RATE = x_curr + sim_data['ocp_to_mpc_ratio'] * (x_pred - x_curr) 160 | u_ref_MPC_RATE = u_curr 161 | f_ref_MPC_RATE = f_curr + sim_data['ocp_to_mpc_ratio'] * (f_pred - f_curr) 162 | if(mpc_cycle==0): 163 | sim_data['state_des_MPC_RATE'][mpc_cycle, :] = x_curr 164 | sim_data['ctrl_des_MPC_RATE'][mpc_cycle, :] = u_ref_MPC_RATE 165 | sim_data['state_des_MPC_RATE'][mpc_cycle+1, :] = x_ref_MPC_RATE 166 | sim_data['force_des_MPC_RATE'][mpc_cycle, :] = f_ref_MPC_RATE 167 | 168 | # Increment planning counter 169 | mpc_cycle += 1 170 | 171 | 172 | # Select reference control and state for the current SIMU cycle 173 | x_ref_SIM_RATE = x_curr + sim_data['ocp_to_mpc_ratio'] * (x_pred - x_curr) 174 | u_ref_SIM_RATE = u_curr 175 | f_ref_SIM_RATE = f_curr + sim_data['ocp_to_mpc_ratio'] * (f_pred - f_curr) 176 | 177 | # First prediction = measurement = initialization of MPC 178 | if(i==0): 179 | sim_data['state_des_SIM_RATE'][i, :] = x_curr 180 | sim_data['ctrl_des_SIM_RATE'][i, :] = u_ref_SIM_RATE 181 | sim_data['state_des_SIM_RATE'][i+1, :] = x_ref_SIM_RATE 182 | sim_data['force_des_SIM_RATE'][i, :] = f_ref_SIM_RATE 183 | 184 | #  Send output of actuation torque to the RBD simulator 185 | robot_simulator.send_joint_command(u_ref_SIM_RATE) 186 | env.step() 187 | # Measure new state from simulation 188 | q_mea_SIM_RATE, v_mea_SIM_RATE = robot_simulator.get_state() 189 | # Update pinocchio model 190 | robot_simulator.forward_robot(q_mea_SIM_RATE, v_mea_SIM_RATE) 191 | f_mea_SIM_RATE = mpc_utils.get_contact_wrench(robot_simulator, sim_data['id_endeff']) 192 | # Record data (unnoised) 193 | x_mea_SIM_RATE = np.concatenate([q_mea_SIM_RATE, v_mea_SIM_RATE]).T 194 | sim_data['state_mea_SIM_RATE'][i+1, :] = x_mea_SIM_RATE 195 | sim_data['force_mea_SIM_RATE'][i, :] = f_mea_SIM_RATE 196 | 197 | plot_data = mpc_utils.extract_plot_data_from_sim_data(sim_data) 198 | 199 | mpc_utils.plot_mpc_results(plot_data, which_plots=['all'], PLOT_PREDICTIONS=True, pred_plot_sampling=int(sim_params['mpc_freq']/100)) 200 | -------------------------------------------------------------------------------- /mpc_kuka_reaching.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Example script : MPC simulation with KUKA arm 3 | static target reaching task 4 | ''' 5 | 6 | import crocoddyl 7 | import mim_solvers 8 | import numpy as np 9 | np.set_printoptions(precision=4, linewidth=180) 10 | 11 | import pin_utils, mpc_utils 12 | 13 | from mim_robots.pybullet.env import BulletEnvWithGround 14 | from mim_robots.robot_loader import load_bullet_wrapper 15 | import pybullet as p 16 | 17 | 18 | # # # # # # # # # # # # # # # # # # # 19 | ### LOAD ROBOT MODEL and SIMU ENV ###  20 | # # # # # # # # # # # # # # # # # # #  21 | # Simulation environment 22 | env = BulletEnvWithGround(p.GUI, dt=1e-3) 23 | # Robot simulator 24 | robot_simulator = load_bullet_wrapper("iiwa") 25 | env.add_robot(robot_simulator) 26 | 27 | # Extract robot model 28 | nq = robot_simulator.pin_robot.model.nq 29 | nv = robot_simulator.pin_robot.model.nv 30 | nu = nq; nx = nq+nv 31 | q0 = np.array([0.1, 0.7, 0., 0.7, -0.5, 1.5, 0.]) 32 | v0 = np.zeros(nv) 33 | x0 = np.concatenate([q0, v0]) 34 | # Add robot to simulation and initialize 35 | env.add_robot(robot_simulator) 36 | robot_simulator.reset_state(q0, v0) 37 | robot_simulator.forward_robot(q0, v0) 38 | print("[PyBullet] Created robot (id = "+str(robot_simulator.robotId)+")") 39 | 40 | 41 | # # # # # # # # # # # # # # # 42 | ### SETUP CROCODDYL OCP ### 43 | # # # # # # # # # # # # # # # 44 | # State and actuation model 45 | state = crocoddyl.StateMultibody(robot_simulator.pin_robot.model) 46 | actuation = crocoddyl.ActuationModelFull(state) 47 | # Running and terminal cost models 48 | runningCostModel = crocoddyl.CostModelSum(state) 49 | terminalCostModel = crocoddyl.CostModelSum(state) 50 | # Create cost terms 51 | # Control regularization cost 52 | uResidual = crocoddyl.ResidualModelControlGrav(state) 53 | uRegCost = crocoddyl.CostModelResidual(state, uResidual) 54 | # State regularization cost 55 | xResidual = crocoddyl.ResidualModelState(state, x0) 56 | xRegCost = crocoddyl.CostModelResidual(state, xResidual) 57 | # endeff frame translation cost 58 | endeff_frame_id = robot_simulator.pin_robot.model.getFrameId("contact") 59 | # endeff_translation = robot.data.oMf[endeff_frame_id].translation.copy() 60 | endeff_translation = np.array([-0.4, 0.3, 0.7]) # move endeff +10 cm along x in WORLD frame 61 | frameTranslationResidual = crocoddyl.ResidualModelFrameTranslation(state, endeff_frame_id, endeff_translation) 62 | frameTranslationCost = crocoddyl.CostModelResidual(state, frameTranslationResidual) 63 | # Add costs 64 | runningCostModel.addCost("stateReg", xRegCost, 1e-1) 65 | runningCostModel.addCost("ctrlRegGrav", uRegCost, 1e-4) 66 | runningCostModel.addCost("translation", frameTranslationCost, 10) 67 | terminalCostModel.addCost("stateReg", xRegCost, 1e-1) 68 | terminalCostModel.addCost("translation", frameTranslationCost, 10) 69 | # Create Differential Action Model (DAM), i.e. continuous dynamics and cost functions 70 | running_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics(state, actuation, runningCostModel) 71 | terminal_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics(state, actuation, terminalCostModel) 72 | # Create Integrated Action Model (IAM), i.e. Euler integration of continuous dynamics and cost 73 | dt = 1e-2 74 | runningModel = crocoddyl.IntegratedActionModelEuler(running_DAM, dt) 75 | terminalModel = crocoddyl.IntegratedActionModelEuler(terminal_DAM, 0.) 76 | # Create the shooting problem 77 | T = 100 78 | problem = crocoddyl.ShootingProblem(x0, [runningModel] * T, terminalModel) 79 | # Create solver + callbacks 80 | solver = mim_solvers.SolverSQP(problem) 81 | # solver.setCallbacks([crocoddyl.CallbackLogger(), 82 | # crocoddyl.CallbackVerbose()]) 83 | # Warm start : initial state + gravity compensation 84 | xs_init = [x0 for i in range(T+1)] 85 | us_init = solver.problem.quasiStatic(xs_init[:-1]) 86 | # Solve 87 | solver.termination_tolerance = 1e-4 88 | solver.with_callbacks = True 89 | solver.solve(xs_init, us_init, 100) 90 | solver.with_callbacks = False 91 | 92 | 93 | # # # # # # # # # # # # 94 | ### MPC SIMULATION ### 95 | # # # # # # # # # # # # 96 | # OCP parameters 97 | ocp_params = {} 98 | ocp_params['N_h'] = T 99 | ocp_params['dt'] = dt 100 | ocp_params['maxiter'] = 10 101 | ocp_params['pin_model'] = robot_simulator.pin_robot.model 102 | ocp_params['armature'] = runningModel.differential.armature 103 | ocp_params['id_endeff'] = endeff_frame_id 104 | ocp_params['active_costs'] = solver.problem.runningModels[0].differential.costs.active.tolist() 105 | # Simu parameters 106 | sim_params = {} 107 | sim_params['sim_freq'] = int(1./env.dt) 108 | sim_params['mpc_freq'] = 1000 109 | sim_params['T_sim'] = 2. 110 | log_rate = 100 111 | # Initialize simulation data 112 | sim_data = mpc_utils.init_sim_data(sim_params, ocp_params, x0) 113 | # Display target 114 | mpc_utils.display_ball(endeff_translation, RADIUS=.05, COLOR=[1.,0.,0.,.6]) 115 | # Simulate 116 | mpc_cycle = 0 117 | for i in range(sim_data['N_sim']): 118 | 119 | if(i%log_rate==0): 120 | print("\n SIMU step "+str(i)+"/"+str(sim_data['N_sim'])+"\n") 121 | 122 | # Solve OCP if we are in a planning cycle (MPC/planning frequency) 123 | if(i%int(sim_params['sim_freq']/sim_params['mpc_freq']) == 0): 124 | # Set x0 to measured state 125 | solver.problem.x0 = sim_data['state_mea_SIM_RATE'][i, :] 126 | # Warm start using previous solution 127 | xs_init = list(solver.xs[1:]) + [solver.xs[-1]] 128 | xs_init[0] = sim_data['state_mea_SIM_RATE'][i, :] 129 | us_init = list(solver.us[1:]) + [solver.us[-1]] 130 | 131 | # Solve OCP & record MPC predictions 132 | solver.solve(xs_init, us_init, ocp_params['maxiter']) 133 | sim_data['state_pred'][mpc_cycle, :, :] = np.array(solver.xs) 134 | sim_data['ctrl_pred'][mpc_cycle, :, :] = np.array(solver.us) 135 | # Extract relevant predictions for interpolations 136 | x_curr = sim_data['state_pred'][mpc_cycle, 0, :] # x0* = measured state (q^, v^ ) 137 | x_pred = sim_data['state_pred'][mpc_cycle, 1, :] # x1* = predicted state (q1*, v1*) 138 | u_curr = sim_data['ctrl_pred'][mpc_cycle, 0, :] # u0* = optimal control (tau0*) 139 | # Record costs references 140 | q = sim_data['state_pred'][mpc_cycle, 0, :sim_data['nq']] 141 | sim_data['ctrl_ref'][mpc_cycle, :] = pin_utils.get_u_grav(q, solver.problem.runningModels[0].differential.pinocchio, ocp_params['armature']) 142 | sim_data['state_ref'][mpc_cycle, :] = solver.problem.runningModels[0].differential.costs.costs['stateReg'].cost.residual.reference 143 | sim_data['lin_pos_ee_ref'][mpc_cycle, :] = solver.problem.runningModels[0].differential.costs.costs['translation'].cost.residual.reference 144 | 145 | 146 | # Select reference control and state for the current MPC cycle 147 | x_ref_MPC_RATE = x_curr + sim_data['ocp_to_mpc_ratio'] * (x_pred - x_curr) 148 | u_ref_MPC_RATE = u_curr 149 | if(mpc_cycle==0): 150 | sim_data['state_des_MPC_RATE'][mpc_cycle, :] = x_curr 151 | sim_data['ctrl_des_MPC_RATE'][mpc_cycle, :] = u_ref_MPC_RATE 152 | sim_data['state_des_MPC_RATE'][mpc_cycle+1, :] = x_ref_MPC_RATE 153 | 154 | # Increment planning counter 155 | mpc_cycle += 1 156 | 157 | 158 | # Select reference control and state for the current SIMU cycle 159 | x_ref_SIM_RATE = x_curr + sim_data['ocp_to_mpc_ratio'] * (x_pred - x_curr) 160 | u_ref_SIM_RATE = u_curr 161 | 162 | # First prediction = measurement = initialization of MPC 163 | if(i==0): 164 | sim_data['state_des_SIM_RATE'][i, :] = x_curr 165 | sim_data['ctrl_des_SIM_RATE'][i, :] = u_ref_SIM_RATE 166 | sim_data['state_des_SIM_RATE'][i+1, :] = x_ref_SIM_RATE 167 | 168 | # Send torque to simulator & step simulator 169 | robot_simulator.send_joint_command(u_ref_SIM_RATE) 170 | env.step() 171 | # Measure new state from simulator 172 | q_mea_SIM_RATE, v_mea_SIM_RATE = robot_simulator.get_state() 173 | # Update pinocchio model 174 | robot_simulator.forward_robot(q_mea_SIM_RATE, v_mea_SIM_RATE) 175 | # Record data 176 | x_mea_SIM_RATE = np.concatenate([q_mea_SIM_RATE, v_mea_SIM_RATE]).T 177 | sim_data['state_mea_SIM_RATE'][i+1, :] = x_mea_SIM_RATE 178 | 179 | 180 | plot_data = mpc_utils.extract_plot_data_from_sim_data(sim_data) 181 | 182 | mpc_utils.plot_mpc_results(plot_data, which_plots=['all'], PLOT_PREDICTIONS=True, pred_plot_sampling=int(sim_params['mpc_freq']/10)) 183 | -------------------------------------------------------------------------------- /mpc_utils.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | import pinocchio as pin 4 | from matplotlib.collections import LineCollection 5 | import matplotlib.pyplot as plt 6 | import matplotlib 7 | import pin_utils 8 | 9 | import pybullet as p 10 | import os 11 | 12 | 13 | # Get contact wrench from robot simulator 14 | def get_contact_wrench(pybullet_simulator, id_endeff): 15 | ''' 16 | Get contact wrench in LOCAL contact frame 17 | ''' 18 | contact_points = p.getContactPoints() 19 | force = np.zeros(6) 20 | for ci in reversed(contact_points): 21 | p_ct = np.array(ci[6]) 22 | contact_normal = ci[7] 23 | normal_force = ci[9] 24 | lateral_friction_direction_1 = ci[11] 25 | lateral_friction_force_1 = ci[10] 26 | lateral_friction_direction_2 = ci[13] 27 | lateral_friction_force_2 = ci[12] 28 | # Wrench in LOCAL contact frame 29 | linear_LOCAL = np.array([normal_force, lateral_friction_force_1, lateral_friction_force_2]) 30 | wrench_LOCAL = np.concatenate([linear_LOCAL, np.zeros(3)]) 31 | # LOCAL contact placement 32 | R_ct = np.vstack([np.array(contact_normal), np.array(lateral_friction_direction_1), np.array(lateral_friction_direction_2)]).T 33 | M_ct = pin.SE3(R_ct, p_ct) 34 | # wrench LOCAL(p)-->WORLD 35 | wrench_WORLD = M_ct.act(pin.Force(wrench_LOCAL)) 36 | # wrench WORLD-->LOCAL(EE) 37 | wrench_croco = -pybullet_simulator.pin_robot.data.oMf[id_endeff].actInv(wrench_WORLD) 38 | force =+ wrench_croco.vector 39 | return force 40 | 41 | # Display 42 | def display_ball(p_des, RADIUS=.05, COLOR=[1.,1.,1.,1.]): 43 | ''' 44 | Create a sphere visual object in PyBullet (no collision) 45 | Transformed because reference p_des is in pinocchio WORLD frame, which is different 46 | than PyBullet WORLD frame if the base placement in the simulator is not (eye(3), zeros(3)) 47 | INPUT: 48 | p_des : desired position of the ball in pinocchio.WORLD 49 | robot_base_pose : initial pose of the robot BASE in bullet.WORLD 50 | RADIUS : radius of the ball 51 | COLOR : color of the ball 52 | ''' 53 | # pose of the sphere in bullet WORLD 54 | M = pin.SE3(np.eye(3), p_des) # ok for talos reduced since pin.W = bullet.W but careful with talos_arm if base is moved 55 | quat = pin.SE3ToXYZQUAT(M) 56 | visualBallId = p.createVisualShape(shapeType=p.GEOM_SPHERE, 57 | radius=RADIUS, 58 | rgbaColor=COLOR, 59 | visualFramePosition=quat[:3], 60 | visualFrameOrientation=quat[3:]) 61 | ballId = p.createMultiBody(baseMass=0., 62 | baseInertialFramePosition=[0.,0.,0.], 63 | baseVisualShapeIndex=visualBallId, 64 | basePosition=[0.,0.,0.], 65 | useMaximalCoordinates=True) 66 | 67 | return ballId 68 | 69 | # Load contact surface in PyBullet for contact experiments 70 | def display_contact_surface(M, robotId=1, radius=.25, length=0.0, with_collision=False, TILT=[0., 0., 0.]): 71 | ''' 72 | Creates contact surface object in PyBullet as a flat cylinder 73 | M : contact placement (with z_LOCAL coinciding with cylinder axis) 74 | robotId : id of the robot 75 | ''' 76 | # Tilt contact surface (default 0) 77 | TILT_rotation = pin.utils.rpyToMatrix(TILT[0], TILT[1], TILT[2]) 78 | M.rotation = TILT_rotation.dot(M.rotation) 79 | # Get quaternion 80 | quat = pin.SE3ToXYZQUAT(M) 81 | visualShapeId = p.createVisualShape(shapeType=p.GEOM_CYLINDER, 82 | radius=radius, 83 | length=length, 84 | rgbaColor=[.1, .8, .1, .5], 85 | visualFramePosition=quat[:3], 86 | visualFrameOrientation=quat[3:]) 87 | # With collision 88 | if(with_collision): 89 | collisionShapeId = p.createCollisionShape(shapeType=p.GEOM_CYLINDER, 90 | radius=radius, 91 | height=length, 92 | collisionFramePosition=quat[:3], 93 | collisionFrameOrientation=quat[3:]) 94 | contactId = p.createMultiBody(baseMass=0., 95 | baseInertialFramePosition=[0.,0.,0.], 96 | baseCollisionShapeIndex=collisionShapeId, 97 | baseVisualShapeIndex=visualShapeId, 98 | basePosition=[0.,0.,0.], 99 | useMaximalCoordinates=True) 100 | 101 | # Desactivate collisions for all links except end-effector of robot 102 | # TODO: do not hard-code the PyBullet EE id 103 | for i in range(p.getNumJoints(robotId)): 104 | p.setCollisionFilterPair(contactId, robotId, -1, i, 0) 105 | p.setCollisionFilterPair(contactId, robotId, -1, 8, 1) 106 | 107 | return contactId 108 | # Without collisions 109 | else: 110 | contactId = p.createMultiBody(baseMass=0., 111 | baseInertialFramePosition=[0.,0.,0.], 112 | baseVisualShapeIndex=visualShapeId, 113 | basePosition=[0.,0.,0.], 114 | useMaximalCoordinates=True) 115 | return contactId 116 | 117 | 118 | 119 | # Initialize simulation data for MPC simulation 120 | def init_sim_data(sim_params, ocp_params, x0): 121 | ''' 122 | Initialize simulation data from config file 123 | sim_params : dict of sim params 124 | N_h : number of nodes in OCP horizon 125 | x0 : initial state of the 126 | ''' 127 | sim_data = {} 128 | # MPC & simulation parameters 129 | sim_data['T_sim'] = sim_params['T_sim'] # Total duration of simulation (s) 130 | sim_data['sim_freq'] = sim_params['sim_freq'] # Simulation frequency 131 | sim_data['mpc_freq'] = sim_params['mpc_freq'] # Planning frequency (OCP solution update rate) 132 | sim_data['N_mpc'] = int(sim_data['T_sim']*sim_data['mpc_freq']) # Total number of planning steps in the simulation 133 | sim_data['N_sim'] = int(sim_data['T_sim']*sim_data['sim_freq']) # Total number of simulation steps 134 | sim_data['dt_mpc'] = float(1./sim_data['mpc_freq']) # Duration of 1 planning cycle (s) 135 | sim_data['dt_sim'] = float(1./sim_data['sim_freq']) # Duration of 1 simulation cycle (s) 136 | sim_data['ocp_to_mpc_ratio'] = sim_data['dt_mpc']/ocp_params['dt'] 137 | sim_data['ocp_to_sim_ratio'] = sim_data['dt_sim']/ocp_params['dt'] 138 | # Copy OCP params 139 | sim_data['nq'] = ocp_params['pin_model'].nq 140 | sim_data['nv'] = ocp_params['pin_model'].nv 141 | sim_data['nu'] = ocp_params['pin_model'].nq 142 | sim_data['nx'] = sim_data['nq'] + sim_data['nv'] 143 | sim_data['pin_model'] = ocp_params['pin_model'] 144 | sim_data['id_endeff'] = ocp_params['id_endeff'] 145 | sim_data['armature'] = ocp_params['armature'] 146 | sim_data['T_h'] = ocp_params['N_h']*ocp_params['dt'] # Duration of the MPC horizon (s) 147 | sim_data['N_h'] = ocp_params['N_h'] # Number of nodes in MPC horizon 148 | sim_data['active_costs'] = ocp_params['active_costs'] # List of ative costs names 149 | # Cost references 150 | sim_data['ctrl_ref'] = np.zeros((sim_data['N_mpc'], sim_data['nu'])) 151 | sim_data['state_ref'] = np.zeros((sim_data['N_mpc'], sim_data['nx'])) 152 | sim_data['lin_pos_ee_ref'] = np.zeros((sim_data['N_mpc'], 3)) 153 | sim_data['lin_vel_ee_ref'] = np.zeros((sim_data['N_mpc'], 3)) 154 | sim_data['ang_pos_ee_ref'] = np.zeros((sim_data['N_mpc'], 3)) 155 | sim_data['ang_vel_ee_ref'] = np.zeros((sim_data['N_mpc'], 3)) 156 | sim_data['f_ee_ref'] = np.zeros((sim_data['N_mpc'], 6)) 157 | # Predictions 158 | sim_data['state_pred'] = np.zeros((sim_data['N_mpc'], ocp_params['N_h']+1, sim_data['nx'])) # Predicted states ( ddp.xs : {x* = (q*, v*)} ) 159 | sim_data['ctrl_pred'] = np.zeros((sim_data['N_mpc'], ocp_params['N_h'], sim_data['nu'])) # Predicted torques ( ddp.us : {u*} ) 160 | sim_data['force_pred'] = np.zeros((sim_data['N_mpc'], ocp_params['N_h'], 6)) # Predicted EE contact forces 161 | sim_data['state_des_MPC_RATE'] = np.zeros((sim_data['N_mpc']+1, sim_data['nx'])) # Predicted states at planner frequency ( x* interpolated at PLAN freq ) 162 | sim_data['ctrl_des_MPC_RATE'] = np.zeros((sim_data['N_mpc'], sim_data['nu'])) # Predicted torques at planner frequency ( u* interpolated at PLAN freq ) 163 | sim_data['force_des_MPC_RATE'] = np.zeros((sim_data['N_mpc'], 6)) # Predicted EE contact forces planner frequency 164 | sim_data['state_des_SIM_RATE'] = np.zeros((sim_data['N_sim']+1, sim_data['nx'])) # Reference state at actuation freq ( x* interpolated at SIMU freq ) 165 | sim_data['ctrl_des_SIM_RATE'] = np.zeros((sim_data['N_sim'], sim_data['nu'])) # Reference input at actuation freq ( u* interpolated at SIMU freq ) 166 | sim_data['force_des_SIM_RATE'] = np.zeros((sim_data['N_sim'], 6)) # Reference EE contact force at actuation freq 167 | # Measurements 168 | sim_data['state_mea_SIM_RATE'] = np.zeros((sim_data['N_sim']+1, sim_data['nx'])) # Measured states ( x^mea = (q, v) from actuator & PyB at SIMU freq ) 169 | sim_data['force_mea_SIM_RATE'] = np.zeros((sim_data['N_sim'], 6)) 170 | sim_data['state_mea_SIM_RATE'][0, :] = x0 171 | print('') 172 | print(' *************************') 173 | print(' ** Simulation is ready **') 174 | print(' *************************') 175 | print("-------------------------------------------------------------------") 176 | print('- Total simulation duration : T_sim = '+str(sim_data['T_sim'])+' s') 177 | print('- Simulation frequency : f_simu = '+str(float(sim_data['sim_freq']/1000.))+' kHz') 178 | print('- Replanning frequency : f_plan = '+str(float(sim_data['mpc_freq']/1000.))+' kHz') 179 | print('- Total # of simulation steps : N_sim = '+str(sim_data['N_sim'])) 180 | print('- Total # of planning steps : N_mpc = '+str(sim_data['N_mpc'])) 181 | print('- Duration of MPC horizon : T_ocp = '+str(sim_data['T_h'])+' s') 182 | print('- OCP integration step : dt = '+str(ocp_params['dt'])+' s') 183 | print("-------------------------------------------------------------------") 184 | print('') 185 | # time.sleep(2) 186 | 187 | return sim_data 188 | 189 | 190 | # Extract MPC simu-specific plotting data from sim data 191 | def extract_plot_data_from_sim_data(sim_data): 192 | ''' 193 | Extract plot data from simu data 194 | ''' 195 | 196 | plot_data = {} 197 | # Get costs 198 | plot_data['active_costs'] = sim_data['active_costs'] 199 | # Robot model & params 200 | plot_data['pin_model'] = sim_data['pin_model'] 201 | nq = plot_data['pin_model'].nq; plot_data['nq'] = nq 202 | nv = plot_data['pin_model'].nv; plot_data['nv'] = nv 203 | nx = nq+nv; plot_data['nx'] = nx 204 | nu = nq 205 | # MPC params 206 | plot_data['T_sim'] = sim_data['T_sim'] 207 | plot_data['N_sim'] = sim_data['N_sim']; plot_data['N_mpc'] = sim_data['N_mpc'] 208 | plot_data['dt_mpc'] = sim_data['dt_mpc']; plot_data['dt_sim'] = sim_data['dt_sim'] 209 | plot_data['T_h'] = sim_data['T_h']; 210 | plot_data['N_h'] = sim_data['N_h'] 211 | # plot_data['alpha'] = sim_data['alpha']; plot_data['beta'] = sim_data['beta'] 212 | # Record cost references 213 | plot_data['ctrl_ref'] = sim_data['ctrl_ref'] 214 | plot_data['state_ref'] = sim_data['state_ref'] 215 | plot_data['lin_pos_ee_ref'] = sim_data['lin_pos_ee_ref'] 216 | plot_data['f_ee_ref'] = sim_data['f_ee_ref'] 217 | # Control predictions 218 | plot_data['u_pred'] = sim_data['ctrl_pred'] 219 | plot_data['u_des_MPC_RATE'] = sim_data['ctrl_des_MPC_RATE'] 220 | plot_data['u_des_SIM_RATE'] = sim_data['ctrl_des_SIM_RATE'] 221 | # State predictions (at PLAN freq) 222 | plot_data['q_pred'] = sim_data['state_pred'][:,:,:nq] 223 | plot_data['v_pred'] = sim_data['state_pred'][:,:,nq:nq+nv] 224 | plot_data['q_des_MPC_RATE'] = sim_data['state_des_MPC_RATE'][:,:nq] 225 | plot_data['v_des_MPC_RATE'] = sim_data['state_des_MPC_RATE'][:,nq:nq+nv] 226 | plot_data['q_des_SIM_RATE'] = sim_data['state_des_SIM_RATE'][:,:nq] 227 | plot_data['v_des_SIM_RATE'] = sim_data['state_des_SIM_RATE'][:,nq:nq+nv] 228 | # State measurements (at SIMU freq) 229 | plot_data['q_mea'] = sim_data['state_mea_SIM_RATE'][:,:nq] 230 | plot_data['v_mea'] = sim_data['state_mea_SIM_RATE'][:,nq:nq+nv] 231 | # Extract gravity torques 232 | plot_data['grav'] = np.zeros((sim_data['N_sim']+1, plot_data['nq'])) 233 | for i in range(plot_data['N_sim']+1): 234 | plot_data['grav'][i,:] = pin_utils.get_u_grav(plot_data['q_mea'][i,:], plot_data['pin_model'], sim_data['armature']) 235 | # EE predictions (at PLAN freq) 236 | # Linear position velocity of EE 237 | plot_data['lin_pos_ee_pred'] = np.zeros((sim_data['N_mpc'], sim_data['N_h']+1, 3)) 238 | plot_data['lin_vel_ee_pred'] = np.zeros((sim_data['N_mpc'], sim_data['N_h']+1, 3)) 239 | # Angular position velocity of EE 240 | plot_data['ang_pos_ee_pred'] = np.zeros((sim_data['N_mpc'], sim_data['N_h']+1, 3)) 241 | plot_data['ang_vel_ee_pred'] = np.zeros((sim_data['N_mpc'], sim_data['N_h']+1, 3)) 242 | for node_id in range(sim_data['N_h']+1): 243 | plot_data['lin_pos_ee_pred'][:, node_id, :] = pin_utils.get_p_(plot_data['q_pred'][:, node_id, :], plot_data['pin_model'], sim_data['id_endeff']) 244 | plot_data['lin_vel_ee_pred'][:, node_id, :] = pin_utils.get_v_(plot_data['q_pred'][:, node_id, :], plot_data['v_pred'][:, node_id, :], plot_data['pin_model'], sim_data['id_endeff']) 245 | plot_data['ang_pos_ee_pred'][:, node_id, :] = pin_utils.get_rpy_(plot_data['q_pred'][:, node_id, :], plot_data['pin_model'], sim_data['id_endeff']) 246 | plot_data['ang_vel_ee_pred'][:, node_id, :] = pin_utils.get_w_(plot_data['q_pred'][:, node_id, :], plot_data['v_pred'][:, node_id, :], plot_data['pin_model'], sim_data['id_endeff']) 247 | # EE measurements (at SIMU freq) 248 | # Linear 249 | plot_data['lin_pos_ee_mea'] = pin_utils.get_p_(plot_data['q_mea'], sim_data['pin_model'], sim_data['id_endeff']) 250 | plot_data['lin_vel_ee_mea'] = pin_utils.get_v_(plot_data['q_mea'], plot_data['v_mea'], sim_data['pin_model'], sim_data['id_endeff']) 251 | # Angular 252 | plot_data['ang_pos_ee_mea'] = pin_utils.get_rpy_(plot_data['q_mea'], sim_data['pin_model'], sim_data['id_endeff']) 253 | plot_data['ang_vel_ee_mea'] = pin_utils.get_w_(plot_data['q_mea'], plot_data['v_mea'], sim_data['pin_model'], sim_data['id_endeff']) 254 | # EE des 255 | # Linear 256 | plot_data['lin_pos_ee_des_MPC_RATE'] = pin_utils.get_p_(plot_data['q_des_MPC_RATE'], sim_data['pin_model'], sim_data['id_endeff']) 257 | plot_data['lin_vel_ee_des_MPC_RATE'] = pin_utils.get_v_(plot_data['q_des_MPC_RATE'], plot_data['v_des_MPC_RATE'], sim_data['pin_model'], sim_data['id_endeff']) 258 | plot_data['lin_pos_ee_des_SIM_RATE'] = pin_utils.get_p_(plot_data['q_des_SIM_RATE'], sim_data['pin_model'], sim_data['id_endeff']) 259 | plot_data['lin_vel_ee_des_SIM_RATE'] = pin_utils.get_v_(plot_data['q_des_SIM_RATE'], plot_data['v_des_SIM_RATE'], sim_data['pin_model'], sim_data['id_endeff']) 260 | # Angular 261 | plot_data['ang_pos_ee_des_MPC_RATE'] = pin_utils.get_rpy_(plot_data['q_des_MPC_RATE'], sim_data['pin_model'], sim_data['id_endeff']) 262 | plot_data['ang_vel_ee_des_MPC_RATE'] = pin_utils.get_w_(plot_data['q_des_MPC_RATE'], plot_data['v_des_MPC_RATE'], sim_data['pin_model'], sim_data['id_endeff']) 263 | plot_data['ang_pos_ee_des_SIM_RATE'] = pin_utils.get_rpy_(plot_data['q_des_SIM_RATE'], sim_data['pin_model'], sim_data['id_endeff']) 264 | plot_data['ang_vel_ee_des_SIM_RATE'] = pin_utils.get_w_(plot_data['q_des_SIM_RATE'], plot_data['v_des_SIM_RATE'], sim_data['pin_model'], sim_data['id_endeff']) 265 | # Extract EE force 266 | plot_data['f_ee_pred'] = sim_data['force_pred'] 267 | plot_data['f_ee_mea'] = sim_data['force_mea_SIM_RATE'] 268 | plot_data['f_ee_des_MPC_RATE'] = sim_data['force_des_MPC_RATE'] 269 | plot_data['f_ee_des_SIM_RATE'] = sim_data['force_des_SIM_RATE'] 270 | 271 | return plot_data 272 | 273 | 274 | # Plot from MPC simulation 275 | def plot_mpc_results(plot_data, which_plots=None, PLOT_PREDICTIONS=False, 276 | pred_plot_sampling=100, 277 | SAVE=False, SAVE_DIR=None, SAVE_NAME=None, 278 | SHOW=True, 279 | AUTOSCALE=False): 280 | ''' 281 | Plot sim data 282 | Input: 283 | plot_data : plotting data 284 | PLOT_PREDICTIONS : True or False 285 | pred_plot_sampling : plot every pred_plot_sampling prediction 286 | to avoid huge amount of plotted data 287 | ("1" = plot all) 288 | SAVE, SAVE_DIR, SAVE_NAME : save plots as .png 289 | SHOW : show plots 290 | AUTOSCALE : rescale y-axis of endeff plot 291 | based on maximum value taken 292 | ''' 293 | 294 | plots = {} 295 | 296 | if('x' in which_plots or which_plots is None or which_plots =='all' or 'all' in which_plots): 297 | plots['x'] = plot_mpc_state(plot_data, PLOT_PREDICTIONS=PLOT_PREDICTIONS, 298 | pred_plot_sampling=pred_plot_sampling, 299 | SAVE=SAVE, SAVE_DIR=SAVE_DIR, SAVE_NAME=SAVE_NAME, 300 | SHOW=False) 301 | 302 | if('u' in which_plots or which_plots is None or which_plots =='all' or 'all' in which_plots): 303 | plots['u'] = plot_mpc_control(plot_data, PLOT_PREDICTIONS=PLOT_PREDICTIONS, 304 | pred_plot_sampling=pred_plot_sampling, 305 | SAVE=SAVE, SAVE_DIR=SAVE_DIR, SAVE_NAME=SAVE_NAME, 306 | SHOW=False) 307 | 308 | if('ee' in which_plots or which_plots is None or which_plots =='all' or 'all' in which_plots): 309 | plots['ee_lin'] = plot_mpc_endeff_linear(plot_data, PLOT_PREDICTIONS=PLOT_PREDICTIONS, 310 | pred_plot_sampling=pred_plot_sampling, 311 | SAVE=SAVE, SAVE_DIR=SAVE_DIR, SAVE_NAME=SAVE_NAME, 312 | SHOW=False, AUTOSCALE=AUTOSCALE) 313 | plots['ee_ang'] = plot_mpc_endeff_angular(plot_data, PLOT_PREDICTIONS=PLOT_PREDICTIONS, 314 | pred_plot_sampling=pred_plot_sampling, 315 | SAVE=SAVE, SAVE_DIR=SAVE_DIR, SAVE_NAME=SAVE_NAME, 316 | SHOW=False, AUTOSCALE=AUTOSCALE) 317 | 318 | if('f' in which_plots or which_plots is None or which_plots =='all' or 'all' in which_plots): 319 | plots['f'] = plot_mpc_force(plot_data, PLOT_PREDICTIONS=PLOT_PREDICTIONS, 320 | pred_plot_sampling=pred_plot_sampling, 321 | SAVE=SAVE, SAVE_DIR=SAVE_DIR, SAVE_NAME=SAVE_NAME, 322 | SHOW=False, AUTOSCALE=AUTOSCALE) 323 | if(SHOW): 324 | plt.show() 325 | plt.close('all') 326 | 327 | # Plot state data 328 | def plot_mpc_state(plot_data, PLOT_PREDICTIONS=False, 329 | pred_plot_sampling=100, 330 | SAVE=False, SAVE_DIR=None, SAVE_NAME=None, 331 | SHOW=True): 332 | ''' 333 | Plot state data 334 | Input: 335 | plot_data : plotting data 336 | PLOT_PREDICTIONS : True or False 337 | pred_plot_sampling : plot every pred_plot_sampling prediction 338 | to avoid huge amount of plotted data 339 | ("1" = plot all) 340 | SAVE, SAVE_DIR, SAVE_NAME : save plots as .png 341 | SHOW : show plots 342 | ''' 343 | T_sim = plot_data['T_sim'] 344 | N_sim = plot_data['N_sim'] 345 | N_mpc = plot_data['N_mpc'] 346 | dt_mpc = plot_data['dt_mpc'] 347 | nq = plot_data['nq'] 348 | nx = plot_data['nx'] 349 | T_h = plot_data['T_h'] 350 | N_h = plot_data['N_h'] 351 | # Create time spans for X and U + Create figs and subplots 352 | t_span_simu = np.linspace(0, T_sim, N_sim+1) 353 | t_span_plan = np.linspace(0, T_sim, N_mpc+1) 354 | fig_x, ax_x = plt.subplots(nq, 2, figsize=(19.2,10.8), sharex='col') 355 | # For each joint 356 | for i in range(nq): 357 | 358 | if(PLOT_PREDICTIONS): 359 | 360 | # Extract state predictions of i^th joint 361 | q_pred_i = plot_data['q_pred'][:,:,i] 362 | v_pred_i = plot_data['v_pred'][:,:,i] 363 | # For each planning step in the trajectory 364 | for j in range(0, N_mpc, pred_plot_sampling): 365 | # Receding horizon = [j,j+N_h] 366 | t0_horizon = j*dt_mpc 367 | tspan_x_pred = np.linspace(t0_horizon, t0_horizon + T_h, N_h+1) 368 | tspan_u_pred = np.linspace(t0_horizon, t0_horizon + T_h - dt_mpc, N_h) 369 | # Set up lists of (x,y) points for predicted positions and velocities 370 | points_q = np.array([tspan_x_pred, q_pred_i[j,:]]).transpose().reshape(-1,1,2) 371 | points_v = np.array([tspan_x_pred, v_pred_i[j,:]]).transpose().reshape(-1,1,2) 372 | # Set up lists of segments 373 | segs_q = np.concatenate([points_q[:-1], points_q[1:]], axis=1) 374 | segs_v = np.concatenate([points_v[:-1], points_v[1:]], axis=1) 375 | # Make collections segments 376 | cm = plt.get_cmap('Greys_r') 377 | lc_q = LineCollection(segs_q, cmap=cm, zorder=-1) 378 | lc_v = LineCollection(segs_v, cmap=cm, zorder=-1) 379 | lc_q.set_array(tspan_x_pred) 380 | lc_v.set_array(tspan_x_pred) 381 | # Customize 382 | lc_q.set_linestyle('-') 383 | lc_v.set_linestyle('-') 384 | lc_q.set_linewidth(1) 385 | lc_v.set_linewidth(1) 386 | # Plot collections 387 | ax_x[i,0].add_collection(lc_q) 388 | ax_x[i,1].add_collection(lc_v) 389 | # Scatter to highlight points 390 | colors = np.r_[np.linspace(0.1, 1, N_h), 1] 391 | my_colors = cm(colors) 392 | ax_x[i,0].scatter(tspan_x_pred, q_pred_i[j,:], s=10, zorder=1, c=my_colors, cmap=matplotlib.cm.Greys) #c='black', 393 | ax_x[i,1].scatter(tspan_x_pred, v_pred_i[j,:], s=10, zorder=1, c=my_colors, cmap=matplotlib.cm.Greys) #c='black', 394 | 395 | # Joint position 396 | ax_x[i,0].plot(t_span_plan, plot_data['q_des_MPC_RATE'][:,i], color='b', linestyle='-', marker='.', label='Predicted', alpha=0.1) 397 | # ax_x[i,0].plot(t_span_simu, plot_data['q_des_SIM_RATE'][:,i], color='y', linestyle='-', marker='.', label='Predicted (SIMU rate)', alpha=0.5) 398 | ax_x[i,0].plot(t_span_simu, plot_data['q_mea'][:,i], 'r-', label='Measured', linewidth=1, alpha=0.3) 399 | # Plot joint position regularization reference 400 | if('stateReg' in plot_data['active_costs']): 401 | ax_x[i,0].plot(t_span_plan[:-1], plot_data['state_ref'][:, i], linestyle='-.', color='k', marker=None, label='xReg_ref', alpha=0.5) 402 | ax_x[i,0].set_ylabel('$q_{}$'.format(i), fontsize=12) 403 | ax_x[i,0].yaxis.set_major_locator(plt.MaxNLocator(2)) 404 | ax_x[i,0].yaxis.set_major_formatter(plt.FormatStrFormatter('%.2e')) 405 | ax_x[i,0].grid(True) 406 | 407 | # Joint velocity 408 | ax_x[i,1].plot(t_span_plan, plot_data['v_des_MPC_RATE'][:,i], color='b', linestyle='-', marker='.', label='Predicted', alpha=0.5) 409 | # ax_x[i,1].plot(t_span_simu, plot_data['v_des_SIM_RATE'][:,i], color='y', linestyle='-', marker='.', label='Predicted (SIMU)', alpha=0.5) 410 | ax_x[i,1].plot(t_span_simu, plot_data['v_mea'][:,i], 'r-', label='Measured', linewidth=1, alpha=0.3) 411 | if('stateReg' in plot_data['active_costs']): 412 | ax_x[i,1].plot(t_span_plan[:-1], plot_data['state_ref'][:, i+nq], linestyle='-.', color='k', marker=None, label='xReg_ref', alpha=0.5) 413 | ax_x[i,1].set_ylabel('$v_{}$'.format(i), fontsize=12) 414 | ax_x[i,1].yaxis.set_major_locator(plt.MaxNLocator(2)) 415 | ax_x[i,1].yaxis.set_major_formatter(plt.FormatStrFormatter('%.2e')) 416 | ax_x[i,1].grid(True) 417 | 418 | # Add xlabel on bottom plot of each column 419 | if(i == nq-1): 420 | ax_x[i,0].set_xlabel('t(s)', fontsize=16) 421 | ax_x[i,1].set_xlabel('t(s)', fontsize=16) 422 | # Legend 423 | handles_x, labels_x = ax_x[i,0].get_legend_handles_labels() 424 | fig_x.legend(handles_x, labels_x, loc='upper right', prop={'size': 16}) 425 | # y axis labels 426 | fig_x.text(0.05, 0.5, 'Joint position (rad)', va='center', rotation='vertical', fontsize=16) 427 | fig_x.text(0.49, 0.5, 'Joint velocity (rad/s)', va='center', rotation='vertical', fontsize=16) 428 | fig_x.subplots_adjust(wspace=0.27) 429 | # Titles 430 | fig_x.suptitle('State = joint positions, velocities', size=18) 431 | # Save fig 432 | if(SAVE): 433 | figs = {'x': fig_x} 434 | if(SAVE_DIR is None): 435 | print("SAVE FIGURES IN HOME") 436 | SAVE_DIR = os.environ['HOME'] 437 | if(SAVE_NAME is None): 438 | SAVE_NAME = 'testfig' 439 | for name, fig in figs.items(): 440 | fig.savefig(SAVE_DIR + '/' +str(name) + '_' + SAVE_NAME +'.png') 441 | 442 | if(SHOW): 443 | plt.show() 444 | 445 | return fig_x 446 | 447 | # Plot control data 448 | def plot_mpc_control(plot_data, PLOT_PREDICTIONS=False, 449 | pred_plot_sampling=100, 450 | SAVE=False, SAVE_DIR=None, SAVE_NAME=None, 451 | SHOW=True, 452 | AUTOSCALE=False): 453 | ''' 454 | Plot control data 455 | Input: 456 | plot_data : plotting data 457 | PLOT_PREDICTIONS : True or False 458 | pred_plot_sampling : plot every pred_plot_sampling prediction 459 | to avoid huge amount of plotted data 460 | ("1" = plot all) 461 | SAVE, SAVE_DIR, SAVE_NAME : save plots as .png 462 | SHOW : show plots 463 | ''' 464 | T_sim = plot_data['T_sim'] 465 | N_sim = plot_data['N_sim'] 466 | N_mpc = plot_data['N_mpc'] 467 | dt_mpc = plot_data['dt_mpc'] 468 | dt_sim = plot_data['dt_sim'] 469 | nq = plot_data['nq'] 470 | T_h = plot_data['T_h'] 471 | N_h = plot_data['N_h'] 472 | # Create time spans for X and U + Create figs and subplots 473 | t_span_simu = np.linspace(0, T_sim-dt_sim, N_sim) 474 | t_span_plan = np.linspace(0, T_sim-dt_mpc, N_mpc) 475 | fig_u, ax_u = plt.subplots(nq, 1, figsize=(19.2,10.8), sharex='col') 476 | # For each joint 477 | for i in range(nq): 478 | 479 | if(PLOT_PREDICTIONS): 480 | 481 | # Extract state predictions of i^th joint 482 | u_pred_i = plot_data['u_pred'][:,:,i] 483 | 484 | # For each planning step in the trajectory 485 | for j in range(0, N_mpc, pred_plot_sampling): 486 | # Receding horizon = [j,j+N_h] 487 | t0_horizon = j*dt_mpc 488 | tspan_u_pred = np.linspace(t0_horizon, t0_horizon + T_h - dt_mpc, N_h) 489 | # Set up lists of (x,y) points for predicted positions and velocities 490 | points_u = np.array([tspan_u_pred, u_pred_i[j,:]]).transpose().reshape(-1,1,2) 491 | # Set up lists of segments 492 | segs_u = np.concatenate([points_u[:-1], points_u[1:]], axis=1) 493 | # Make collections segments 494 | cm = plt.get_cmap('Greys_r') 495 | lc_u = LineCollection(segs_u, cmap=cm, zorder=-1) 496 | lc_u.set_array(tspan_u_pred) 497 | # Customize 498 | lc_u.set_linestyle('-') 499 | lc_u.set_linewidth(1) 500 | # Plot collections 501 | ax_u[i].add_collection(lc_u) 502 | # Scatter to highlight points 503 | colors = np.r_[np.linspace(0.1, 1, N_h), 1] 504 | my_colors = cm(colors) 505 | ax_u[i].scatter(tspan_u_pred, u_pred_i[j,:], s=10, zorder=1, c=cm(np.r_[np.linspace(0.1, 1, N_h-1), 1] ), cmap=matplotlib.cm.Greys) #c='black' 506 | 507 | # Joint torques 508 | ax_u[i].plot(t_span_plan, plot_data['u_pred'][:,0,i], color='r', marker=None, linestyle='-', label='Optimal control u0*', alpha=0.6) 509 | ax_u[i].plot(t_span_plan, plot_data['u_des_MPC_RATE'][:,i], color='b', linestyle='-', marker='.', label='Predicted', alpha=0.1) 510 | # ax_u[i].plot(t_span_simu, plot_data['u_des_SIM_RATE'][:,i], color='y', linestyle='-', marker='.', label='Prediction (SIMU)', alpha=0.6) 511 | # ax_u[i].plot(t_span_simu, plot_data['grav'][:-1,i], color=[0.,1.,0.,0.], marker=None, linestyle='-.', label='Gravity torque', alpha=0.9) 512 | # Plot reference 513 | if('ctrlReg' or 'ctrlRegGrav' in plot_data['active_costs']): 514 | ax_u[i].plot(t_span_plan, plot_data['ctrl_ref'][:, i], linestyle='-.', color='k', marker=None, label='uReg_ref', alpha=0.5) 515 | ax_u[i].set_ylabel('$u_{}$'.format(i), fontsize=12) 516 | ax_u[i].yaxis.set_major_locator(plt.MaxNLocator(2)) 517 | ax_u[i].yaxis.set_major_formatter(plt.FormatStrFormatter('%.3e')) 518 | ax_u[i].grid(True) 519 | # Last x axis label 520 | if(i == nq-1): 521 | ax_u[i].set_xlabel('t (s)', fontsize=16) 522 | # LEgend 523 | handles_u, labels_u = ax_u[i].get_legend_handles_labels() 524 | fig_u.legend(handles_u, labels_u, loc='upper right', prop={'size': 16}) 525 | # Sup-y label 526 | fig_u.text(0.04, 0.5, 'Joint torque (Nm)', va='center', rotation='vertical', fontsize=16) 527 | # Titles 528 | fig_u.suptitle('Control = joint torques', size=18) 529 | # Save figs 530 | if(SAVE): 531 | figs = {'u': fig_u} 532 | if(SAVE_DIR is None): 533 | print("SAVE FIGURES IN HOME") 534 | SAVE_DIR = os.environ['HOME'] 535 | if(SAVE_NAME is None): 536 | SAVE_NAME = 'testfig' 537 | for name, fig in figs.items(): 538 | fig.savefig(SAVE_DIR + '/' +str(name) + '_' + SAVE_NAME +'.png') 539 | 540 | if(SHOW): 541 | plt.show() 542 | 543 | return fig_u 544 | 545 | # Plot end-eff data 546 | def plot_mpc_endeff_linear(plot_data, PLOT_PREDICTIONS=False, 547 | pred_plot_sampling=100, 548 | SAVE=False, SAVE_DIR=None, SAVE_NAME=None, 549 | SHOW=True, 550 | AUTOSCALE=False): 551 | ''' 552 | Plot endeff data (linear position and velocity) 553 | Input: 554 | plot_data : plotting data 555 | PLOT_PREDICTIONS : True or False 556 | pred_plot_sampling : plot every pred_plot_sampling prediction 557 | to avoid huge amount of plotted data 558 | ("1" = plot all) 559 | SAVE, SAVE_DIR, SAVE_NAME : save plots as .png 560 | SHOW : show plots 561 | AUTOSCALE : rescale y-axis of endeff plot 562 | based on maximum value taken 563 | ''' 564 | T_sim = plot_data['T_sim'] 565 | N_sim = plot_data['N_sim'] 566 | N_mpc = plot_data['N_mpc'] 567 | dt_mpc = plot_data['dt_mpc'] 568 | T_h = plot_data['T_h'] 569 | N_h = plot_data['N_h'] 570 | # Create time spans for X and U + Create figs and subplots 571 | t_span_simu = np.linspace(0, T_sim, N_sim+1) 572 | t_span_plan = np.linspace(0, T_sim, N_mpc+1) 573 | fig, ax = plt.subplots(3, 2, figsize=(19.2,10.8), sharex='col') 574 | # Plot endeff 575 | xyz = ['x', 'y', 'z'] 576 | for i in range(3): 577 | 578 | if(PLOT_PREDICTIONS): 579 | lin_pos_ee_pred_i = plot_data['lin_pos_ee_pred'][:, :, i] 580 | lin_vel_ee_pred_i = plot_data['lin_vel_ee_pred'][:, :, i] 581 | # For each planning step in the trajectory 582 | for j in range(0, N_mpc, pred_plot_sampling): 583 | # Receding horizon = [j,j+N_h] 584 | t0_horizon = j*dt_mpc 585 | tspan_x_pred = np.linspace(t0_horizon, t0_horizon + T_h, N_h+1) 586 | # Set up lists of (x,y) points for predicted positions 587 | points_p = np.array([tspan_x_pred, lin_pos_ee_pred_i[j,:]]).transpose().reshape(-1,1,2) 588 | points_v = np.array([tspan_x_pred, lin_vel_ee_pred_i[j,:]]).transpose().reshape(-1,1,2) 589 | # Set up lists of segments 590 | segs_p = np.concatenate([points_p[:-1], points_p[1:]], axis=1) 591 | segs_v = np.concatenate([points_v[:-1], points_v[1:]], axis=1) 592 | # Make collections segments 593 | cm = plt.get_cmap('Greys_r') 594 | lc_p = LineCollection(segs_p, cmap=cm, zorder=-1) 595 | lc_v = LineCollection(segs_v, cmap=cm, zorder=-1) 596 | lc_p.set_array(tspan_x_pred) 597 | lc_v.set_array(tspan_x_pred) 598 | # Customize 599 | lc_p.set_linestyle('-') 600 | lc_v.set_linestyle('-') 601 | lc_p.set_linewidth(1) 602 | lc_v.set_linewidth(1) 603 | # Plot collections 604 | ax[i,0].add_collection(lc_p) 605 | ax[i,1].add_collection(lc_v) 606 | # Scatter to highlight points 607 | colors = np.r_[np.linspace(0.1, 1, N_h), 1] 608 | my_colors = cm(colors) 609 | ax[i,0].scatter(tspan_x_pred, lin_pos_ee_pred_i[j,:], s=10, zorder=1, c=my_colors, cmap=matplotlib.cm.Greys) 610 | ax[i,1].scatter(tspan_x_pred, lin_vel_ee_pred_i[j,:], s=10, zorder=1, c=my_colors, cmap=matplotlib.cm.Greys) 611 | 612 | # EE position 613 | ax[i,0].plot(t_span_plan, plot_data['lin_pos_ee_des_MPC_RATE'][:,i], color='b', linestyle='-', marker='.', label='Predicted ', alpha=0.1) 614 | # ax[i,0].plot(t_span_simu, plot_data['lin_pos_ee_des_SIM_RATE'][:,i], color='y', linestyle='-', marker='.', label='Predicted (SIMU rate)', alpha=0.5) 615 | ax[i,0].plot(t_span_simu, plot_data['lin_pos_ee_mea'][:,i], 'r-', label='Measured (WITH noise)', linewidth=1, alpha=0.3) 616 | # Plot reference 617 | if('translation' in plot_data['active_costs']): 618 | ax[i,0].plot(t_span_plan[:-1], plot_data['lin_pos_ee_ref'][:,i], color='k', linestyle='-.', linewidth=2., label='Reference', alpha=0.9) 619 | ax[i,0].set_ylabel('$P^{EE}_%s$ (m)'%xyz[i], fontsize=16) 620 | ax[i,0].yaxis.set_major_locator(plt.MaxNLocator(2)) 621 | ax[i,0].yaxis.set_major_formatter(plt.FormatStrFormatter('%.3e')) 622 | ax[i,0].grid(True) 623 | 624 | # EE velocity 625 | ax[i,1].plot(t_span_plan, plot_data['lin_vel_ee_des_MPC_RATE'][:,i], color='b', linestyle='-', marker='.', label='Predicted ', alpha=0.1) 626 | # ax[i,1].plot(t_span_simu, plot_data['lin_vel_ee_des_SIM_RATE'][:,i], color='y', linestyle='-', marker='.', label='Predicted (SIMU rate)', alpha=0.5) 627 | ax[i,1].plot(t_span_simu, plot_data['lin_vel_ee_mea'][:,i], 'r-', label='Measured (WITH noise)', linewidth=1, alpha=0.3) 628 | # Plot reference 629 | if('velocity' in plot_data['active_costs']): 630 | ax[i,1].plot(t_span_plan, [0.]*(N_mpc+1), color='k', linestyle='-.', linewidth=2., label='Reference', alpha=0.9) 631 | ax[i,1].set_ylabel('$V^{EE}_%s$ (m)'%xyz[i], fontsize=16) 632 | ax[i,1].yaxis.set_major_locator(plt.MaxNLocator(2)) 633 | ax[i,1].yaxis.set_major_formatter(plt.FormatStrFormatter('%.3e')) 634 | ax[i,1].grid(True) 635 | 636 | 637 | # Align 638 | fig.align_ylabels(ax[:,0]) 639 | fig.align_ylabels(ax[:,1]) 640 | ax[i,0].set_xlabel('t (s)', fontsize=16) 641 | ax[i,1].set_xlabel('t (s)', fontsize=16) 642 | # Set ylim if any 643 | TOL = 1e-3 644 | if(AUTOSCALE): 645 | ax_p_ylim = 1.1*max(np.max(np.abs(plot_data['lin_pos_ee_mea'])), TOL) 646 | ax_v_ylim = 1.1*max(np.max(np.abs(plot_data['lin_vel_ee_mea'])), TOL) 647 | for i in range(3): 648 | ax[i,0].set_ylim(-ax_p_ylim, ax_p_ylim) 649 | ax[i,1].set_ylim(-ax_v_ylim, ax_v_ylim) 650 | 651 | handles_p, labels_p = ax[0,0].get_legend_handles_labels() 652 | fig.legend(handles_p, labels_p, loc='upper right', prop={'size': 16}) 653 | # Titles 654 | fig.suptitle('End-effector trajectories', size=18) 655 | # Save figs 656 | if(SAVE): 657 | figs = {'ee_lin': fig} 658 | if(SAVE_DIR is None): 659 | print("SAVE FIGURES IN HOME") 660 | SAVE_DIR = os.environ['HOME'] 661 | if(SAVE_NAME is None): 662 | SAVE_NAME = 'testfig' 663 | for name, fig in figs.items(): 664 | fig.savefig(SAVE_DIR + '/' +str(name) + '_' + SAVE_NAME +'.png') 665 | 666 | if(SHOW): 667 | plt.show() 668 | 669 | return fig, ax 670 | 671 | # Plot end-eff data 672 | def plot_mpc_endeff_angular(plot_data, PLOT_PREDICTIONS=False, 673 | pred_plot_sampling=100, 674 | SAVE=False, SAVE_DIR=None, SAVE_NAME=None, 675 | SHOW=True, 676 | AUTOSCALE=False): 677 | ''' 678 | Plot endeff data (angular position and velocity) 679 | Input: 680 | plot_data : plotting data 681 | PLOT_PREDICTIONS : True or False 682 | pred_plot_sampling : plot every pred_plot_sampling prediction 683 | to avoid huge amount of plotted data 684 | ("1" = plot all) 685 | SAVE, SAVE_DIR, SAVE_NAME : save plots as .png 686 | SHOW : show plots 687 | AUTOSCALE : rescale y-axis of endeff plot 688 | based on maximum value taken 689 | ''' 690 | T_sim = plot_data['T_sim'] 691 | N_sim = plot_data['N_sim'] 692 | N_mpc = plot_data['N_mpc'] 693 | dt_mpc = plot_data['dt_mpc'] 694 | T_h = plot_data['T_h'] 695 | N_h = plot_data['N_h'] 696 | # Create time spans for X and U + Create figs and subplots 697 | t_span_simu = np.linspace(0, T_sim, N_sim+1) 698 | t_span_plan = np.linspace(0, T_sim, N_mpc+1) 699 | fig, ax = plt.subplots(3, 2, figsize=(19.2,10.8), sharex='col') 700 | # Plot endeff 701 | xyz = ['x', 'y', 'z'] 702 | for i in range(3): 703 | 704 | if(PLOT_PREDICTIONS): 705 | ang_pos_ee_pred_i = plot_data['ang_pos_ee_pred'][:, :, i] 706 | ang_vel_ee_pred_i = plot_data['ang_vel_ee_pred'][:, :, i] 707 | # For each planning step in the trajectory 708 | for j in range(0, N_mpc, pred_plot_sampling): 709 | # Receding horizon = [j,j+N_h] 710 | t0_horizon = j*dt_mpc 711 | tspan_x_pred = np.linspace(t0_horizon, t0_horizon + T_h, N_h+1) 712 | # Set up lists of (x,y) points for predicted positions 713 | points_p = np.array([tspan_x_pred, ang_pos_ee_pred_i[j,:]]).transpose().reshape(-1,1,2) 714 | points_v = np.array([tspan_x_pred, ang_vel_ee_pred_i[j,:]]).transpose().reshape(-1,1,2) 715 | # Set up lists of segments 716 | segs_p = np.concatenate([points_p[:-1], points_p[1:]], axis=1) 717 | segs_v = np.concatenate([points_v[:-1], points_v[1:]], axis=1) 718 | # Make collections segments 719 | cm = plt.get_cmap('Greys_r') 720 | lc_p = LineCollection(segs_p, cmap=cm, zorder=-1) 721 | lc_v = LineCollection(segs_v, cmap=cm, zorder=-1) 722 | lc_p.set_array(tspan_x_pred) 723 | lc_v.set_array(tspan_x_pred) 724 | # Customize 725 | lc_p.set_linestyle('-') 726 | lc_v.set_linestyle('-') 727 | lc_p.set_linewidth(1) 728 | lc_v.set_linewidth(1) 729 | # Plot collections 730 | ax[i,0].add_collection(lc_p) 731 | ax[i,1].add_collection(lc_v) 732 | # Scatter to highlight points 733 | colors = np.r_[np.linspace(0.1, 1, N_h), 1] 734 | my_colors = cm(colors) 735 | ax[i,0].scatter(tspan_x_pred, ang_pos_ee_pred_i[j,:], s=10, zorder=1, c=my_colors, cmap=matplotlib.cm.Greys) 736 | ax[i,1].scatter(tspan_x_pred, ang_vel_ee_pred_i[j,:], s=10, zorder=1, c=my_colors, cmap=matplotlib.cm.Greys) 737 | 738 | # EE position 739 | ax[i,0].plot(t_span_plan, plot_data['ang_pos_ee_des_MPC_RATE'][:,i], color='b', linestyle='-', marker='.', label='Predicted ', alpha=0.1) 740 | # ax[i,0].plot(t_span_simu, plot_data['ang_pos_ee_des_SIM_RATE'][:,i], color='y', linestyle='-', marker='.', label='Predicted (SIMU rate)', alpha=0.5) 741 | ax[i,0].plot(t_span_simu, plot_data['ang_pos_ee_mea'][:,i], 'r-', label='Measured (WITH noise)', linewidth=1, alpha=0.3) 742 | # Plot reference 743 | if('rotation' in plot_data['active_costs']): 744 | ax[i,0].plot(t_span_plan[:-1], plot_data['ang_pos_ee_ref'][:,i], 'm-.', linewidth=2., label='Reference', alpha=0.9) 745 | ax[i,0].set_ylabel('$RPY^{EE}_%s$ (m)'%xyz[i], fontsize=16) 746 | ax[i,0].yaxis.set_major_locator(plt.MaxNLocator(2)) 747 | ax[i,0].yaxis.set_major_formatter(plt.FormatStrFormatter('%.3e')) 748 | ax[i,0].grid(True) 749 | 750 | # EE velocity 751 | ax[i,1].plot(t_span_plan, plot_data['ang_vel_ee_des_MPC_RATE'][:,i], color='b', linestyle='-', marker='.', label='Predicted ', alpha=0.1) 752 | # ax[i,1].plot(t_span_simu, plot_data['ang_vel_ee_des_SIM_RATE'][:,i], color='y', linestyle='-', marker='.', label='Predicted (SIMU rate)', alpha=0.5) 753 | ax[i,1].plot(t_span_simu, plot_data['ang_vel_ee_mea'][:,i], 'r-', label='Measured (WITH noise)', linewidth=1, alpha=0.3) 754 | # Plot reference 755 | if('velocity' in plot_data['active_costs']): 756 | ax[i,1].plot(t_span_plan, [0.]*(N_mpc+1), 'm-.', linewidth=2., label='Reference', alpha=0.9) 757 | ax[i,1].set_ylabel('$W^{EE}_%s$ (m)'%xyz[i], fontsize=16) 758 | ax[i,1].yaxis.set_major_locator(plt.MaxNLocator(2)) 759 | ax[i,1].yaxis.set_major_formatter(plt.FormatStrFormatter('%.3e')) 760 | ax[i,1].grid(True) 761 | 762 | 763 | # Align 764 | fig.align_ylabels(ax[:,0]) 765 | fig.align_ylabels(ax[:,1]) 766 | ax[i,0].set_xlabel('t (s)', fontsize=16) 767 | ax[i,1].set_xlabel('t (s)', fontsize=16) 768 | # Set ylim if any 769 | TOL = 1e-3 770 | if(AUTOSCALE): 771 | ax_p_ylim = 1.1*max(np.max(np.abs(plot_data['ang_pos_ee_mea'])), TOL) 772 | ax_v_ylim = 1.1*max(np.max(np.abs(plot_data['ang_vel_ee_mea'])), TOL) 773 | for i in range(3): 774 | ax[i,0].set_ylim(-ax_p_ylim, ax_p_ylim) 775 | ax[i,1].set_ylim(-ax_v_ylim, ax_v_ylim) 776 | 777 | handles_p, labels_p = ax[0,0].get_legend_handles_labels() 778 | fig.legend(handles_p, labels_p, loc='upper right', prop={'size': 16}) 779 | # Titles 780 | fig.suptitle('End-effector frame orientation (RPY) and angular velocity', size=18) 781 | # Save figs 782 | if(SAVE): 783 | figs = {'ee_ang': fig} 784 | if(SAVE_DIR is None): 785 | print("SAVE FIGURES IN HOME") 786 | SAVE_DIR = os.environ['HOME'] 787 | if(SAVE_NAME is None): 788 | SAVE_NAME = 'testfig' 789 | for name, fig in figs.items(): 790 | fig.savefig(SAVE_DIR + '/' +str(name) + '_' + SAVE_NAME +'.png') 791 | 792 | if(SHOW): 793 | plt.show() 794 | 795 | return fig, ax 796 | 797 | # Plot end-eff data 798 | def plot_mpc_force(plot_data, PLOT_PREDICTIONS=False, 799 | pred_plot_sampling=100, 800 | SAVE=False, SAVE_DIR=None, SAVE_NAME=None, 801 | SHOW=True, 802 | AUTOSCALE=False): 803 | ''' 804 | Plot EE force data 805 | Input: 806 | plot_data : plotting data 807 | PLOT_PREDICTIONS : True or False 808 | pred_plot_sampling : plot every pred_plot_sampling prediction 809 | to avoid huge amount of plotted data 810 | ("1" = plot all) 811 | SAVE, SAVE_DIR, SAVE_NAME : save plots as .png 812 | SHOW : show plots 813 | AUTOSCALE : rescale y-axis of endeff plot 814 | based on maximum value taken 815 | ''' 816 | T_sim = plot_data['T_sim'] 817 | N_sim = plot_data['N_sim'] 818 | N_mpc = plot_data['N_mpc'] 819 | dt_mpc = plot_data['dt_mpc'] 820 | dt_sim = plot_data['dt_sim'] 821 | T_h = plot_data['T_h'] 822 | N_h = plot_data['N_h'] 823 | # Create time spans for X and U + Create figs and subplots 824 | t_span_simu = np.linspace(0, T_sim - dt_sim, N_sim) 825 | t_span_plan = np.linspace(0, T_sim - dt_mpc, N_mpc) 826 | fig, ax = plt.subplots(3, 2, figsize=(19.2,10.8), sharex='col') 827 | # Plot endeff 828 | xyz = ['x', 'y', 'z'] 829 | for i in range(3): 830 | 831 | if(PLOT_PREDICTIONS): 832 | f_ee_pred_i = plot_data['f_ee_pred'][:, :, i] 833 | # For each planning step in the trajectory 834 | for j in range(0, N_mpc, pred_plot_sampling): 835 | # Receding horizon = [j,j+N_h] 836 | t0_horizon = j*dt_mpc 837 | tspan_x_pred = np.linspace(t0_horizon, t0_horizon + T_h - dt_mpc, N_h) 838 | # Set up lists of (x,y) points for predicted positions 839 | points_f = np.array([tspan_x_pred, f_ee_pred_i[j,:]]).transpose().reshape(-1,1,2) 840 | # Set up lists of segments 841 | segs_f = np.concatenate([points_f[:-1], points_f[1:]], axis=1) 842 | # Make collections segments 843 | cm = plt.get_cmap('Greys_r') 844 | lc_f = LineCollection(segs_f, cmap=cm, zorder=-1) 845 | lc_f.set_array(tspan_x_pred) 846 | # Customize 847 | lc_f.set_linestyle('-') 848 | lc_f.set_linewidth(1) 849 | # Plot collections 850 | ax[i,0].add_collection(lc_f) 851 | # Scatter to highlight points 852 | colors = np.r_[np.linspace(0.1, 1, N_h-1), 1] 853 | my_colors = cm(colors) 854 | ax[i,0].scatter(tspan_x_pred, f_ee_pred_i[j,:], s=10, zorder=1, c=my_colors, cmap=matplotlib.cm.Greys) 855 | 856 | # EE linear force 857 | ax[i,0].plot(t_span_plan, plot_data['f_ee_des_MPC_RATE'][:,i], color='b', linestyle='-', marker='.', label='Predicted ', alpha=0.1) 858 | # ax[i,0].plot(t_span_simu, plot_data['f_ee_des_SIM_RATE'][:,i], color='y', linestyle='-', marker='.', label='Predicted (SIMU rate)', alpha=0.5) 859 | ax[i,0].plot(t_span_simu, plot_data['f_ee_mea'][:,i], 'r-', label='Measured', linewidth=2, alpha=0.6) 860 | # Plot reference 861 | if('force' in plot_data['active_costs']): 862 | ax[i,0].plot(t_span_plan, plot_data['f_ee_ref'][:,i], color=[0.,1.,0.,0.], linestyle='-.', linewidth=2., label='Reference', alpha=0.9) 863 | ax[i,0].set_ylabel('$\\lambda^{EE}_%s$ (N)'%xyz[i], fontsize=16) 864 | ax[i,0].yaxis.set_major_locator(plt.MaxNLocator(2)) 865 | ax[i,0].yaxis.set_major_formatter(plt.FormatStrFormatter('%.3e')) 866 | ax[i,0].grid(True) 867 | 868 | # EE angular force 869 | ax[i,1].plot(t_span_plan, plot_data['f_ee_des_MPC_RATE'][:,3+i], color='b', linestyle='-', marker='.', label='Predicted ', alpha=0.1) 870 | # ax[i,1].plot(t_span_simu, plot_data['f_ee_des_SIM_RATE'][:,3+i], color='y', linestyle='-', marker='.', label='Predicted (SIMU rate)', alpha=0.5) 871 | ax[i,1].plot(t_span_simu, plot_data['f_ee_mea'][:,3+i], 'r-', label='Measured', linewidth=2, alpha=0.6) 872 | # Plot reference 873 | if('force' in plot_data['active_costs']): 874 | ax[i,1].plot(t_span_plan, plot_data['f_ee_ref'][:,3+i], color=[0.,1.,0.,0.], linestyle='-.', linewidth=2., label='Reference', alpha=0.9) 875 | ax[i,1].set_ylabel('$\\tau^{EE}_%s$ (Nm)'%xyz[i], fontsize=16) 876 | ax[i,1].yaxis.set_major_locator(plt.MaxNLocator(2)) 877 | ax[i,1].yaxis.set_major_formatter(plt.FormatStrFormatter('%.3e')) 878 | ax[i,1].grid(True) 879 | 880 | # Align 881 | fig.align_ylabels(ax[:,0]) 882 | fig.align_ylabels(ax[:,1]) 883 | ax[i,0].set_xlabel('t (s)', fontsize=16) 884 | ax[i,1].set_xlabel('t (s)', fontsize=16) 885 | # Set ylim if any 886 | TOL = 1e-3 887 | if(AUTOSCALE): 888 | ax_ylim = 1.1*max( np.max(np.abs(plot_data['f_ee_mea'])), TOL ) 889 | ax_ylim = 1.1*max( np.max(np.abs(plot_data['f_ee_mea'])), TOL ) 890 | for i in range(3): 891 | ax[i,0].set_ylim(-ax_ylim, ax_ylim) 892 | # ax[i,0].set_ylim(-30, 10) 893 | ax[i,1].set_ylim(-ax_ylim, ax_ylim) 894 | 895 | handles_p, labels_p = ax[0,0].get_legend_handles_labels() 896 | fig.legend(handles_p, labels_p, loc='upper right', prop={'size': 16}) 897 | # Titles 898 | fig.suptitle('End-effector forces', size=18) 899 | # Save figs 900 | if(SAVE): 901 | figs = {'f': fig} 902 | if(SAVE_DIR is None): 903 | print("SAVE FIGURES IN HOME") 904 | SAVE_DIR = os.environ['HOME'] 905 | if(SAVE_NAME is None): 906 | SAVE_NAME = 'testfig' 907 | for name, fig in figs.items(): 908 | fig.savefig(SAVE_DIR + '/' +str(name) + '_' + SAVE_NAME +'.png') 909 | 910 | if(SHOW): 911 | plt.show() 912 | 913 | return fig, ax 914 | -------------------------------------------------------------------------------- /ocp_kuka_avoid.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Example script : Crocoddyl OCP with KUKA arm with KUKA arm 3 | static target reaching task while avoiding an obstacle. 4 | Based from the mpc_kuka_reaching.py code. 5 | 6 | Written by Arthur Haffemayer. 7 | ''' 8 | 9 | import crocoddyl 10 | import numpy as np 11 | import ocp_utils 12 | import pin_utils 13 | import hppfcl 14 | import pinocchio as pin 15 | import mim_solvers 16 | try: 17 | from colmpc import ResidualDistanceCollision 18 | except: 19 | print( 20 | """It needs to be built with COLMPC https://github.com/agimus-project/colmpc 21 | and HPPFCL the devel branch for it to work.""" 22 | ) 23 | # # # # # # # # # # # # # 24 | ### LOAD ROBOT MODEL ### 25 | # # # # # # # # # # # # # 26 | 27 | from mim_robots.robot_loader import load_pinocchio_wrapper 28 | robot = load_pinocchio_wrapper("iiwa_convex") 29 | 30 | model = robot.model 31 | cmodel = pin_utils.transform_model_into_capsules(robot.collision_model) 32 | nq = model.nq; nv = model.nv; nu = nq; nx = nq+nv 33 | q0 = np.array([0.1, 0.7, 0., 0.7, -0.5, 1.5, 0.]) 34 | v0 = np.zeros(nv) 35 | x0 = np.concatenate([q0, v0]) 36 | robot.framesForwardKinematics(q0) 37 | robot.computeJointJacobians(q0) 38 | 39 | # # # # # # # # # # # # 40 | ### ADDING OBSTACLE ### 41 | # # # # # # # # # # # # 42 | 43 | OBSTACLE_POSE = pin.SE3(pin.utils.rotate("x", np.pi), np.array([-0.2, 0.15, 0.7])) 44 | OBSTACLE = hppfcl.Sphere(1e-1) 45 | OBSTACLE_GEOM_OBJECT = pin.GeometryObject( 46 | "obstacle", 47 | 0, 48 | 0, 49 | OBSTACLE, 50 | OBSTACLE_POSE, 51 | ) 52 | ID_OBSTACLE = cmodel.addGeometryObject( 53 | OBSTACLE_GEOM_OBJECT 54 | ) 55 | 56 | shapes_avoiding_collision_with_obstacles = [ 57 | "L3_capsule_0", 58 | "L4_capsule_0", 59 | "L4_capsule_1", 60 | "L5_capsule_0", 61 | "L6_capsule_0", 62 | "L6_capsule_1", 63 | "L7_capsule_0", 64 | ] 65 | 66 | # Adding the collisions pairs to the collision model 67 | for shape in shapes_avoiding_collision_with_obstacles: 68 | cmodel.addCollisionPair( 69 | pin.CollisionPair( 70 | cmodel.getGeometryId(shape), 71 | cmodel.getGeometryId("obstacle"), 72 | ) 73 | ) 74 | 75 | # # # # # # # # # # # # # # # 76 | ### SETUP CROCODDYL OCP ### 77 | # # # # # # # # # # # # # # # 78 | 79 | # State and actuation model 80 | state = crocoddyl.StateMultibody(model) 81 | actuation = crocoddyl.ActuationModelFull(state) 82 | 83 | # Running and terminal cost models 84 | runningCostModel = crocoddyl.CostModelSum(state) 85 | terminalCostModel = crocoddyl.CostModelSum(state) 86 | 87 | # Constraint model managers 88 | runningConstraintModelManager = crocoddyl.ConstraintModelManager(state, actuation.nu) 89 | terminalConstraintModelManager = crocoddyl.ConstraintModelManager(state, actuation.nu) 90 | 91 | # Create collision avoidance constraints 92 | 93 | # Creating the residual 94 | if len(cmodel.collisionPairs) != 0: 95 | for col_idx in range(len(cmodel.collisionPairs)): 96 | # obstacleDistanceResidual = ResidualCollision( 97 | # state, cmodel, cdata, col_idx 98 | # ) 99 | obstacleDistanceResidual = ResidualDistanceCollision(state, 7, cmodel, col_idx) 100 | 101 | # Creating the inequality constraint 102 | constraint = crocoddyl.ConstraintModelResidual( 103 | state, 104 | obstacleDistanceResidual, 105 | np.array([2e-2]), 106 | np.array([np.inf]), 107 | ) 108 | 109 | # Adding the constraint to the constraint manager 110 | runningConstraintModelManager.addConstraint("col_" + str(col_idx), constraint) 111 | terminalConstraintModelManager.addConstraint("col_term_" + str(col_idx), constraint) 112 | 113 | # Create cost terms 114 | # Control regularization cost 115 | uResidual = crocoddyl.ResidualModelControlGrav(state) 116 | uRegCost = crocoddyl.CostModelResidual(state, uResidual) 117 | # State regularization cost 118 | xResidual = crocoddyl.ResidualModelState(state, x0) 119 | xRegCost = crocoddyl.CostModelResidual(state, xResidual) 120 | # endeff frame translation cost 121 | endeff_frame_id = model.getFrameId("contact") 122 | endeff_translation = np.array([-0.4, 0.3, 0.7]) # move endeff +10 cm along x in WORLD frame 123 | frameTranslationResidual = crocoddyl.ResidualModelFrameTranslation(state, endeff_frame_id, endeff_translation) 124 | frameTranslationCost = crocoddyl.CostModelResidual(state, frameTranslationResidual) 125 | 126 | 127 | # Add costs 128 | runningCostModel.addCost("stateReg", xRegCost, 1e-1) 129 | runningCostModel.addCost("ctrlRegGrav", uRegCost, 1e-4) 130 | runningCostModel.addCost("translation", frameTranslationCost, 10) 131 | terminalCostModel.addCost("stateReg", xRegCost, 1e-1) 132 | terminalCostModel.addCost("translation", frameTranslationCost, 10) 133 | 134 | # Create Differential Action Model (DAM), i.e. continuous dynamics and cost functions 135 | running_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics(state, actuation, runningCostModel, runningConstraintModelManager) 136 | terminal_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics(state, actuation, terminalCostModel, terminalConstraintModelManager) 137 | 138 | # Create Integrated Action Model (IAM), i.e. Euler integration of continuous dynamics and cost 139 | dt = 1e-2 140 | runningModel = crocoddyl.IntegratedActionModelEuler(running_DAM, dt) 141 | terminalModel = crocoddyl.IntegratedActionModelEuler(terminal_DAM, 0.) 142 | 143 | # Optionally add armature to take into account actuator's inertia 144 | # runningModel.differential.armature = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.]) 145 | # terminalModel.differential.armature = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.]) 146 | 147 | # Create the shooting problem 148 | T = 250 149 | problem = crocoddyl.ShootingProblem(x0, [runningModel] * T, terminalModel) 150 | 151 | # Create solver + callbacks 152 | solver = mim_solvers.SolverCSQP(problem) 153 | xs_init = [x0 for i in range(T + 1)] 154 | us_init = solver.problem.quasiStatic(xs_init[:-1]) 155 | # Solve 156 | solver.termination_tolerance = 1e-4 157 | solver.with_callbacks = True 158 | solver.solve(xs_init, us_init, 100) 159 | solver.with_callbacks = False 160 | 161 | # Extract DDP data and plot 162 | ddp_data = ocp_utils.extract_ocp_data(solver, ee_frame_name='contact') 163 | ocp_utils.plot_ocp_results(ddp_data, which_plots='all', labels=None, markers=['.'], colors=['b'], sampling_plot=1, SHOW=True) 164 | 165 | # Display in Meshcat 166 | import time 167 | from pinocchio.visualize import MeshcatVisualizer 168 | 169 | robot.visual_model.addGeometryObject(OBSTACLE_GEOM_OBJECT) 170 | viz = MeshcatVisualizer(robot.model, robot.collision_model, robot.visual_model) 171 | 172 | viz.initViewer(open=True) 173 | viz.loadViewerModel() 174 | viz.display(q0) 175 | viz.displayCollisions(False) 176 | 177 | input("Press Enter to start the simulation") 178 | for t in range(T): 179 | viz.display(solver.xs[t][:model.nq]) 180 | time.sleep(0.1) 181 | -------------------------------------------------------------------------------- /ocp_kuka_contact.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Example script : Crocoddyl OCP with KUKA arm 3 | contact force task 4 | ''' 5 | import crocoddyl 6 | import pinocchio 7 | import numpy as np 8 | import pin_utils, ocp_utils 9 | 10 | # # # # # # # # # # # # # 11 | ### LOAD ROBOT MODEL ### 12 | # # # # # # # # # # # # # 13 | 14 | from mim_robots.robot_loader import load_pinocchio_wrapper 15 | robot = load_pinocchio_wrapper("iiwa") 16 | 17 | model = robot.model 18 | nq = model.nq; nv = model.nv; nu = nq; nx = nq+nv 19 | q0 = np.array([0.1, 0.7, 0., 0.7, -0.5, 1.5, 0.]) 20 | v0 = np.zeros(nv) 21 | x0 = np.concatenate([q0, v0]) 22 | robot.framesForwardKinematics(q0) 23 | robot.computeJointJacobians(q0) 24 | 25 | 26 | # # # # # # # # # # # # # # # 27 | ### SETUP CROCODDYL OCP ### 28 | # # # # # # # # # # # # # # # 29 | 30 | # State and actuation model 31 | state = crocoddyl.StateMultibody(model) 32 | actuation = crocoddyl.ActuationModelFull(state) 33 | 34 | # Running and terminal cost models 35 | runningCostModel = crocoddyl.CostModelSum(state) 36 | terminalCostModel = crocoddyl.CostModelSum(state) 37 | 38 | # Contact model 39 | contactModel = crocoddyl.ContactModelMultiple(state, actuation.nu) 40 | 41 | # Create 3D contact on the en-effector frame 42 | contact_frame_id = model.getFrameId("contact") 43 | contact_position = robot.data.oMf[contact_frame_id].copy() 44 | baumgarte_gains = np.array([0., 50.]) 45 | pinRef = pinocchio.LOCAL_WORLD_ALIGNED 46 | contact3d = crocoddyl.ContactModel6D(state, contact_frame_id, contact_position, pinRef, baumgarte_gains) 47 | 48 | # Populate contact model with contacts 49 | contactModel.addContact("contact", contact3d, active=True) 50 | 51 | # Create cost terms 52 | # Control regularization cost 53 | uResidual = crocoddyl.ResidualModelContactControlGrav(state) 54 | uRegCost = crocoddyl.CostModelResidual(state, uResidual) 55 | # State regularization cost 56 | xResidual = crocoddyl.ResidualModelState(state, x0) 57 | xRegCost = crocoddyl.CostModelResidual(state, xResidual) 58 | # End-effector frame force cost 59 | desired_wrench = np.array([0., 0., -20., 0., 0., 0.]) 60 | frameForceResidual = crocoddyl.ResidualModelContactForce(state, contact_frame_id, pinocchio.Force(desired_wrench), 6, actuation.nu) 61 | contactForceCost = crocoddyl.CostModelResidual(state, frameForceResidual) 62 | 63 | # Populate cost models with cost terms 64 | runningCostModel.addCost("stateReg", xRegCost, 1e-2) 65 | runningCostModel.addCost("ctrlRegGrav", uRegCost, 1e-4) 66 | runningCostModel.addCost("force", contactForceCost, 10.) 67 | terminalCostModel.addCost("stateReg", xRegCost, 1e-2) 68 | 69 | # Create Differential Action Model (DAM), i.e. continuous dynamics and cost functions 70 | running_DAM = crocoddyl.DifferentialActionModelContactFwdDynamics(state, actuation, contactModel, runningCostModel, inv_damping=0., enable_force=True) 71 | terminal_DAM = crocoddyl.DifferentialActionModelContactFwdDynamics(state, actuation, contactModel, terminalCostModel, inv_damping=0., enable_force=True) 72 | 73 | # Create Integrated Action Model (IAM), i.e. Euler integration of continuous dynamics and cost 74 | dt = 1e-2 75 | runningModel = crocoddyl.IntegratedActionModelEuler(running_DAM, dt) 76 | terminalModel = crocoddyl.IntegratedActionModelEuler(terminal_DAM, 0.) 77 | 78 | # Optionally add armature to take into account actuator's inertia 79 | runningModel.differential.armature = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.]) 80 | terminalModel.differential.armature = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.]) 81 | 82 | # Create the shooting problem 83 | T = 250 84 | problem = crocoddyl.ShootingProblem(x0, [runningModel] * T, terminalModel) 85 | 86 | # Create solver + callbacks 87 | ddp = crocoddyl.SolverFDDP(problem) 88 | ddp.setCallbacks([crocoddyl.CallbackLogger(), 89 | crocoddyl.CallbackVerbose()]) 90 | # Warm start : initial state + gravity compensation 91 | xs_init = [x0 for i in range(T+1)] 92 | us_init = ddp.problem.quasiStatic(xs_init[:-1]) 93 | 94 | # Solve 95 | ddp.solve(xs_init, us_init, maxiter=100, is_feasible=False) 96 | 97 | 98 | # Extract DDP data and plot 99 | ddp_data = {} 100 | ddp_data = ocp_utils.extract_ocp_data(ddp, ee_frame_name='contact', ct_frame_name='contact') 101 | ocp_utils.plot_ocp_results(ddp_data, which_plots='all', labels=None, markers=['.'], colors=['b'], sampling_plot=1, SHOW=True) 102 | 103 | 104 | # Display solution in Gepetto Viewer 105 | display = crocoddyl.GepettoDisplay(robot, frameNames=['contact']) 106 | display.displayFromSolver(ddp, factor=1) 107 | 108 | 109 | -------------------------------------------------------------------------------- /ocp_kuka_obstacle.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Example script : Crocoddyl OCP with KUKA arm 3 | static target reaching task 4 | ''' 5 | 6 | import crocoddyl 7 | import numpy as np 8 | import pinocchio as pin 9 | import ocp_utils 10 | try: 11 | import hppfcl 12 | except: 13 | print("Need to install hpp-flc") 14 | 15 | # # # # # # # # # # # # # 16 | ### LOAD ROBOT MODEL ### 17 | # # # # # # # # # # # # # 18 | 19 | 20 | from mim_robots.robot_loader import load_pinocchio_wrapper 21 | robot = load_pinocchio_wrapper("iiwa") 22 | 23 | model = robot.model 24 | nq = model.nq; nv = model.nv; nu = nq; nx = nq+nv 25 | q0 = np.array([0.1, 0.7, 0., 0.7, -0.5, 1.5, 0.]) 26 | v0 = np.zeros(nv) 27 | x0 = np.concatenate([q0, v0]) 28 | robot.framesForwardKinematics(q0) 29 | robot.computeJointJacobians(q0) 30 | 31 | # # # # # # # # # # # # # # # 32 | ### SETUP CROCODDYL OCP ### 33 | # # # # # # # # # # # # # # # 34 | 35 | # State and actuation model 36 | state = crocoddyl.StateMultibody(model) 37 | actuation = crocoddyl.ActuationModelFull(state) 38 | 39 | # Running and terminal cost models 40 | runningCostModel = crocoddyl.CostModelSum(state) 41 | terminalCostModel = crocoddyl.CostModelSum(state) 42 | 43 | 44 | # Create cost terms 45 | # Control regularization cost 46 | uResidual = crocoddyl.ResidualModelControlGrav(state) 47 | uRegCost = crocoddyl.CostModelResidual(state, uResidual) 48 | # State regularization cost 49 | xResidual = crocoddyl.ResidualModelState(state, x0) 50 | xRegCost = crocoddyl.CostModelResidual(state, xResidual) 51 | # endeff frame translation cost 52 | endeff_frame_id = model.getFrameId("contact") 53 | endeff_translation = np.array([-0.4, 0.3, 0.7]) # move endeff +10 cm along x in WORLD frame 54 | frameTranslationResidual = crocoddyl.ResidualModelFrameTranslation(state, endeff_frame_id, endeff_translation) 55 | frameTranslationCost = crocoddyl.CostModelResidual(state, frameTranslationResidual) 56 | 57 | # COLLISION COST 58 | # Create a capsule for the arm 59 | link_names = ["A1", "A2", "A3", "A4", "A5", "A6", "A7"] 60 | ig_link_names = [] 61 | for i,ln in enumerate(link_names): 62 | pin_link_id = model.getFrameId(ln) 63 | pin_joint_id = model.getJointId(ln) 64 | geomModel = pin.GeometryModel() 65 | ig_link_names = geomModel.addGeometryObject(pin.GeometryObject("arm_link_"+str(i), 66 | pin_link_id, 67 | model.frames[model.getFrameId(ln)].parent, 68 | hppfcl.Capsule(0, 0.5), 69 | pin.SE3.Identity())) 70 | 71 | # Create obstacle in the world 72 | obsPose = pin.SE3.Identity() 73 | obsPose.translation = np.array([0., 0., 1.]) 74 | obsObj = pin.GeometryObject("obstacle", 75 | model.getFrameId("universe"), 76 | model.frames[model.getFrameId("universe")].parent, 77 | hppfcl.Box(0.1, 0.1, 0.1), 78 | obsPose) 79 | ig_obs = geomModel.addGeometryObject(obsObj) 80 | geomModel.addCollisionPair(pin.CollisionPair(ig_link_names,ig_obs)) 81 | # Add collision cost to the OCP 82 | collision_radius = 0.2 83 | activationCollision = crocoddyl.ActivationModel2NormBarrier(3, collision_radius) 84 | residualCollision = crocoddyl.ResidualModelPairCollision(state, nu, geomModel, 0, pin_joint_id) 85 | costCollision = crocoddyl.CostModelResidual(state, activationCollision, residualCollision) 86 | 87 | 88 | # Add costs 89 | runningCostModel.addCost("stateReg", xRegCost, 1e-1) 90 | runningCostModel.addCost("ctrlRegGrav", uRegCost, 1e-4) 91 | runningCostModel.addCost("translation", frameTranslationCost, 10) 92 | runningCostModel.addCost("collision", costCollision, 100) 93 | terminalCostModel.addCost("stateReg", xRegCost, 1e-1) 94 | terminalCostModel.addCost("translation", frameTranslationCost, 10) 95 | terminalCostModel.addCost("collision", costCollision, 100) 96 | 97 | 98 | # Create Differential Action Model (DAM), i.e. continuous dynamics and cost functions 99 | running_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics(state, actuation, runningCostModel) 100 | terminal_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics(state, actuation, terminalCostModel) 101 | 102 | # Create Integrated Action Model (IAM), i.e. Euler integration of continuous dynamics and cost 103 | dt = 1e-2 104 | runningModel = crocoddyl.IntegratedActionModelEuler(running_DAM, dt) 105 | terminalModel = crocoddyl.IntegratedActionModelEuler(terminal_DAM, 0.) 106 | 107 | # Optionally add armature to take into account actuator's inertia 108 | runningModel.differential.armature = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.]) 109 | terminalModel.differential.armature = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.]) 110 | 111 | # Create the shooting problem 112 | T = 250 113 | problem = crocoddyl.ShootingProblem(x0, [runningModel] * T, terminalModel) 114 | 115 | 116 | # Create solver + callbacks 117 | ddp = crocoddyl.SolverFDDP(problem) 118 | ddp.setCallbacks([crocoddyl.CallbackLogger(), 119 | crocoddyl.CallbackVerbose()]) 120 | # Warm start : initial state + gravity compensation 121 | xs_init = [x0 for i in range(T+1)] 122 | us_init = ddp.problem.quasiStatic(xs_init[:-1]) 123 | 124 | # Solve 125 | ddp.solve(xs_init, us_init, maxiter=100, is_feasible=False) 126 | 127 | # # Extract DDP data and plot 128 | ddp_data = ocp_utils.extract_ocp_data(ddp, ee_frame_name='contact') 129 | ocp_utils.plot_ocp_results(ddp_data, which_plots='all', labels=None, markers=['.'], colors=['b'], sampling_plot=1, SHOW=True) 130 | 131 | # Display in Meshcat 132 | import time 133 | from pinocchio.visualize import MeshcatVisualizer 134 | robot.visual_model.addGeometryObject(obsObj) 135 | 136 | viz = MeshcatVisualizer(robot.model, robot.collision_model, robot.visual_model) 137 | viz.initViewer(open=True) 138 | viz.loadViewerModel() 139 | viz.display(q0) 140 | 141 | viz.displayCollisions(True) 142 | viz.displayVisuals(True) 143 | 144 | input("Press Enter to start the simulation") 145 | for t in range(T): 146 | viz.display(ddp.xs[t][:model.nq]) 147 | time.sleep(0.1) -------------------------------------------------------------------------------- /ocp_kuka_reaching.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Example script : Crocoddyl OCP with KUKA arm 3 | static target reaching task 4 | ''' 5 | 6 | import crocoddyl 7 | import numpy as np 8 | import ocp_utils 9 | 10 | # # # # # # # # # # # # # 11 | ### LOAD ROBOT MODEL ### 12 | # # # # # # # # # # # # # 13 | 14 | from mim_robots.robot_loader import load_pinocchio_wrapper 15 | robot = load_pinocchio_wrapper("iiwa") 16 | 17 | model = robot.model 18 | nq = model.nq; nv = model.nv; nu = nq; nx = nq+nv 19 | q0 = np.array([0.1, 0.7, 0., 0.7, -0.5, 1.5, 0.]) 20 | v0 = np.zeros(nv) 21 | x0 = np.concatenate([q0, v0]) 22 | robot.framesForwardKinematics(q0) 23 | robot.computeJointJacobians(q0) 24 | 25 | 26 | # # # # # # # # # # # # # # # 27 | ### SETUP CROCODDYL OCP ### 28 | # # # # # # # # # # # # # # # 29 | 30 | # State and actuation model 31 | state = crocoddyl.StateMultibody(model) 32 | actuation = crocoddyl.ActuationModelFull(state) 33 | 34 | # Running and terminal cost models 35 | runningCostModel = crocoddyl.CostModelSum(state) 36 | terminalCostModel = crocoddyl.CostModelSum(state) 37 | 38 | 39 | # Create cost terms 40 | # Control regularization cost 41 | uResidual = crocoddyl.ResidualModelControlGrav(state) 42 | uRegCost = crocoddyl.CostModelResidual(state, uResidual) 43 | # State regularization cost 44 | xResidual = crocoddyl.ResidualModelState(state, x0) 45 | xRegCost = crocoddyl.CostModelResidual(state, xResidual) 46 | # endeff frame translation cost 47 | endeff_frame_id = model.getFrameId("contact") 48 | endeff_translation = np.array([-0.4, 0.3, 0.7]) # move endeff +10 cm along x in WORLD frame 49 | frameTranslationResidual = crocoddyl.ResidualModelFrameTranslation(state, endeff_frame_id, endeff_translation) 50 | frameTranslationCost = crocoddyl.CostModelResidual(state, frameTranslationResidual) 51 | 52 | 53 | # Add costs 54 | runningCostModel.addCost("stateReg", xRegCost, 1e-1) 55 | runningCostModel.addCost("ctrlRegGrav", uRegCost, 1e-4) 56 | runningCostModel.addCost("translation", frameTranslationCost, 10) 57 | terminalCostModel.addCost("stateReg", xRegCost, 1e-1) 58 | terminalCostModel.addCost("translation", frameTranslationCost, 10) 59 | 60 | # Create Differential Action Model (DAM), i.e. continuous dynamics and cost functions 61 | running_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics(state, actuation, runningCostModel) 62 | terminal_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics(state, actuation, terminalCostModel) 63 | 64 | # Create Integrated Action Model (IAM), i.e. Euler integration of continuous dynamics and cost 65 | dt = 1e-2 66 | runningModel = crocoddyl.IntegratedActionModelEuler(running_DAM, dt) 67 | terminalModel = crocoddyl.IntegratedActionModelEuler(terminal_DAM, 0.) 68 | 69 | # Optionally add armature to take into account actuator's inertia 70 | runningModel.differential.armature = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.]) 71 | terminalModel.differential.armature = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.]) 72 | 73 | # Create the shooting problem 74 | T = 250 75 | problem = crocoddyl.ShootingProblem(x0, [runningModel] * T, terminalModel) 76 | 77 | # Create solver + callbacks 78 | ddp = crocoddyl.SolverFDDP(problem) 79 | ddp.setCallbacks([crocoddyl.CallbackLogger(), 80 | crocoddyl.CallbackVerbose()]) 81 | 82 | # Warm start : initial state + gravity compensation 83 | xs_init = [x0 for i in range(T+1)] 84 | us_init = ddp.problem.quasiStatic(xs_init[:-1]) 85 | 86 | # Solve 87 | ddp.solve(xs_init, us_init, maxiter=100, is_feasible=False) 88 | 89 | # Extract DDP data and plot 90 | ddp_data = ocp_utils.extract_ocp_data(ddp, ee_frame_name='contact') 91 | ocp_utils.plot_ocp_results(ddp_data, which_plots='all', labels=None, markers=['.'], colors=['b'], sampling_plot=1, SHOW=True) 92 | 93 | # Display solution in Gepetto Viewer 94 | display = crocoddyl.GepettoDisplay(robot) 95 | display.displayFromSolver(ddp, factor=1) 96 | 97 | -------------------------------------------------------------------------------- /ocp_utils.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | import pinocchio as pin 4 | import matplotlib.pyplot as plt 5 | import pin_utils 6 | 7 | 8 | # Extract data from solver 9 | def extract_ocp_data(ddp, ee_frame_name='contact', ct_frame_name='contact'): 10 | ''' 11 | Record relevant data from ddp solver in order to plot 12 | ee_frame_name = name of frame for which ee plots will be generated 13 | by default 'contact' as in KUKA urdf model (Tennis ball) 14 | ct_frame_name = name of frame for which force plots will be generated 15 | by default 'contact' as in KUKA urdf model (Tennis ball) 16 | ''' 17 | # Store data 18 | ddp_data = {} 19 | # OCP params 20 | ddp_data['T'] = ddp.problem.T 21 | ddp_data['dt'] = ddp.problem.runningModels[0].dt 22 | ddp_data['nq'] = ddp.problem.runningModels[0].state.nq 23 | ddp_data['nv'] = ddp.problem.runningModels[0].state.nv 24 | ddp_data['nu'] = ddp.problem.runningModels[0].differential.actuation.nu 25 | ddp_data['nx'] = ddp.problem.runningModels[0].state.nx 26 | # Pin model 27 | ddp_data['pin_model'] = ddp.problem.runningModels[0].differential.pinocchio 28 | ddp_data['armature'] = ddp.problem.runningModels[0].differential.armature 29 | ddp_data['frame_id'] = ddp_data['pin_model'].getFrameId(ee_frame_name) 30 | # Solution trajectories 31 | ddp_data['xs'] = ddp.xs 32 | ddp_data['us'] = ddp.us 33 | ddp_data['CONTACT_TYPE'] = None 34 | # Extract force at EE frame and contact info 35 | if(hasattr(ddp.problem.runningModels[0].differential, 'contacts')): 36 | # Get refs for contact model 37 | contactModelRef0 = ddp.problem.runningModels[0].differential.contacts.contacts[ct_frame_name].contact.reference 38 | # Case 6D contact (x,y,z,Ox,Oy,Oz) 39 | if(hasattr(contactModelRef0, 'rotation')): 40 | ddp_data['contact_rotation'] = [ddp.problem.runningModels[i].differential.contacts.contacts[ct_frame_name].contact.reference.rotation for i in range(ddp.problem.T)] 41 | ddp_data['contact_rotation'].append(ddp.problem.terminalModel.differential.contacts.contacts[ct_frame_name].contact.reference.rotation) 42 | ddp_data['contact_translation'] = [ddp.problem.runningModels[i].differential.contacts.contacts[ct_frame_name].contact.reference.translation for i in range(ddp.problem.T)] 43 | ddp_data['contact_translation'].append(ddp.problem.terminalModel.differential.contacts.contacts[ct_frame_name].contact.reference.translation) 44 | ddp_data['CONTACT_TYPE'] = '6D' 45 | # Case 3D contact (x,y,z) 46 | elif(np.size(contactModelRef0)==3): 47 | # Get ref translation for 3D 48 | ddp_data['contact_translation'] = [ddp.problem.runningModels[i].differential.contacts.contacts[ct_frame_name].contact.reference for i in range(ddp.problem.T)] 49 | ddp_data['contact_translation'].append(ddp.problem.terminalModel.differential.contacts.contacts[ct_frame_name].contact.reference) 50 | ddp_data['CONTACT_TYPE'] = '3D' 51 | # Case 1D contact (z) 52 | elif(np.size(contactModelRef0)==1): 53 | ddp_data['contact_translation'] = [ddp.problem.runningModels[i].differential.contacts.contacts[ct_frame_name].contact.reference for i in range(ddp.problem.T)] 54 | ddp_data['contact_translation'].append(ddp.problem.terminalModel.differential.contacts.contacts[ct_frame_name].contact.reference) 55 | ddp_data['CONTACT_TYPE'] = '1D' 56 | # Get contact force 57 | datas = [ddp.problem.runningDatas[i].differential.multibody.contacts.contacts[ct_frame_name] for i in range(ddp.problem.T)] 58 | # data.f = force exerted at parent joint expressed in WORLD frame (oMi) 59 | # express it in LOCAL contact frame using jMf 60 | ee_forces = [data.jMf.actInv(data.f).vector for data in datas] 61 | ddp_data['fs'] = [ee_forces[i] for i in range(ddp.problem.T)] 62 | # Extract refs for active costs 63 | # TODO : active costs may change along horizon : how to deal with that when plotting? 64 | ddp_data['active_costs'] = ddp.problem.runningModels[0].differential.costs.active.tolist() 65 | if('stateReg' in ddp_data['active_costs']): 66 | ddp_data['stateReg_ref'] = [ddp.problem.runningModels[i].differential.costs.costs['stateReg'].cost.residual.reference for i in range(ddp.problem.T)] 67 | ddp_data['stateReg_ref'].append(ddp.problem.terminalModel.differential.costs.costs['stateReg'].cost.residual.reference) 68 | if('ctrlReg' in ddp_data['active_costs']): 69 | ddp_data['ctrlReg_ref'] = [ddp.problem.runningModels[i].differential.costs.costs['ctrlReg'].cost.residual.reference for i in range(ddp.problem.T)] 70 | if('ctrlRegGrav' in ddp_data['active_costs']): 71 | ddp_data['ctrlRegGrav_ref'] = [pin_utils.get_u_grav(ddp.xs[i][:ddp_data['nq']], ddp_data['pin_model'], ddp_data['armature']) for i in range(ddp.problem.T)] 72 | if('stateLim' in ddp_data['active_costs']): 73 | ddp_data['stateLim_ub'] = [ddp.problem.runningModels[i].differential.costs.costs['stateLim'].cost.activation.bounds.ub for i in range(ddp.problem.T)] 74 | ddp_data['stateLim_lb'] = [ddp.problem.runningModels[i].differential.costs.costs['stateLim'].cost.activation.bounds.lb for i in range(ddp.problem.T)] 75 | ddp_data['stateLim_ub'].append(ddp.problem.terminalModel.differential.costs.costs['stateLim'].cost.activation.bounds.ub) 76 | ddp_data['stateLim_lb'].append(ddp.problem.terminalModel.differential.costs.costs['stateLim'].cost.activation.bounds.lb) 77 | if('ctrlLim' in ddp_data['active_costs']): 78 | ddp_data['ctrlLim_ub'] = [ddp.problem.runningModels[i].differential.costs.costs['ctrlLim'].cost.activation.bounds.ub for i in range(ddp.problem.T)] 79 | ddp_data['ctrlLim_lb'] = [ddp.problem.runningModels[i].differential.costs.costs['ctrlLim'].cost.activation.bounds.lb for i in range(ddp.problem.T)] 80 | ddp_data['ctrlLim_ub'].append(ddp.problem.runningModels[-1].differential.costs.costs['ctrlLim'].cost.activation.bounds.ub) 81 | ddp_data['ctrlLim_lb'].append(ddp.problem.runningModels[-1].differential.costs.costs['ctrlLim'].cost.activation.bounds.lb) 82 | if('placement' in ddp_data['active_costs']): 83 | ddp_data['translation_ref'] = [ddp.problem.runningModels[i].differential.costs.costs['placement'].cost.residual.reference.translation for i in range(ddp.problem.T)] 84 | ddp_data['translation_ref'].append(ddp.problem.terminalModel.differential.costs.costs['placement'].cost.residual.reference.translation) 85 | ddp_data['rotation_ref'] = [ddp.problem.runningModels[i].differential.costs.costs['placement'].cost.residual.reference.rotation for i in range(ddp.problem.T)] 86 | ddp_data['rotation_ref'].append(ddp.problem.terminalModel.differential.costs.costs['placement'].cost.residual.reference.rotation) 87 | if('translation' in ddp_data['active_costs']): 88 | ddp_data['translation_ref'] = [ddp.problem.runningModels[i].differential.costs.costs['translation'].cost.residual.reference for i in range(ddp.problem.T)] 89 | ddp_data['translation_ref'].append(ddp.problem.terminalModel.differential.costs.costs['translation'].cost.residual.reference) 90 | if('velocity' in ddp_data['active_costs']): 91 | ddp_data['velocity_ref'] = [ddp.problem.runningModels[i].differential.costs.costs['velocity'].cost.residual.reference.vector for i in range(ddp.problem.T)] 92 | ddp_data['velocity_ref'].append(ddp.problem.terminalModel.differential.costs.costs['velocity'].cost.residual.reference.vector) 93 | # ddp_data['frame_id'] = ddp.problem.runningModels[0].differential.costs.costs['velocity'].cost.residual.id 94 | if('rotation' in ddp_data['active_costs']): 95 | ddp_data['rotation_ref'] = [ddp.problem.runningModels[i].differential.costs.costs['rotation'].cost.residual.reference for i in range(ddp.problem.T)] 96 | ddp_data['rotation_ref'].append(ddp.problem.terminalModel.differential.costs.costs['rotation'].cost.residual.reference) 97 | if('force' in ddp_data['active_costs']): 98 | ddp_data['force_ref'] = [ddp.problem.runningModels[i].differential.costs.costs['force'].cost.residual.reference.vector for i in range(ddp.problem.T)] 99 | return ddp_data 100 | 101 | 102 | # Plot from DDP solver 103 | def plot_ocp_results(DDP_DATA, which_plots='all', labels=None, markers=None, colors=None, sampling_plot=1, SHOW=False): 104 | ''' 105 | Plot ddp results from 1 or several DDP solvers 106 | X, U, EE trajs 107 | INPUT 108 | DDP_DATA : DDP data or list of ddp data (cf. data_utils.extract_ocp_data()) 109 | ''' 110 | if(type(DDP_DATA) != list): 111 | DDP_DATA = [DDP_DATA] 112 | if(labels==None): 113 | labels=[None for k in range(len(DDP_DATA))] 114 | if(markers==None): 115 | markers=[None for k in range(len(DDP_DATA))] 116 | if(colors==None): 117 | colors=[None for k in range(len(DDP_DATA))] 118 | for k,data in enumerate(DDP_DATA): 119 | # If last plot, make legend 120 | make_legend = False 121 | if(k+sampling_plot > len(DDP_DATA)-1): 122 | make_legend=True 123 | # Return figs and axes object in case need to overlay new plots 124 | if(k==0): 125 | if('x' in which_plots or which_plots =='all' or 'all' in which_plots): 126 | if('xs' in data.keys()): 127 | fig_x, ax_x = plot_ocp_state(data, label=labels[k], marker=markers[k], color=colors[k], MAKE_LEGEND=make_legend, SHOW=False) 128 | if('u' in which_plots or which_plots =='all' or 'all' in which_plots): 129 | if('us' in data.keys()): 130 | fig_u, ax_u = plot_ocp_control(data, label=labels[k], marker=markers[k], color=colors[k], MAKE_LEGEND=make_legend, SHOW=False) 131 | if('ee' in which_plots or which_plots =='all' or 'all' in which_plots): 132 | if('xs' in data.keys()): 133 | fig_ee_lin, ax_ee_lin = plot_ocp_endeff_linear(data, label=labels[k], marker=markers[k], color=colors[k], MAKE_LEGEND=make_legend, SHOW=False) 134 | fig_ee_ang, ax_ee_ang = plot_ocp_endeff_angular(data, label=labels[k], marker=markers[k], color=colors[k], MAKE_LEGEND=make_legend, SHOW=False) 135 | if('f' in which_plots or which_plots =='all' or 'all' in which_plots): 136 | if('fs' in data.keys()): 137 | fig_f, ax_f = plot_ocp_force(data, label=labels[k], marker=markers[k], color=colors[k], MAKE_LEGEND=make_legend, SHOW=False) 138 | else: 139 | if(k%sampling_plot==0): 140 | if('x' in which_plots or which_plots =='all' or 'all' in which_plots): 141 | if('xs' in data.keys()): 142 | plot_ocp_state(data, fig=fig_x, ax=ax_x, label=labels[k], marker=markers[k], color=colors[k], MAKE_LEGEND=make_legend, SHOW=False) 143 | if('u' in which_plots or which_plots =='all' or 'all' in which_plots): 144 | if('us' in data.keys()): 145 | plot_ocp_control(data, fig=fig_u, ax=ax_u, label=labels[k], marker=markers[k], color=colors[k], MAKE_LEGEND=make_legend, SHOW=False) 146 | if('ee' in which_plots or which_plots =='all' or 'all' in which_plots): 147 | if('xs' in data.keys()): 148 | plot_ocp_endeff_linear(data, fig=fig_ee_lin, ax=ax_ee_lin, label=labels[k], marker=markers[k], color=colors[k], MAKE_LEGEND=make_legend, SHOW=False) 149 | if('f' in which_plots or which_plots =='all' or 'all' in which_plots): 150 | if('fs' in data.keys()): 151 | plot_ocp_force(data, fig=fig_f, ax=ax_f, label=labels[k], marker=markers[k], color=colors[k], MAKE_LEGEND=make_legend, SHOW=False) 152 | if(SHOW): 153 | plt.show() 154 | 155 | # Record and return if user needs to overlay stuff 156 | fig = {} 157 | ax = {} 158 | if('x' in which_plots or which_plots =='all' or 'all' in which_plots): 159 | if('xs' in data.keys()): 160 | fig['x'] = fig_x 161 | ax['x'] = ax_x 162 | if('u' in which_plots or which_plots =='all' or 'all' in which_plots): 163 | if('us' in data.keys()): 164 | fig['u'] = fig_u 165 | ax['u'] = ax_u 166 | if('ee' in which_plots or which_plots =='all' or 'all' in which_plots): 167 | if('xs' in data.keys()): 168 | fig['ee_lin'] = fig_ee_lin 169 | ax['ee_lin'] = ax_ee_lin 170 | fig['ee_ang'] = fig_ee_ang 171 | ax['ee_ang'] = ax_ee_ang 172 | if('f' in which_plots or which_plots =='all' or 'all' in which_plots): 173 | if('fs' in data.keys()): 174 | fig['f'] = fig_f 175 | ax['f'] = ax_f 176 | 177 | return fig, ax 178 | 179 | def plot_ocp_state(ddp_data, fig=None, ax=None, label=None, marker=None, color=None, alpha=1., MAKE_LEGEND=False, SHOW=True): 180 | ''' 181 | Plot ddp results (state) 182 | ''' 183 | # Parameters 184 | N = ddp_data['T'] 185 | dt = ddp_data['dt'] 186 | nq = ddp_data['nq'] 187 | nv = ddp_data['nv'] 188 | # Extract trajectories 189 | x = np.array(ddp_data['xs']) 190 | q = x[:,:nq] 191 | v = x[:,nv:] 192 | # If state reg cost, 193 | if('stateReg' in ddp_data['active_costs']): 194 | x_reg_ref = np.array(ddp_data['stateReg_ref']) 195 | # Plots 196 | tspan = np.linspace(0, N*dt, N+1) 197 | if(ax is None or fig is None): 198 | fig, ax = plt.subplots(nq, 2, sharex='col') 199 | if(label is None): 200 | label='State' 201 | for i in range(nq): 202 | # Plot positions 203 | ax[i,0].plot(tspan, q[:,i], linestyle='-', marker=marker, label=label, color=color, alpha=alpha) 204 | 205 | # Plot joint position regularization reference 206 | if('stateReg' in ddp_data['active_costs']): 207 | handles, labels = ax[i,0].get_legend_handles_labels() 208 | if('reg_ref' in labels): 209 | handles.pop(labels.index('reg_ref')) 210 | ax[i,0].lines.pop(labels.index('reg_ref')) 211 | labels.remove('reg_ref') 212 | ax[i,0].plot(tspan, x_reg_ref[:,i], linestyle='-.', color='k', marker=None, label='reg_ref', alpha=0.5) 213 | ax[i,0].set_ylabel('$q_%s$'%i, fontsize=16) 214 | ax[i,0].yaxis.set_major_locator(plt.MaxNLocator(2)) 215 | ax[i,0].yaxis.set_major_formatter(plt.FormatStrFormatter('%.2e')) 216 | ax[i,0].grid(True) 217 | 218 | # Plot velocities 219 | ax[i,1].plot(tspan, v[:,i], linestyle='-', marker=marker, label=label, color=color, alpha=alpha) 220 | 221 | # Plot joint velocity regularization reference 222 | if('stateReg' in ddp_data['active_costs']): 223 | handles, labels = ax[i,1].get_legend_handles_labels() 224 | if('reg_ref' in labels): 225 | handles.pop(labels.index('reg_ref')) 226 | ax[i,1].lines.pop(labels.index('reg_ref')) 227 | labels.remove('reg_ref') 228 | ax[i,1].plot(tspan, x_reg_ref[:,nq+i], linestyle='-.', color='k', marker=None, label='reg_ref', alpha=0.5) 229 | 230 | # Labels, tick labels and grid 231 | ax[i,1].set_ylabel('$v_%s$'%i, fontsize=16) 232 | ax[i,1].yaxis.set_major_locator(plt.MaxNLocator(2)) 233 | ax[i,1].yaxis.set_major_formatter(plt.FormatStrFormatter('%.2e')) 234 | ax[i,1].grid(True) 235 | 236 | # Common x-labels + align 237 | ax[-1,0].set_xlabel('Time (s)', fontsize=16) 238 | ax[-1,1].set_xlabel('Time (s)', fontsize=16) 239 | fig.align_ylabels(ax[:, 0]) 240 | fig.align_ylabels(ax[:, 1]) 241 | 242 | if(MAKE_LEGEND): 243 | handles, labels = ax[0,0].get_legend_handles_labels() 244 | fig.legend(handles, labels, loc='upper right', prop={'size': 16}) 245 | fig.align_ylabels() 246 | fig.suptitle('State trajectories', size=18) 247 | if(SHOW): 248 | plt.show() 249 | return fig, ax 250 | 251 | def plot_ocp_control(ddp_data, fig=None, ax=None, label=None, marker=None, color=None, alpha=1., MAKE_LEGEND=False, SHOW=True): 252 | ''' 253 | Plot ddp results (control) 254 | ''' 255 | # Parameters 256 | N = ddp_data['T'] 257 | dt = ddp_data['dt'] 258 | nu = ddp_data['nu'] 259 | # Extract trajectory 260 | u = np.array(ddp_data['us']) 261 | if('ctrlReg' in ddp_data['active_costs']): 262 | ureg_ref = np.array(ddp_data['ctrlReg_ref']) 263 | if('ctrlRegGrav' in ddp_data['active_costs']): 264 | ureg_grav = np.array(ddp_data['ctrlRegGrav_ref']) 265 | 266 | tspan = np.linspace(0, N*dt-dt, N) 267 | if(ax is None or fig is None): 268 | fig, ax = plt.subplots(nu, 1, sharex='col') 269 | if(label is None): 270 | label='Control' 271 | 272 | for i in range(nu): 273 | # Plot optimal control trajectory 274 | ax[i].plot(tspan, u[:,i], linestyle='-', marker=marker, label=label, color=color, alpha=alpha) 275 | 276 | # Plot control regularization reference 277 | if('ctrlReg' in ddp_data['active_costs']): 278 | handles, labels = ax[i].get_legend_handles_labels() 279 | if('u_reg' in labels): 280 | handles.pop(labels.index('u_reg')) 281 | ax[i].lines.pop(labels.index('u_reg')) 282 | labels.remove('u_reg') 283 | ax[i].plot(tspan, ureg_ref[:,i], linestyle='-.', color='k', marker=None, label='u_reg', alpha=0.5) 284 | 285 | # Plot gravity compensation torque 286 | if('ctrlRegGrav' in ddp_data['active_costs']): 287 | handles, labels = ax[i].get_legend_handles_labels() 288 | if('grav(q)' in labels): 289 | handles.pop(labels.index('u_grav(q)')) 290 | ax[i].lines.pop(labels.index('u_grav(q)')) 291 | labels.remove('u_grav(q)') 292 | ax[i].plot(tspan, ureg_grav[:,i], linestyle='-.', color='m', marker=None, label='u_grav(q)', alpha=0.5) 293 | 294 | # Labels, tick labels, grid 295 | ax[i].set_ylabel('$u_%s$'%i, fontsize=16) 296 | ax[i].yaxis.set_major_locator(plt.MaxNLocator(2)) 297 | ax[i].yaxis.set_major_formatter(plt.FormatStrFormatter('%.2e')) 298 | ax[i].grid(True) 299 | 300 | # Set x label + align 301 | ax[-1].set_xlabel('Time (s)', fontsize=16) 302 | fig.align_ylabels(ax[:]) 303 | # Legend 304 | if(MAKE_LEGEND): 305 | handles, labels = ax[i].get_legend_handles_labels() 306 | fig.legend(handles, labels, loc='upper right', prop={'size': 16}) 307 | fig.suptitle('Control trajectories', size=18) 308 | if(SHOW): 309 | plt.show() 310 | return fig, ax 311 | 312 | def plot_ocp_endeff_linear(ddp_data, fig=None, ax=None, label=None, marker=None, color=None, alpha=1., 313 | MAKE_LEGEND=False, SHOW=True, AUTOSCALE=True): 314 | ''' 315 | Plot ddp results (endeff linear position, velocity) 316 | ''' 317 | # Parameters 318 | N = ddp_data['T'] 319 | dt = ddp_data['dt'] 320 | nq = ddp_data['nq'] 321 | nv = ddp_data['nv'] 322 | # Extract EE traj 323 | x = np.array(ddp_data['xs']) 324 | q = x[:,:nq] 325 | v = x[:,nq:nq+nv] 326 | lin_pos_ee = pin_utils.get_p_(q, ddp_data['pin_model'], ddp_data['frame_id']) 327 | lin_vel_ee = pin_utils.get_v_(q, v, ddp_data['pin_model'], ddp_data['frame_id']) 328 | # Cost reference frame translation if any, or initial one 329 | if('translation' in ddp_data['active_costs'] or 'placement' in ddp_data['active_costs']): 330 | lin_pos_ee_ref = np.array(ddp_data['translation_ref']) 331 | else: 332 | lin_pos_ee_ref = np.array([lin_pos_ee[0,:] for i in range(N+1)]) 333 | # Cost reference frame linear velocity if any, or initial one 334 | if('velocity' in ddp_data['active_costs']): 335 | lin_vel_ee_ref = np.array(ddp_data['velocity_ref'])[:,:3] # linear part 336 | else: 337 | lin_vel_ee_ref = np.array([lin_vel_ee[0,:] for i in range(N+1)]) 338 | # Contact reference translation if CONTACT 339 | if(ddp_data['CONTACT_TYPE'] is not None): 340 | lin_pos_ee_contact = np.array(ddp_data['contact_translation']) 341 | # Plots 342 | tspan = np.linspace(0, N*dt, N+1) 343 | if(ax is None or fig is None): 344 | fig, ax = plt.subplots(3, 2, sharex='col') 345 | if(label is None): 346 | label='OCP solution' 347 | xyz = ['x', 'y', 'z'] 348 | for i in range(3): 349 | # Plot EE position in WORLD frame 350 | ax[i,0].plot(tspan, lin_pos_ee[:,i], linestyle='-', marker=marker, label=label, color=color, alpha=alpha) 351 | # Plot EE target frame translation in WORLD frame 352 | if('translation' or 'placement' in ddp_data['active_costs']): 353 | handles, labels = ax[i,0].get_legend_handles_labels() 354 | if('reference' in labels): 355 | handles.pop(labels.index('reference')) 356 | ax[i,0].lines.pop(labels.index('reference')) 357 | labels.remove('reference') 358 | ax[i,0].plot(tspan, lin_pos_ee_ref[:,i], linestyle='--', color='k', marker=None, label='reference', alpha=0.5) 359 | # Plot CONTACT reference frame translation in WORLD frame 360 | if(ddp_data['CONTACT_TYPE'] is not None): 361 | handles, labels = ax[i,0].get_legend_handles_labels() 362 | if('Baumgarte stab. ref.' in labels): 363 | handles.pop(labels.index('Baumgarte stab. ref.')) 364 | ax[i,0].lines.pop(labels.index('Baumgarte stab. ref.')) 365 | labels.remove('Baumgarte stab. ref.') 366 | # Exception for 1D contact: plot only along z-axis 367 | if(ddp_data['CONTACT_TYPE']=='1D'): 368 | if(i==2): 369 | ax[i,0].plot(tspan, lin_pos_ee_contact, linestyle=':', color='r', marker=None, label='Baumgarte stab. ref.', alpha=0.3) 370 | else: 371 | ax[i,0].plot(tspan, lin_pos_ee_contact[:,i], linestyle=':', color='r', marker=None, label='Baumgarte stab. ref.', alpha=0.3) 372 | # Labels, tick labels, grid 373 | ax[i,0].set_ylabel('$P^{EE}_%s$ (m)'%xyz[i], fontsize=16) 374 | ax[i,0].yaxis.set_major_locator(plt.MaxNLocator(2)) 375 | ax[i,0].yaxis.set_major_formatter(plt.FormatStrFormatter('%.2e')) 376 | ax[i,0].grid(True) 377 | 378 | # Plot EE (linear) velocities in WORLD frame 379 | ax[i,1].plot(tspan, lin_vel_ee[:,i], linestyle='-', marker=marker, label=label, color=color, alpha=alpha) 380 | # Plot EE target frame (linear) velocity in WORLD frame 381 | if('velocity' in ddp_data['active_costs']): 382 | handles, labels = ax[i,1].get_legend_handles_labels() 383 | if('reference' in labels): 384 | handles.pop(labels.index('reference')) 385 | ax[i,1].lines.pop(labels.index('reference')) 386 | labels.remove('reference') 387 | ax[i,1].plot(tspan, lin_vel_ee_ref[:,i], linestyle='--', color='k', marker=None, label='reference', alpha=0.5) 388 | # Labels, tick labels, grid 389 | ax[i,1].set_ylabel('$V^{EE}_%s$ (m/s)'%xyz[i], fontsize=16) 390 | ax[i,1].yaxis.set_major_locator(plt.MaxNLocator(2)) 391 | ax[i,1].yaxis.set_major_formatter(plt.FormatStrFormatter('%.2e')) 392 | ax[i,1].grid(True) 393 | 394 | #x-label + align 395 | fig.align_ylabels(ax[:,0]) 396 | fig.align_ylabels(ax[:,1]) 397 | ax[i,0].set_xlabel('t (s)', fontsize=16) 398 | ax[i,1].set_xlabel('t (s)', fontsize=16) 399 | 400 | # Set ylim if any 401 | if(AUTOSCALE): 402 | TOL = 0.1 403 | ax_p_ylim = 1.1*max(np.max(np.abs(lin_pos_ee)), TOL) 404 | ax_v_ylim = 1.1*max(np.max(np.abs(lin_vel_ee)), TOL) 405 | for i in range(3): 406 | ax[i,0].set_ylim(lin_pos_ee_ref[0,i]-ax_p_ylim, lin_pos_ee_ref[0,i]+ax_p_ylim) 407 | ax[i,1].set_ylim(lin_vel_ee_ref[0,i]-ax_v_ylim, lin_vel_ee_ref[0,i]+ax_v_ylim) 408 | 409 | if(MAKE_LEGEND): 410 | handles, labels = ax[2,0].get_legend_handles_labels() 411 | fig.legend(handles, labels, loc='upper right', prop={'size': 16}) 412 | fig.suptitle('End-effector frame position and linear velocity', size=18) 413 | if(SHOW): 414 | plt.show() 415 | return fig, ax 416 | 417 | def plot_ocp_endeff_angular(ddp_data, fig=None, ax=None, label=None, marker=None, color=None, alpha=1., 418 | MAKE_LEGEND=False, SHOW=True, AUTOSCALE=False): 419 | ''' 420 | Plot ddp results (endeff angular position, velocity) 421 | ''' 422 | # Parameters 423 | N = ddp_data['T'] 424 | dt = ddp_data['dt'] 425 | nq = ddp_data['nq'] 426 | nv = ddp_data['nv'] 427 | # Extract EE traj 428 | x = np.array(ddp_data['xs']) 429 | q = x[:,:nq] 430 | v = x[:,nq:nq+nv] 431 | rpy_ee = pin_utils.get_rpy_(q, ddp_data['pin_model'], ddp_data['frame_id']) 432 | w_ee = pin_utils.get_w_(q, v, ddp_data['pin_model'], ddp_data['frame_id']) 433 | # Cost reference frame orientation if any, or initial one 434 | if('rotation' in ddp_data['active_costs'] or 'placement' in ddp_data['active_costs']): 435 | rpy_ee_ref = np.array([pin.utils.matrixToRpy(np.array(R)) for R in ddp_data['rotation_ref']]) 436 | else: 437 | rpy_ee_ref = np.array([rpy_ee[0,:] for i in range(N+1)]) 438 | # Cost reference angular velocity if any, or initial one 439 | if('velocity' in ddp_data['active_costs']): 440 | w_ee_ref = np.array(ddp_data['velocity_ref'])[:,3:] # angular part 441 | else: 442 | w_ee_ref = np.array([w_ee[0,:] for i in range(N+1)]) 443 | # Contact reference orientation (6D) 444 | if(ddp_data['CONTACT_TYPE']=='6D'): 445 | rpy_ee_contact = np.array([pin.utils.matrixToRpy(R) for R in ddp_data['contact_rotation']]) 446 | # Plots 447 | tspan = np.linspace(0, N*dt, N+1) 448 | if(ax is None or fig is None): 449 | fig, ax = plt.subplots(3, 2, sharex='col') 450 | if(label is None): 451 | label='OCP solution' 452 | xyz = ['x', 'y', 'z'] 453 | for i in range(3): 454 | # Plot EE orientation in WORLD frame 455 | ax[i,0].plot(tspan, rpy_ee[:,i], linestyle='-', marker=marker, label=label, color=color, alpha=alpha) 456 | 457 | # Plot EE target frame orientation in WORLD frame 458 | if('rotation' or 'placement' in ddp_data['active_costs']): 459 | handles, labels = ax[i,0].get_legend_handles_labels() 460 | if('reference' in labels): 461 | handles.pop(labels.index('reference')) 462 | ax[i,0].lines.pop(labels.index('reference')) 463 | labels.remove('reference') 464 | ax[i,0].plot(tspan, rpy_ee_ref[:,i], linestyle='--', color='k', marker=None, label='reference', alpha=0.5) 465 | 466 | # Plot CONTACT reference frame rotation in WORLD frame 467 | if(ddp_data['CONTACT_TYPE']=='6D'): 468 | handles, labels = ax[i,0].get_legend_handles_labels() 469 | if('contact' in labels): 470 | handles.pop(labels.index('contact')) 471 | ax[i,0].lines.pop(labels.index('contact')) 472 | labels.remove('contact') 473 | ax[i,0].plot(tspan, rpy_ee_contact[:,i], linestyle=':', color='r', marker=None, label='Baumgarte stab. ref.', alpha=0.3) 474 | 475 | # Labels, tick labels, grid 476 | ax[i,0].set_ylabel('$RPY^{EE}_%s$ (rad)'%xyz[i], fontsize=16) 477 | ax[i,0].yaxis.set_major_locator(plt.MaxNLocator(2)) 478 | ax[i,0].yaxis.set_major_formatter(plt.FormatStrFormatter('%.2e')) 479 | ax[i,0].grid(True) 480 | 481 | # Plot EE 'linear) velocities in WORLD frame 482 | ax[i,1].plot(tspan, w_ee[:,i], linestyle='-', marker=marker, label=label, color=color, alpha=alpha) 483 | 484 | # Plot EE target frame (linear) velocity in WORLD frame 485 | if('velocity' in ddp_data['active_costs']): 486 | handles, labels = ax[i,1].get_legend_handles_labels() 487 | if('reference' in labels): 488 | handles.pop(labels.index('reference')) 489 | ax[i,1].lines.pop(labels.index('reference')) 490 | labels.remove('reference') 491 | ax[i,1].plot(tspan, w_ee_ref[:,i], linestyle='--', color='k', marker=None, label='reference', alpha=0.5) 492 | 493 | # Labels, tick labels, grid 494 | ax[i,1].set_ylabel('$W^{EE}_%s$ (rad/s)'%xyz[i], fontsize=16) 495 | ax[i,1].yaxis.set_major_locator(plt.MaxNLocator(2)) 496 | ax[i,1].yaxis.set_major_formatter(plt.FormatStrFormatter('%.2e')) 497 | ax[i,1].grid(True) 498 | 499 | #x-label + align 500 | fig.align_ylabels(ax[:,0]) 501 | fig.align_ylabels(ax[:,1]) 502 | ax[i,0].set_xlabel('t (s)', fontsize=16) 503 | ax[i,1].set_xlabel('t (s)', fontsize=16) 504 | 505 | # Set ylim if any 506 | if(AUTOSCALE): 507 | TOL = 0.1 508 | ax_p_ylim = 1.1*max(np.max(np.abs(rpy_ee)), TOL) 509 | ax_v_ylim = 1.1*max(np.max(np.abs(w_ee)), TOL) 510 | for i in range(3): 511 | ax[i,0].set_ylim(-ax_p_ylim, +ax_p_ylim) 512 | ax[i,1].set_ylim(-ax_v_ylim, +ax_v_ylim) 513 | 514 | if(MAKE_LEGEND): 515 | handles, labels = ax[0,0].get_legend_handles_labels() 516 | fig.legend(handles, labels, loc='upper right', prop={'size': 16}) 517 | fig.suptitle('End-effector frame orientation and angular velocity', size=18) 518 | if(SHOW): 519 | plt.show() 520 | return fig, ax 521 | 522 | def plot_ocp_force(ddp_data, fig=None, ax=None, label=None, marker=None, color=None, alpha=1., 523 | MAKE_LEGEND=False, SHOW=True, AUTOSCALE=True): 524 | ''' 525 | Plot ddp results (force) 526 | ''' 527 | # Parameters 528 | N = ddp_data['T'] 529 | dt = ddp_data['dt'] 530 | # Extract EE traj 531 | f = np.array(ddp_data['fs']) 532 | f_ee_lin = f[:,:3] 533 | f_ee_ang = f[:,3:] 534 | # Get desired contact wrench (linear, angular) 535 | if('force' in ddp_data['active_costs']): 536 | f_ee_ref = np.array(ddp_data['force_ref']) 537 | else: 538 | f_ee_ref = np.zeros((N,6)) 539 | f_ee_lin_ref = f_ee_ref[:,:3] 540 | f_ee_ang_ref = f_ee_ref[:,3:] 541 | # Plots 542 | tspan = np.linspace(0, N*dt, N) 543 | if(ax is None or fig is None): 544 | fig, ax = plt.subplots(3, 2, sharex='col') 545 | if(label is None): 546 | label='End-effector force' 547 | xyz = ['x', 'y', 'z'] 548 | for i in range(3): 549 | # Plot contact linear wrench (force) in LOCAL frame 550 | ax[i,0].plot(tspan, f_ee_lin[:,i], linestyle='-', marker=marker, label=label, color=color, alpha=alpha) 551 | 552 | # Plot desired contact linear wrench (force) in LOCAL frame 553 | if('force' in ddp_data['active_costs']): 554 | handles, labels = ax[i,0].get_legend_handles_labels() 555 | if('reference' in labels): 556 | handles.pop(labels.index('reference')) 557 | ax[i,0].lines.pop(labels.index('reference')) 558 | labels.remove('reference') 559 | ax[i,0].plot(tspan, f_ee_lin_ref[:,i], linestyle='-.', color='k', marker=None, label='reference', alpha=0.5) 560 | 561 | # Labels, tick labels+ grid 562 | ax[i,0].set_ylabel('$\\lambda^{lin}_%s$ (N)'%xyz[i], fontsize=16) 563 | ax[i,0].yaxis.set_major_locator(plt.MaxNLocator(2)) 564 | ax[i,0].yaxis.set_major_formatter(plt.FormatStrFormatter('%.2e')) 565 | ax[i,0].grid(True) 566 | 567 | # Plot contact angular wrench (torque) in LOCAL frame 568 | ax[i,1].plot(tspan, f_ee_ang[:,i], linestyle='-', marker=marker, label=label, color=color, alpha=alpha) 569 | 570 | # Plot desired contact anguler wrench (torque) in LOCAL frame 571 | if('force' in ddp_data['active_costs']): 572 | handles, labels = ax[i,1].get_legend_handles_labels() 573 | if('reference' in labels): 574 | handles.pop(labels.index('reference')) 575 | ax[i,1].lines.pop(labels.index('reference')) 576 | labels.remove('reference') 577 | ax[i,1].plot(tspan, f_ee_ang_ref[:,i], linestyle='-.', color='k', marker=None, label='reference', alpha=0.5) 578 | 579 | # Labels, tick labels+ grid 580 | ax[i,1].set_ylabel('$\\lambda^{ang}_%s$ (Nm)'%xyz[i], fontsize=16) 581 | ax[i,1].yaxis.set_major_locator(plt.MaxNLocator(2)) 582 | ax[i,1].yaxis.set_major_formatter(plt.FormatStrFormatter('%.2e')) 583 | ax[i,1].grid(True) 584 | 585 | # x-label + align 586 | fig.align_ylabels(ax[:,0]) 587 | fig.align_ylabels(ax[:,1]) 588 | ax[i,0].set_xlabel('t (s)', fontsize=16) 589 | ax[i,1].set_xlabel('t (s)', fontsize=16) 590 | 591 | # Set ylim if any 592 | if(AUTOSCALE): 593 | TOL = 1e-1 594 | ax_lin_ylim = 1.1*max(np.max(np.abs(f_ee_lin)), TOL) 595 | ax_ang_ylim = 1.1*max(np.max(np.abs(f_ee_ang)), TOL) 596 | for i in range(3): 597 | ax[i,0].set_ylim(f_ee_lin_ref[0,i]-ax_lin_ylim, f_ee_lin_ref[0,i]+ax_lin_ylim) 598 | ax[i,1].set_ylim(f_ee_ang_ref[0,i]-ax_ang_ylim, f_ee_ang_ref[0,i]+ax_ang_ylim) 599 | 600 | if(MAKE_LEGEND): 601 | handles, labels = ax[0,0].get_legend_handles_labels() 602 | fig.legend(handles, labels, loc='upper right', prop={'size': 16}) 603 | fig.suptitle('End-effector forces: linear and angular', size=18) 604 | if(SHOW): 605 | plt.show() 606 | return fig, ax 607 | 608 | 609 | -------------------------------------------------------------------------------- /pin_utils.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | import pinocchio as pin 4 | import eigenpy 5 | from numpy.linalg import pinv 6 | import time 7 | import matplotlib.pyplot as plt 8 | import hppfcl 9 | 10 | RED = np.array([249, 136, 126, 125]) / 255 11 | 12 | # Rotate placement 13 | def rotate(se3_placement, rpy=[0., 0., 0.]): 14 | ''' 15 | Rotates se3_placement.rotation by rpy (LOCAL) 16 | input : 17 | se3_placement : pin.SE3 18 | rpy : RPY orientation in LOCAL frame 19 | RPY 20 | ''' 21 | se3_placement_rotated = se3_placement.copy() 22 | R = pin.rpy.rpyToMatrix(rpy[0], rpy[1], rpy[2]) 23 | se3_placement_rotated.rotation = se3_placement_rotated.rotation.copy().dot(R) 24 | return se3_placement_rotated 25 | 26 | 27 | 28 | # Get frame position 29 | def get_p(q, pin_robot, id_endeff): 30 | ''' 31 | Returns end-effector positions given q trajectory 32 | q : joint positions 33 | robot : pinocchio wrapper 34 | id_endeff : id of EE frame 35 | ''' 36 | return get_p_(q, pin_robot.model, id_endeff) 37 | 38 | def get_p_(q, model, id_endeff): 39 | ''' 40 | Returns end-effector positions given q trajectory 41 | q : joint positions 42 | model : pinocchio model 43 | id_endeff : id of EE frame 44 | ''' 45 | 46 | data = model.createData() 47 | if(type(q)==np.ndarray and len(q.shape)==1): 48 | pin.forwardKinematics(model, data, q) 49 | pin.updateFramePlacements(model, data) 50 | p = data.oMf[id_endeff].translation.T 51 | else: 52 | N = np.shape(q)[0] 53 | p = np.empty((N,3)) 54 | for i in range(N): 55 | pin.forwardKinematics(model, data, q[i]) 56 | pin.updateFramePlacements(model, data) 57 | p[i,:] = data.oMf[id_endeff].translation.T 58 | return p 59 | 60 | 61 | 62 | # Get frame linear velocity 63 | def get_v(q, dq, pin_robot, id_endeff, ref=pin.LOCAL): 64 | ''' 65 | Returns end-effector velocities given q,dq trajectory 66 | q : joint positions 67 | dq : joint velocities 68 | pin_robot : pinocchio wrapper 69 | id_endeff : id of EE frame 70 | ''' 71 | return get_v_(q, dq, pin_robot.model, id_endeff, ref) 72 | 73 | def get_v_(q, dq, model, id_endeff, ref=pin.LOCAL): 74 | ''' 75 | Returns end-effector velocities given q,dq trajectory 76 | q : joint positions 77 | dq : joint velocities 78 | model : pinocchio model 79 | id_endeff : id of EE frame 80 | ''' 81 | data = model.createData() 82 | if(len(q) != len(dq)): 83 | print("q and dq must have the same size !") 84 | if(type(q)==np.ndarray and len(q.shape)==1): 85 | # J = pin.computeFrameJacobian(model, data, q, id_endeff) 86 | # v = J.dot(dq)[:3] 87 | pin.forwardKinematics(model, data, q, dq) 88 | spatial_vel = pin.getFrameVelocity(model, data, id_endeff, ref) 89 | v = spatial_vel.linear 90 | else: 91 | N = np.shape(q)[0] 92 | v = np.empty((N,3)) 93 | for i in range(N): 94 | # J = pin.computeFrameJacobian(model, data, q[i,:], id_endeff) 95 | # v[i,:] = J.dot(dq[i])[:3] 96 | pin.forwardKinematics(model, data, q[i], dq[i]) 97 | spatial_vel = pin.getFrameVelocity(model, data, id_endeff, ref) 98 | v[i,:] = spatial_vel.linear 99 | return v 100 | 101 | 102 | 103 | # Get frame orientation (rotation) 104 | def get_R(q, pin_robot, id_endeff): 105 | ''' 106 | Returns end-effector rotation matrices given q trajectory 107 | q : joint positions 108 | pin_robot : pinocchio wrapper 109 | id_endeff : id of EE frame 110 | ''' 111 | return get_R_(q, pin_robot.model, id_endeff) 112 | 113 | def get_R_(q, model, id_endeff): 114 | ''' 115 | Returns end-effector rotation matrices given q trajectory 116 | q : joint positions 117 | model : pinocchio model 118 | id_endeff : id of EE frame 119 | Output : single 3x3 array (or list of 3x3 arrays) 120 | ''' 121 | data = model.createData() 122 | if(type(q)==np.ndarray and len(q.shape)==1): 123 | pin.framesForwardKinematics(model, data, q) 124 | R = data.oMf[id_endeff].rotation.copy() 125 | else: 126 | N = np.shape(q)[0] 127 | R = [] 128 | for i in range(N): 129 | pin.framesForwardKinematics(model, data, q[i]) 130 | R.append(data.oMf[id_endeff].rotation.copy()) 131 | return R 132 | 133 | 134 | 135 | # Get frame orientation (RPY) 136 | def get_rpy(q, pin_robot, id_endeff): 137 | ''' 138 | Returns RPY angles of end-effector frame given q trajectory 139 | q : joint positions 140 | model : pinocchio wrapper 141 | id_endeff : id of EE frame 142 | ''' 143 | return get_rpy_(q, pin_robot.model, id_endeff) 144 | 145 | def get_rpy_(q, model, id_endeff): 146 | ''' 147 | Returns RPY angles of end-effector frame given q trajectory 148 | q : joint positions 149 | model : pinocchio model 150 | id_endeff : id of EE frame 151 | ''' 152 | R = get_R_(q, model, id_endeff) 153 | if(type(R)==list): 154 | N = np.shape(q)[0] 155 | rpy = np.empty((N,3)) 156 | for i in range(N): 157 | rpy[i,:] = pin.rpy.matrixToRpy(R[i]) #%(2*np.pi) 158 | else: 159 | rpy = pin.rpy.matrixToRpy(R) #%(2*np.pi) 160 | return rpy 161 | 162 | 163 | 164 | # Get frame angular velocity 165 | def get_w(q, dq, pin_robot, id_endeff, ref=pin.LOCAL): 166 | ''' 167 | Returns end-effector angular velocity given q,dq trajectory 168 | q : joint positions 169 | dq : joint velocities 170 | pin_robot : pinocchio wrapper 171 | id_endeff : id of EE frame 172 | ''' 173 | return get_w_(q, dq, pin_robot.model, id_endeff, ref) 174 | 175 | def get_w_(q, dq, model, id_endeff, ref=pin.LOCAL): 176 | ''' 177 | Returns end-effector angular velocity given q,dq trajectory 178 | q : joint positions 179 | dq : joint velocities 180 | pin_robot : pinocchio wrapper 181 | id_endeff : id of EE frame 182 | ''' 183 | data = model.createData() 184 | if(len(q) != len(dq)): 185 | print("q and dq must have the same size !") 186 | if(type(q)==np.ndarray and len(q.shape)==1): 187 | pin.forwardKinematics(model, data, q, dq) 188 | spatial_vel = pin.getFrameVelocity(model, data, id_endeff, ref) 189 | w = spatial_vel.angular 190 | else: 191 | N = np.shape(q)[0] 192 | w = np.empty((N,3)) 193 | for i in range(N): 194 | pin.forwardKinematics(model, data, q[i], dq[i]) 195 | spatial_vel = pin.getFrameVelocity(model, data, id_endeff, ref) 196 | w[i,:] = spatial_vel.angular 197 | return w 198 | 199 | 200 | # Get gravity joint torque 201 | def get_u_grav(q, model, armature): 202 | ''' 203 | Return gravity torque at q 204 | ''' 205 | data = model.createData() 206 | data.M += np.diag(armature) 207 | return pin.computeGeneralizedGravity(model, data, q) 208 | 209 | 210 | # Transform collision model shapes of spheres and cylinders into capsules 211 | def transform_model_into_capsules(collision_model): 212 | """Modifying the collision model to transform the spheres/cylinders into capsules which makes it easier to have a fully constrained robot.""" 213 | collision_model_copy = collision_model.copy() 214 | list_names_capsules = [] 215 | 216 | # Going through all the goemetry objects in the collision model 217 | for geom_object in collision_model_copy.geometryObjects: 218 | if isinstance(geom_object.geometry, hppfcl.Cylinder): 219 | # Sometimes for one joint there are two cylinders, which need to be defined by two capsules for the same link. 220 | # Hence the name convention here. 221 | if (geom_object.name[:-1] + "capsule_0") in list_names_capsules: 222 | name = geom_object.name[:-1] + "capsule_" + "1" 223 | else: 224 | name = geom_object.name[:-1] + "capsule_" + "0" 225 | list_names_capsules.append(name) 226 | placement = geom_object.placement 227 | parentJoint = geom_object.parentJoint 228 | parentFrame = geom_object.parentFrame 229 | geometry = geom_object.geometry 230 | geom = pin.GeometryObject( 231 | name, 232 | parentFrame, 233 | parentJoint, 234 | hppfcl.Capsule(geometry.radius, geometry.halfLength), 235 | placement, 236 | ) 237 | geom.meshColor = RED 238 | collision_model.addGeometryObject(geom) 239 | collision_model.removeGeometryObject(geom_object.name) 240 | elif isinstance(geom_object.geometry, hppfcl.Sphere): 241 | collision_model.removeGeometryObject(geom_object.name) 242 | return collision_model 243 | -------------------------------------------------------------------------------- /test_riccati_gains_mpc.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Example script : MPC simulation with KUKA arm 3 | static target reaching task 4 | ''' 5 | 6 | import crocoddyl 7 | import numpy as np 8 | np.set_printoptions(precision=4, linewidth=180) 9 | 10 | import pin_utils, mpc_utils 11 | 12 | from mim_robots.pybullet.env import BulletEnvWithGround 13 | from mim_robots.robot_loader import load_bullet_wrapper 14 | import pybullet as p 15 | 16 | 17 | # # # # # # # # # # # # # # # # # # # 18 | ### LOAD ROBOT MODEL and SIMU ENV ###  19 | # # # # # # # # # # # # # # # # # # #  20 | # Simulation environment 21 | env = BulletEnvWithGround(p.GUI, dt=1e-3) 22 | # Robot simulator 23 | robot_simulator = load_bullet_wrapper("iiwa") 24 | # Extract robot model 25 | nq = robot_simulator.pin_robot.model.nq 26 | nv = robot_simulator.pin_robot.model.nv 27 | nu = nq; nx = nq+nv 28 | q0 = np.array([0.1, 0.7, 0., 0.7, -0.5, 1.5, 0.]) 29 | v0 = np.zeros(nv) 30 | x0 = np.concatenate([q0, v0]) 31 | # Add robot to simulation and initialize 32 | env.add_robot(robot_simulator) 33 | robot_simulator.reset_state(q0, v0) 34 | robot_simulator.forward_robot(q0, v0) 35 | print("[PyBullet] Created robot (id = "+str(robot_simulator.robotId)+")") 36 | 37 | 38 | # # # # # # # # # # # # # # # 39 | ### SETUP CROCODDYL OCP ### 40 | # # # # # # # # # # # # # # # 41 | # State and actuation model 42 | state = crocoddyl.StateMultibody(robot_simulator.pin_robot.model) 43 | actuation = crocoddyl.ActuationModelFull(state) 44 | # Running and terminal cost models 45 | runningCostModel = crocoddyl.CostModelSum(state) 46 | terminalCostModel = crocoddyl.CostModelSum(state) 47 | # Create cost terms 48 | # Control regularization cost 49 | # uResidual = crocoddyl.ResidualModelControlGrav(state) 50 | uResidual = crocoddyl.ResidualModelControlGrav(state) 51 | uRegCost = crocoddyl.CostModelResidual(state, uResidual) 52 | # State regularization cost 53 | xResidual = crocoddyl.ResidualModelState(state, x0) 54 | xRegCost = crocoddyl.CostModelResidual(state, xResidual) 55 | # endeff frame translation cost 56 | endeff_frame_id = robot_simulator.pin_robot.model.getFrameId("contact") 57 | # endeff_translation = robot.data.oMf[endeff_frame_id].translation.copy() 58 | endeff_translation = np.array([0.4, 0.3, 0.7]) # move endeff +10 cm along x in WORLD frame 59 | frameTranslationResidual = crocoddyl.ResidualModelFrameTranslation(state, endeff_frame_id, endeff_translation) 60 | frameTranslationCost = crocoddyl.CostModelResidual(state, frameTranslationResidual) 61 | # Add costs 62 | runningCostModel.addCost("stateReg", xRegCost, 1e-1) 63 | runningCostModel.addCost("ctrlRegGrav", uRegCost, 1e-3) 64 | runningCostModel.addCost("translation", frameTranslationCost, 0.5) 65 | terminalCostModel.addCost("stateReg", xRegCost, 1e-1) 66 | terminalCostModel.addCost("translation", frameTranslationCost, 0.5) 67 | # Create Differential Action Model (DAM), i.e. continuous dynamics and cost functions 68 | running_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics(state, actuation, runningCostModel) 69 | terminal_DAM = crocoddyl.DifferentialActionModelFreeFwdDynamics(state, actuation, terminalCostModel) 70 | # Create Integrated Action Model (IAM), i.e. Euler integration of continuous dynamics and cost 71 | dt = 0.04 72 | runningModel = crocoddyl.IntegratedActionModelEuler(running_DAM, dt) 73 | terminalModel = crocoddyl.IntegratedActionModelEuler(terminal_DAM, 0.) 74 | # Optionally add armature to take into account actuator's inertia 75 | # runningModel.differential.armature = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.]) 76 | # terminalModel.differential.armature = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.]) 77 | # Create the shooting problem 78 | T = 100 79 | problem = crocoddyl.ShootingProblem(x0, [runningModel] * T, terminalModel) 80 | # Create solver + callbacks 81 | ddp = crocoddyl.SolverFDDP(problem) 82 | # ddp.setCallbacks([crocoddyl.CallbackLogger(), 83 | # crocoddyl.CallbackVerbose()]) 84 | # Warm start : initial state + gravity compensation 85 | xs_init = [x0 for i in range(T+1)] 86 | us_init = ddp.problem.quasiStatic(xs_init[:-1]) 87 | # Solve 88 | ddp.solve(xs_init, us_init, maxiter=100, is_feasible=False) 89 | 90 | 91 | 92 | # # # # # # # # # # # # 93 | ### MPC SIMULATION ### 94 | # # # # # # # # # # # # 95 | # OCP parameters 96 | ocp_params = {} 97 | ocp_params['N_h'] = T 98 | ocp_params['dt'] = dt 99 | ocp_params['maxiter'] = 100 100 | ocp_params['pin_model'] = robot_simulator.pin_robot.model 101 | ocp_params['armature'] = runningModel.differential.armature 102 | ocp_params['id_endeff'] = endeff_frame_id 103 | ocp_params['active_costs'] = ddp.problem.runningModels[0].differential.costs.active.tolist() 104 | # Simu parameters 105 | sim_params = {} 106 | sim_params['sim_freq'] = int(1./env.dt) 107 | sim_params['mpc_freq'] = 1 #100 #25 108 | sim_params['T_sim'] = 10. 109 | log_rate = 100 110 | # Initialize simulation data 111 | sim_data = mpc_utils.init_sim_data(sim_params, ocp_params, x0) 112 | # Display target 113 | mpc_utils.display_ball(endeff_translation, RADIUS=.05, COLOR=[1.,0.,0.,.6]) 114 | # Simulate 115 | mpc_cycle = 0 116 | ctrl_cycle = 0 117 | # RICCATI_TYPE = 1 118 | # INTERP_FEEDFORWARD = False 119 | # Kv = 0.*np.eye(nv) 120 | NAIVE = False 121 | T_force = 0 122 | dt_force = 100 123 | linkIndex = 7 124 | force = np.array([0., 100., 0]) 125 | 126 | for i in range(sim_data['N_sim']): 127 | # External push 128 | if T_force <= i and i < T_force + dt_force : 129 | position = p.getLinkState(robot_simulator.robotId, linkIndex)[0] 130 | p.applyExternalForce(objectUniqueId=robot_simulator.robotId, linkIndex=linkIndex, forceObj=force, posObj=position, flags=p.WORLD_FRAME) 131 | 132 | if(i%log_rate==0): 133 | print("\n SIMU step "+str(i)+"/"+str(sim_data['N_sim'])+"\n") 134 | 135 | # Solve OCP if we are in a planning cycle (MPC/planning frequency) 136 | if(i%int(sim_params['sim_freq']/sim_params['mpc_freq']) == 0): 137 | # Set x0 to measured state 138 | ddp.problem.x0 = sim_data['state_mea_SIM_RATE'][i, :] 139 | # Warm start using previous solution 140 | xs_init = list(ddp.xs[1:]) + [ddp.xs[-1]] 141 | xs_init[0] = sim_data['state_mea_SIM_RATE'][i, :] 142 | us_init = list(ddp.us[1:]) + [ddp.us[-1]] 143 | 144 | # Solve OCP & record MPC predictions 145 | solved = ddp.solve(xs_init, us_init, maxiter=ocp_params['maxiter'], is_feasible=False) 146 | print(solved) 147 | sim_data['state_pred'][mpc_cycle, :, :] = np.array(ddp.xs) 148 | sim_data['ctrl_pred'][mpc_cycle, :, :] = np.array(ddp.us) 149 | # Extract relevant predictions for interpolations 150 | x_curr = sim_data['state_pred'][mpc_cycle, 0, :] # x0* = measured state (q^, v^ ) 151 | x_pred = sim_data['state_pred'][mpc_cycle, 1, :] # x1* = predicted state (q1*, v1*) 152 | u_curr = sim_data['ctrl_pred'][mpc_cycle, 0, :] # u0* = optimal control (tau0*) 153 | u_pred = sim_data['ctrl_pred'][mpc_cycle, 1, :] # u1* = predicted optimal control (tau1*) 154 | # Record costs references 155 | q = sim_data['state_pred'][mpc_cycle, 0, :sim_data['nq']] 156 | sim_data['ctrl_ref'][mpc_cycle, :] = pin_utils.get_u_grav(q, ddp.problem.runningModels[0].differential.pinocchio, ocp_params['armature']) 157 | sim_data['state_ref'][mpc_cycle, :] = ddp.problem.runningModels[0].differential.costs.costs['stateReg'].cost.residual.reference 158 | sim_data['lin_pos_ee_ref'][mpc_cycle, :] = ddp.problem.runningModels[0].differential.costs.costs['translation'].cost.residual.reference 159 | 160 | 161 | # Select reference control and state for the current MPC cycle 162 | x_ref_MPC_RATE = x_curr + sim_data['ocp_to_mpc_ratio'] * (x_pred - x_curr) 163 | u_ref_MPC_RATE = u_curr 164 | if(mpc_cycle==0): 165 | sim_data['state_des_MPC_RATE'][mpc_cycle, :] = x_curr 166 | sim_data['ctrl_des_MPC_RATE'][mpc_cycle, :] = u_ref_MPC_RATE 167 | sim_data['state_des_MPC_RATE'][mpc_cycle+1, :] = x_ref_MPC_RATE 168 | 169 | xs = ddp.xs.tolist() 170 | us = ddp.us.tolist() 171 | K = ddp.K.tolist() 172 | for iii in range(T-1): 173 | if iii == 0: 174 | xs_int = np.linspace(xs[iii].copy(), xs[iii+1].copy(), int(dt/sim_data['dt_sim'])) 175 | us_int = np.linspace(us[iii].copy(), us[iii+1].copy(), int(dt/sim_data['dt_sim'])) 176 | Klist = [K[iii].copy()] * int(dt/sim_data['dt_sim']) 177 | else: 178 | xs_int = np.vstack((xs_int, np.linspace(xs[iii].copy(), xs[iii+1].copy(), int(dt/sim_data['dt_sim'])))) 179 | us_int = np.vstack((us_int, np.linspace(us[iii].copy(), us[iii+1].copy(), int(dt/sim_data['dt_sim'])))) 180 | Klist = Klist + [K[iii].copy()] * int(dt/sim_data['dt_sim']) 181 | 182 | # Increment planning counter 183 | mpc_cycle += 1 184 | ctrl_cycle = 0 185 | 186 | # Select reference control and state for the current SIMU cycle 187 | x_ref_SIM_RATE = x_curr + ctrl_cycle*sim_data['ocp_to_sim_ratio'] * (x_pred - x_curr) 188 | 189 | if(NAIVE): 190 | u_ref_SIM_RATE = u_curr #+ ctrl_cycle*sim_data['ocp_to_sim_ratio'] * (u_pred - u_curr) 191 | else: 192 | u_ref_SIM_RATE = us_int[ctrl_cycle] #u_curr #+ ctrl_cycle*sim_data['ocp_to_sim_ratio'] * (u_pred - u_curr) 193 | # if(INTERP_FEEDFORWARD): 194 | # u_ref_SIM_RATE += ctrl_cycle*sim_data['ocp_to_sim_ratio'] * (u_pred - u_curr) 195 | 196 | # First prediction = measurement = initialization of MPC 197 | if(i==0): 198 | sim_data['state_des_SIM_RATE'][i, :] = x_curr 199 | sim_data['ctrl_des_SIM_RATE'][i, :] = u_ref_SIM_RATE 200 | sim_data['state_des_SIM_RATE'][i+1, :] = x_ref_SIM_RATE 201 | 202 | # Interpolate state 203 | # if(RICCATI_TYPE == 1): 204 | # u_ref_SIM_RATE += ddp.K[0] @ (x_ref_SIM_RATE - sim_data['state_mea_SIM_RATE'][i, :]) 205 | # # Use x0 206 | # elif(RICCATI_TYPE == 2): 207 | 208 | if(NAIVE): 209 | u_ref_SIM_RATE += ddp.K[0] @ (ddp.problem.x0 - sim_data['state_mea_SIM_RATE'][i, :]) 210 | else: 211 | u_ref_SIM_RATE += Klist[ctrl_cycle] @ (xs_int[ctrl_cycle] - sim_data['state_mea_SIM_RATE'][i, :]) 212 | 213 | # # Interpolate gains + state 214 | # elif(RICCATI_TYPE == 3): 215 | # Kinterp = ddp.K[0] + ctrl_cycle*sim_data['ocp_to_sim_ratio'] * (ddp.K[1] - ddp.K[0]) 216 | # u_ref_SIM_RATE += Kinterp @ (x_ref_SIM_RATE - sim_data['state_mea_SIM_RATE'][i, :]) 217 | # # Interpolate the whole feedabck 218 | # elif(RICCATI_TYPE == 4): 219 | # ufb0 = ddp.K[0] @ (ddp.problem.x0 - sim_data['state_mea_SIM_RATE'][i, :]) 220 | # ufb1 = ddp.K[1] @ (x_pred - sim_data['state_mea_SIM_RATE'][i, :]) 221 | # ufb = ufb0 + ctrl_cycle*sim_data['ocp_to_sim_ratio'] * (ufb1 - ufb0) 222 | # u_ref_SIM_RATE += ufb #Kinterp @ (x_ref_SIM_RATE - sim_data['state_mea_SIM_RATE'][i, :]) 223 | # else: 224 | # u_ref_SIM_RATE -= Kv @ sim_data['state_mea_SIM_RATE'][i, nq:] 225 | 226 | # Simulate actuation 227 | # u_ref_SIM_RATE = 0.9 * u_ref_SIM_RATE 228 | # Send torque to simulator & step simulator 229 | robot_simulator.send_joint_command(u_ref_SIM_RATE) 230 | env.step() 231 | # Measure new state from simulator 232 | q_mea_SIM_RATE, v_mea_SIM_RATE = robot_simulator.get_state() 233 | # Update pinocchio model 234 | robot_simulator.forward_robot(q_mea_SIM_RATE, v_mea_SIM_RATE) 235 | # Record data 236 | x_mea_SIM_RATE = np.concatenate([q_mea_SIM_RATE, v_mea_SIM_RATE]).T 237 | sim_data['state_mea_SIM_RATE'][i+1, :] = x_mea_SIM_RATE 238 | ctrl_cycle += 1 239 | 240 | plot_data = mpc_utils.extract_plot_data_from_sim_data(sim_data) 241 | 242 | mpc_utils.plot_mpc_results(plot_data, which_plots=['all'], PLOT_PREDICTIONS=True, pred_plot_sampling=int(sim_params['mpc_freq'])) 243 | --------------------------------------------------------------------------------