├── .gitignore ├── LICENSE ├── Project.toml ├── README.md ├── animations ├── acrobot_joint_limits.gif ├── acrobot_no_joint_limits.gif ├── cartpole_friction_01.gif ├── cartpole_friction_1.gif ├── cartpole_friction_25.gif ├── cartpole_friction_35.gif ├── hopper_gait_1.gif ├── hopper_gait_2.gif ├── hopper_gait_3.gif ├── planar_push_rotate.gif ├── planar_push_translate.gif └── starship_bellyflop_landing.gif ├── deps ├── .gitignore └── build.jl ├── examples ├── .gitignore ├── Manifest.toml ├── Project.toml ├── README.md ├── acrobot.jl ├── cartpole.jl ├── comparisons │ ├── acrobot │ │ ├── Project.toml │ │ ├── acrobot.jl │ │ ├── acrobot.xml │ │ ├── acrobot_limits.xml │ │ └── mujoco_model.jl │ └── hopper.jl ├── generate_notebooks.jl ├── hopper.jl ├── planar_push.jl └── rocket.jl └── src ├── OptimizationDynamics.jl ├── dynamics.jl ├── gradient_bundle.jl ├── ls.jl └── models ├── acrobot ├── codegen.jl ├── model.jl ├── simulator_impact.jl ├── simulator_nominal.jl └── visuals.jl ├── cartpole ├── codegen.jl ├── model.jl ├── simulator_friction.jl ├── simulator_frictionless.jl └── visuals.jl ├── planar_push ├── codegen.jl ├── model.jl ├── simulator.jl └── visuals.jl ├── rocket ├── codegen.jl ├── dynamics.jl ├── model.jl ├── simulator.jl └── visuals.jl └── visualize.jl /.gitignore: -------------------------------------------------------------------------------- 1 | # Files generated by invoking Julia with --code-coverage 2 | *.jl.cov 3 | *.jl.*.cov 4 | 5 | # Files generated by invoking Julia with --track-allocation 6 | *.jl.mem 7 | 8 | # System-specific files and directories generated by the BinaryProvider and BinDeps packages 9 | # They contain absolute paths specific to the host computer, and so should not be committed 10 | deps/deps.jl 11 | deps/build.log 12 | deps/downloads/ 13 | deps/usr/ 14 | deps/src/ 15 | 16 | # Build artifacts for creating documentation generated by the Documenter package 17 | docs/build/ 18 | docs/site/ 19 | 20 | # File generated by Pkg, the package manager, based on a corresponding Project.toml 21 | # It records a fixed state of all packages used by the project. As such, it should not be 22 | # committed for packages, but should be committed for applications that require a static 23 | # environment. 24 | Manifest.toml 25 | 26 | # Outputs 27 | *.txt 28 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Taylor Howell 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Project.toml: -------------------------------------------------------------------------------- 1 | name = "OptimizationDynamics" 2 | uuid = "b76b7f0e-3267-49b4-a6cd-c01948ce091d" 3 | authors = ["thowell "] 4 | version = "0.1.0" 5 | 6 | [deps] 7 | BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" 8 | Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" 9 | CoordinateTransformations = "150eb455-5306-5404-9cee-2592286d6298" 10 | DirectTrajectoryOptimization = "a9f42406-efe7-414c-8b71-df971cc98041" 11 | GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326" 12 | IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173" 13 | IterativeLQR = "605048dd-e178-462b-beb9-98a09398ef27" 14 | JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" 15 | LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" 16 | MeshCat = "283c5d60-a78f-5afe-a0af-af636b173e11" 17 | RoboDojo = "2859eb88-988f-4613-a49e-c643bedbc906" 18 | Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc" 19 | Scratch = "6c6a2e73-6563-6170-7368-637461726353" 20 | Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" 21 | 22 | [compat] 23 | Colors = "0.12" 24 | CoordinateTransformations = "0.6.2" 25 | DirectTrajectoryOptimization = "0.1.2" 26 | GeometryBasics = "0.3.13" 27 | IfElse = "0.1.1" 28 | IterativeLQR = "0.1.0" 29 | JLD2 = "0.4.15" 30 | MeshCat = "0.13 - 0.14" 31 | RoboDojo = "0.1.2" 32 | Rotations = "1.0.2 - 1.0.2" 33 | Scratch = "1.0" 34 | Symbolics = "0.1.29 - 0.1.29" 35 | julia = "1.6" 36 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Trajectory Optimization with Optimization-Based Dynamics 2 | 3 | This repository contains the implementation and examples from our paper: [Trajectory Optimization with Optimization-Based Dynamics](https://arxiv.org/abs/2109.04928). 4 | 5 | ## Installation 6 | From the Julia REPL, type `]` to enter the Pkg REPL mode and run: 7 | ```julia 8 | pkg> add https://github.com/thowell/optimization_dynamics 9 | ``` 10 | This will install the package. 11 | 12 | Run 13 | ```julia 14 | (OptimizationDynamics) pkg> build 15 | ``` 16 | to build the package. 17 | 18 | [Notebooks](examples) are generated upon installation and can be run for the following examples: 19 | 20 | ## planar push 21 | drawing 22 | 23 | ## acrobot with joint limits 24 | drawing 25 | 26 | ## cart-pole with joint friction 27 | drawing 28 | 29 | ## hopper gait 30 | drawing 31 | 32 | ## rocket with thrust limits 33 | drawing 34 | 35 | Additional comparisons with [MuJoCo](examples/comparisons/acrobot) and [contact-implicit trajectory optimization](examples/comparisons/hopper.jl) are available. 36 | The path-following solver and hopper model can be found in [RoboDojo.jl](https://github.com/thowell/RoboDojo.jl). 37 | 38 | -------------------------------------------------------------------------------- /animations/acrobot_joint_limits.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thowell/optimization_dynamics/fdfa55324037a3571e1efdf1f349e7e770613005/animations/acrobot_joint_limits.gif -------------------------------------------------------------------------------- /animations/acrobot_no_joint_limits.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thowell/optimization_dynamics/fdfa55324037a3571e1efdf1f349e7e770613005/animations/acrobot_no_joint_limits.gif -------------------------------------------------------------------------------- /animations/cartpole_friction_01.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thowell/optimization_dynamics/fdfa55324037a3571e1efdf1f349e7e770613005/animations/cartpole_friction_01.gif -------------------------------------------------------------------------------- /animations/cartpole_friction_1.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thowell/optimization_dynamics/fdfa55324037a3571e1efdf1f349e7e770613005/animations/cartpole_friction_1.gif -------------------------------------------------------------------------------- /animations/cartpole_friction_25.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thowell/optimization_dynamics/fdfa55324037a3571e1efdf1f349e7e770613005/animations/cartpole_friction_25.gif -------------------------------------------------------------------------------- /animations/cartpole_friction_35.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thowell/optimization_dynamics/fdfa55324037a3571e1efdf1f349e7e770613005/animations/cartpole_friction_35.gif -------------------------------------------------------------------------------- /animations/hopper_gait_1.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thowell/optimization_dynamics/fdfa55324037a3571e1efdf1f349e7e770613005/animations/hopper_gait_1.gif -------------------------------------------------------------------------------- /animations/hopper_gait_2.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thowell/optimization_dynamics/fdfa55324037a3571e1efdf1f349e7e770613005/animations/hopper_gait_2.gif -------------------------------------------------------------------------------- /animations/hopper_gait_3.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thowell/optimization_dynamics/fdfa55324037a3571e1efdf1f349e7e770613005/animations/hopper_gait_3.gif -------------------------------------------------------------------------------- /animations/planar_push_rotate.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thowell/optimization_dynamics/fdfa55324037a3571e1efdf1f349e7e770613005/animations/planar_push_rotate.gif -------------------------------------------------------------------------------- /animations/planar_push_translate.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thowell/optimization_dynamics/fdfa55324037a3571e1efdf1f349e7e770613005/animations/planar_push_translate.gif -------------------------------------------------------------------------------- /animations/starship_bellyflop_landing.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/thowell/optimization_dynamics/fdfa55324037a3571e1efdf1f349e7e770613005/animations/starship_bellyflop_landing.gif -------------------------------------------------------------------------------- /deps/.gitignore: -------------------------------------------------------------------------------- 1 | build.log 2 | -------------------------------------------------------------------------------- /deps/build.jl: -------------------------------------------------------------------------------- 1 | append!(empty!(LOAD_PATH), Base.DEFAULT_LOAD_PATH) 2 | using Pkg 3 | 4 | ################################################################################ 5 | # Generate notebooks 6 | ################################################################################ 7 | exampledir = joinpath(@__DIR__, "..", "examples") 8 | Pkg.activate(exampledir) 9 | Pkg.instantiate() 10 | include(joinpath(exampledir, "generate_notebooks.jl")) 11 | 12 | ################################################################################ 13 | # Build simulation environments 14 | ################################################################################ 15 | pkgdir = joinpath(@__DIR__, "..") 16 | Pkg.activate(pkgdir) 17 | 18 | using JLD2 19 | using Symbolics 20 | using LinearAlgebra 21 | using Scratch 22 | using RoboDojo 23 | using Rotations 24 | import RoboDojo: Model, lagrangian_derivatives, IndicesZ, cone_product 25 | 26 | # acrobot 27 | include("../src/models/acrobot/model.jl") 28 | include("../src/models/acrobot/simulator_impact.jl") 29 | include("../src/models/acrobot/simulator_nominal.jl") 30 | include("../src/models/acrobot/codegen.jl") 31 | 32 | # cartpole 33 | include("../src/models/cartpole/model.jl") 34 | include("../src/models/cartpole/simulator_friction.jl") 35 | include("../src/models/cartpole/simulator_frictionless.jl") 36 | include("../src/models/cartpole/codegen.jl") 37 | 38 | # hopper from RoboDojo.jl 39 | 40 | # planar push 41 | include("../src/models/planar_push/model.jl") 42 | include("../src/models/planar_push/simulator.jl") 43 | include("../src/models/planar_push/codegen.jl") 44 | 45 | # rocket 46 | include("../src/models/rocket/model.jl") 47 | include("../src/models/rocket/simulator.jl") 48 | include("../src/models/rocket/codegen.jl") 49 | 50 | -------------------------------------------------------------------------------- /examples/.gitignore: -------------------------------------------------------------------------------- 1 | **/*.ipynb 2 | **/*.md 3 | -------------------------------------------------------------------------------- /examples/Manifest.toml: -------------------------------------------------------------------------------- 1 | # This file is machine-generated - editing it directly is not advised 2 | 3 | [[Base64]] 4 | uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" 5 | 6 | [[Dates]] 7 | deps = ["Printf"] 8 | uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" 9 | 10 | [[IOCapture]] 11 | deps = ["Logging", "Random"] 12 | git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a" 13 | uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89" 14 | version = "0.2.2" 15 | 16 | [[InteractiveUtils]] 17 | deps = ["Markdown"] 18 | uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240" 19 | 20 | [[JSON]] 21 | deps = ["Dates", "Mmap", "Parsers", "Unicode"] 22 | git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37" 23 | uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" 24 | version = "0.21.2" 25 | 26 | [[Literate]] 27 | deps = ["Base64", "IOCapture", "JSON", "REPL"] 28 | git-tree-sha1 = "d3493acfb9e6aa0cff46b09773fc2342327b0feb" 29 | uuid = "98b081ad-f1c9-55d3-8b20-4c87d4299306" 30 | version = "2.9.4" 31 | 32 | [[Logging]] 33 | uuid = "56ddb016-857b-54e1-b83d-db4d58db5568" 34 | 35 | [[Markdown]] 36 | deps = ["Base64"] 37 | uuid = "d6f4376e-aef5-505a-96c1-9c027394607a" 38 | 39 | [[Mmap]] 40 | uuid = "a63ad114-7e13-5084-954f-fe012c677804" 41 | 42 | [[Parsers]] 43 | deps = ["Dates"] 44 | git-tree-sha1 = "ae4bbcadb2906ccc085cf52ac286dc1377dceccc" 45 | uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0" 46 | version = "2.1.2" 47 | 48 | [[Printf]] 49 | deps = ["Unicode"] 50 | uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7" 51 | 52 | [[REPL]] 53 | deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"] 54 | uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb" 55 | 56 | [[Random]] 57 | deps = ["Serialization"] 58 | uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" 59 | 60 | [[Serialization]] 61 | uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b" 62 | 63 | [[Sockets]] 64 | uuid = "6462fe0b-24de-5631-8697-dd941f90decc" 65 | 66 | [[Unicode]] 67 | uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5" 68 | -------------------------------------------------------------------------------- /examples/Project.toml: -------------------------------------------------------------------------------- 1 | [deps] 2 | Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306" 3 | -------------------------------------------------------------------------------- /examples/README.md: -------------------------------------------------------------------------------- 1 | # Examples 2 | 3 | The `.jl` files in each subdirectory are meant to be processed using [Literate.jl](https://github.com/fredrikekre/Literate.jl). 4 | During the build process, the `.jl` files are converted to notebooks. 5 | 6 | 1. [install optimization_dynamics](https://github.com/thowell/optimization_dynamics) 7 | 2. [install IJulia](https://github.com/JuliaLang/IJulia.jl) (`add` it to the default project) 8 | 3. interact with notebooks 9 | ``` 10 | using IJulia, OptimizationDynamics 11 | notebook(dir=joinpath(dirname(pathof(OptimizationDynamics)), "..", "examples")) 12 | ``` -------------------------------------------------------------------------------- /examples/acrobot.jl: -------------------------------------------------------------------------------- 1 | using OptimizationDynamics 2 | const iLQR = OptimizationDynamics.IterativeLQR 3 | using LinearAlgebra 4 | using Random 5 | 6 | # ## visualization 7 | vis = Visualizer() 8 | render(vis) 9 | 10 | # ## mode 11 | MODE = :impact 12 | MODE = :nominal 13 | 14 | # ## state-space model 15 | h = 0.05 16 | T = 101 17 | κ_grad = 1.0e-3 # gradient smoothness 18 | 19 | if MODE == :impact 20 | im_dyn = ImplicitDynamics(acrobot_impact, h, eval(r_acrobot_impact_func), eval(rz_acrobot_impact_func), eval(rθ_acrobot_impact_func); 21 | r_tol=1.0e-8, κ_eval_tol=1.0e-4, 22 | κ_grad_tol=κ_grad, 23 | no_friction=true) 24 | else 25 | im_dyn = ImplicitDynamics(acrobot_nominal, h, eval(r_acrobot_nominal_func), eval(rz_acrobot_nominal_func), eval(rθ_acrobot_nominal_func); 26 | r_tol=1.0e-8, κ_eval_tol=1.0, κ_grad_tol=1.0, no_friction=true) 27 | end 28 | 29 | nx = 2 * acrobot_impact.nq 30 | nu = acrobot_impact.nu 31 | 32 | # ## iLQR model 33 | ilqr_dyn = iLQR.Dynamics((d, x, u, w) -> f(d, im_dyn, x, u, w), 34 | (dx, x, u, w) -> fx(dx, im_dyn, x, u, w), 35 | (du, x, u, w) -> fu(du, im_dyn, x, u, w), 36 | nx, nx, nu) 37 | model = [ilqr_dyn for t = 1:T-1]; 38 | 39 | # ## initial and goal states 40 | q1 = [0.0; 0.0] 41 | q2 = [0.0; 0.0] 42 | qT = [π; 0.0] 43 | q_ref = qT 44 | 45 | x1 = [q1; q2] 46 | xT = [qT; qT] 47 | 48 | # ## objective 49 | function objt(x, u, w) 50 | J = 0.0 51 | 52 | q1 = x[1:acrobot_impact.nq] 53 | q2 = x[acrobot_impact.nq .+ (1:acrobot_impact.nq)] 54 | v1 = (q2 - q1) ./ h 55 | 56 | J += 0.5 * 0.1 * transpose(v1) * v1 57 | J += 0.5 * transpose(u) * u 58 | 59 | return J 60 | end 61 | 62 | function objT(x, u, w) 63 | J = 0.0 64 | 65 | q1 = x[1:acrobot_impact.nq] 66 | q2 = x[acrobot_impact.nq .+ (1:acrobot_impact.nq)] 67 | v1 = (q2 - q1) ./ h 68 | 69 | J += 0.5 * 0.1 * transpose(v1) * v1 70 | 71 | return J 72 | end 73 | 74 | ct = iLQR.Cost(objt, nx, nu) 75 | cT = iLQR.Cost(objT, nx, 0) 76 | obj = [[ct for t = 1:T-1]..., cT]; 77 | 78 | # ## constraints 79 | function terminal_con(x, u, w) 80 | [ 81 | x - xT; # goal 82 | ] 83 | end 84 | 85 | cont = iLQR.Constraint() 86 | conT = iLQR.Constraint(terminal_con, nx, 0) 87 | cons = [[cont for t = 1:T-1]..., conT]; 88 | 89 | # ## rollout 90 | Random.seed!(1) 91 | ū = [1.0e-3 * randn(nu) for t = 1:T-1] 92 | x̄ = iLQR.rollout(model, x1, ū) 93 | q̄ = state_to_configuration(x̄) 94 | visualize!(vis, acrobot_impact, q̄, Δt=h); 95 | 96 | # ## solver 97 | solver = iLQR.solver(model, obj, cons, 98 | opts=iLQR.Options(linesearch = :armijo, 99 | α_min=1.0e-5, 100 | obj_tol=1.0e-5, 101 | grad_tol=1.0e-5, 102 | max_iter=50, 103 | max_al_iter=20, 104 | con_tol=0.001, 105 | ρ_init=1.0, 106 | ρ_scale=10.0, 107 | verbose=false)) 108 | iLQR.initialize_controls!(solver, ū) 109 | iLQR.initialize_states!(solver, x̄); 110 | 111 | # ## solve 112 | iLQR.reset!(solver.s_data) 113 | @time iLQR.solve!(solver); 114 | 115 | @show solver.s_data.iter[1] 116 | @show iLQR.eval_obj(solver.m_data.obj.costs, solver.m_data.x, solver.m_data.u, solver.m_data.w) 117 | @show norm(terminal_con(solver.m_data.x[T], zeros(0), zeros(0)), Inf) 118 | @show solver.s_data.obj[1] # augmented Lagrangian cost 119 | 120 | # ## solution 121 | x_sol, u_sol = iLQR.get_trajectory(solver) 122 | q_sol = state_to_configuration(x_sol) 123 | visualize!(vis, acrobot_impact, q_sol, Δt=h); 124 | 125 | # ## benchmark 126 | using BenchmarkTools 127 | solver.options.verbose = false 128 | @benchmark iLQR.solve!($solver, x̄, ū) setup=(x̄=deepcopy(x̄), ū=deepcopy(ū)); 129 | 130 | -------------------------------------------------------------------------------- /examples/cartpole.jl: -------------------------------------------------------------------------------- 1 | using OptimizationDynamics 2 | const iLQR = OptimizationDynamics.IterativeLQR 3 | using LinearAlgebra 4 | using Random 5 | 6 | # ## visualization 7 | vis = Visualizer() 8 | render(vis) 9 | 10 | # ## mode 11 | MODE = :friction 12 | MODE = :frictionless 13 | 14 | # ## state-space model 15 | h = 0.05 16 | T = 51 17 | 18 | if MODE == :friction 19 | im_dyn = ImplicitDynamics(cartpole_friction, h, eval(r_cartpole_friction_func), eval(rz_cartpole_friction_func), eval(rθ_cartpole_friction_func); 20 | r_tol=1.0e-8, κ_eval_tol=1.0e-4, κ_grad_tol=1.0e-3, no_impact=true) 21 | cartpole_friction.friction .= [0.35; 0.35] 22 | ## cartpole_friction.friction .= [0.25; 0.25] 23 | ## cartpole_friction.friction .= [0.1; 0.1] 24 | ## cartpole_friction.friction .= [0.01; 0.01] 25 | else 26 | im_dyn = ImplicitDynamics(cartpole_frictionless, h, eval(r_cartpole_frictionless_func), eval(rz_cartpole_frictionless_func), eval(rθ_cartpole_frictionless_func); 27 | r_tol=1.0e-8, κ_eval_tol=1.0, κ_grad_tol=1.0, no_impact=true, no_friction=true) 28 | end 29 | 30 | nx = 2 * cartpole_friction.nq 31 | nu = cartpole_friction.nu 32 | 33 | # ## iLQR model 34 | ilqr_dyn = iLQR.Dynamics((d, x, u, w) -> f(d, im_dyn, x, u, w), 35 | (dx, x, u, w) -> fx(dx, im_dyn, x, u, w), 36 | (du, x, u, w) -> fu(du, im_dyn, x, u, w), 37 | nx, nx, nu) 38 | model = [ilqr_dyn for t = 1:T-1]; 39 | 40 | # ## initial conditions 41 | q0 = [0.0; 0.0] 42 | q1 = [0.0; 0.0] 43 | qT = [0.0; π] 44 | q_ref = [0.0; π] 45 | 46 | x1 = [q1; q1] 47 | xT = [qT; qT] 48 | 49 | # ## objective 50 | function objt(x, u, w) 51 | J = 0.0 52 | J += transpose(u) * u 53 | return J 54 | end 55 | 56 | function objT(x, u, w) 57 | J = 0.0 58 | J += transpose(x - xT) * (x - xT) 59 | return J 60 | end 61 | 62 | ct = iLQR.Cost(objt, nx, nu) 63 | cT = iLQR.Cost(objT, nx, 0) 64 | obj = [[ct for t = 1:T-1]..., cT]; 65 | 66 | # ## constraints 67 | function terminal_con(x, u, w) 68 | [ 69 | x - xT; # goal 70 | ] 71 | end 72 | 73 | cont = iLQR.Constraint() 74 | conT = iLQR.Constraint(terminal_con, nx, 0) 75 | cons = [[cont for t = 1:T-1]..., conT]; 76 | 77 | # ## rollout 78 | ū = [(t == 1 ? -1.5 : 0.0) * ones(nu) for t = 1:T-1] # set value to -1.0 when friction coefficient = 0.25 79 | x̄ = iLQR.rollout(model, x1, ū) 80 | q̄ = state_to_configuration(x̄) 81 | visualize!(vis, cartpole_friction, q̄, Δt=h); 82 | 83 | # ## solver 84 | solver = iLQR.solver(model, obj, cons, 85 | opts=iLQR.Options(linesearch = :armijo, 86 | α_min=1.0e-5, 87 | obj_tol=1.0e-5, 88 | grad_tol=1.0e-3, 89 | max_iter=100, 90 | max_al_iter=20, 91 | con_tol=0.005, 92 | ρ_init=1.0, 93 | ρ_scale=10.0, 94 | verbose=false)) 95 | iLQR.initialize_controls!(solver, ū) 96 | iLQR.initialize_states!(solver, x̄); 97 | 98 | # ## solve 99 | iLQR.reset!(solver.s_data) 100 | @time iLQR.solve!(solver); 101 | 102 | @show iLQR.eval_obj(solver.m_data.obj.costs, solver.m_data.x, solver.m_data.u, solver.m_data.w) 103 | @show solver.s_data.iter[1] 104 | @show norm(terminal_con(solver.m_data.x[T], zeros(0), zeros(0)), Inf) 105 | @show solver.s_data.obj[1] # augmented Lagrangian cost 106 | 107 | # ## solution 108 | x_sol, u_sol = iLQR.get_trajectory(solver) 109 | q_sol = state_to_configuration(x_sol) 110 | visualize!(vis, cartpole_friction, q_sol, Δt=h); 111 | 112 | # ## benchmark 113 | solver.options.verbose = false 114 | @benchmark iLQR.solve!($solver, x̄, ū) setup=(x̄=deepcopy(x̄), ū=deepcopy(ū)); -------------------------------------------------------------------------------- /examples/comparisons/acrobot/Project.toml: -------------------------------------------------------------------------------- 1 | [deps] 2 | BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" 3 | Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" 4 | FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" 5 | IterativeLQR = "605048dd-e178-462b-beb9-98a09398ef27" 6 | LyceumMuJoCo = "48b9757e-04b8-4dbf-b6ed-75c13d9e4026" 7 | LyceumMuJoCoViz = "409fd311-7925-44f6-a514-6be572b34fdd" 8 | MuJoCo = "93189219-7048-461c-94ec-443a161ed927" 9 | 10 | [compat] 11 | FiniteDiff = "2.10.1" 12 | Colors = "0.12.8" 13 | IterativeLQR = "0.1.1" 14 | -------------------------------------------------------------------------------- /examples/comparisons/acrobot/acrobot.jl: -------------------------------------------------------------------------------- 1 | using Pkg 2 | Pkg.activate(@__DIR__) 3 | 4 | using MuJoCo 5 | mj_activate("/home/taylor/.mujoco/bin/mjkey.txt") # set location to MuJoCo key path 6 | 7 | using LyceumMuJoCo, LyceumMuJoCoViz 8 | using FiniteDiff 9 | 10 | using IterativeLQR 11 | using LinearAlgebra 12 | using Random 13 | 14 | # ## load MuJoCo model 15 | path = joinpath(@__DIR__, "acrobot.xml") 16 | path = joinpath(@__DIR__, "acrobot_limits.xml") 17 | 18 | include("mujoco_model.jl") 19 | acrobot_mujoco = MuJoCoModel(path) 20 | sim = MJSim(acrobot_mujoco.m, acrobot_mujoco.d); 21 | 22 | # ## horizon 23 | T = 101 24 | 25 | # ## acrobot_mujoco 26 | nx = acrobot_mujoco.nx 27 | nu = acrobot_mujoco.nu 28 | 29 | # ## model 30 | dyn = IterativeLQR.Dynamics( 31 | (y, x, u, w) -> f_mujoco!(y, acrobot_mujoco, x, u), 32 | (dx, x, u, w) -> fx_mujoco!(dx, acrobot_mujoco, x, u), 33 | (du, x, u, w) -> fu_mujoco!(du, acrobot_mujoco, x, u), 34 | nx, nx, nu) 35 | 36 | model = [dyn for t = 1:T-1]; 37 | 38 | # ## initial conditions 39 | x1 = [π; 0.0; 0.0; 0.0] 40 | xT = [0.0; 0.0; 0.0; 0.0] 41 | 42 | # ## objective 43 | function objt(x, u, w) 44 | J = 0.0 45 | v = x[3:4] 46 | J += 0.5 * 0.1 * transpose(v) * v 47 | J += 0.5 * transpose(u) * u 48 | return J 49 | end 50 | 51 | function objT(x, u, w) 52 | J = 0.0 53 | v = x[3:4] 54 | J += 0.5 * 0.1 * transpose(v) * v 55 | return J 56 | end 57 | 58 | ct = IterativeLQR.Cost(objt, acrobot_mujoco.nx, acrobot_mujoco.nu) 59 | cT = IterativeLQR.Cost(objT, acrobot_mujoco.nx, 0) 60 | obj = [[ct for t = 1:T-1]..., cT]; 61 | 62 | # ## constraints 63 | goal(x, u, w) = x - xT 64 | 65 | cont = IterativeLQR.Constraint() 66 | conT = IterativeLQR.Constraint(goal, nx, 0) 67 | cons = [[cont for t = 1:T-1]..., conT]; 68 | 69 | # # rollout 70 | Random.seed!(1) 71 | ū = [1.0e-3 * randn(acrobot_mujoco.nu) for t = 1:T-1] 72 | w = [zeros(0) for t = 1:T-1] 73 | x̄ = IterativeLQR.rollout(model, x1, ū); 74 | 75 | # ## solver 76 | solver = IterativeLQR.solver(model, obj, cons, 77 | opts=iLQR.Options( 78 | linesearch = :armijo, 79 | α_min=1.0e-5, 80 | obj_tol=1.0e-5, 81 | grad_tol=1.0e-5, 82 | max_iter=50, 83 | max_al_iter=10, 84 | con_tol=0.001, 85 | ρ_init=1.0, 86 | ρ_scale=10.0, 87 | verbose=false)) 88 | IterativeLQR.initialize_controls!(solver, ū) 89 | IterativeLQR.initialize_states!(solver, x̄); 90 | 91 | # ## solve 92 | IterativeLQR.reset!(solver.s_data) 93 | @time IterativeLQR.solve!(solver); 94 | 95 | @show solver.s_data.iter[1] 96 | @show IterativeLQR.eval_obj(solver.m_data.obj.costs, solver.m_data.x, solver.m_data.u, solver.m_data.w) 97 | @show norm(goal(solver.m_data.x[T], zeros(0), zeros(0)), Inf) 98 | @show solver.s_data.obj[1] # augmented Lagrangian cost 99 | 100 | # ## benchmark 101 | solver.options.verbose = true 102 | @benchmark IterativeLQR.solve!($solver, x̄, ū) setup=(x̄=deepcopy(x̄), ū=deepcopy(ū)); 103 | 104 | # ## solution 105 | x_sol, u_sol = IterativeLQR.get_trajectory(solver); 106 | 107 | # ## MuJoCo visualizer 108 | states = Array(undef, statespace(sim), T-1) 109 | for t = 1:T-1 110 | sim.d.qpos .= x_sol[t][acrobot_mujoco.idx_pos] 111 | sim.d.qvel .= x_sol[t][acrobot_mujoco.idx_vel] 112 | sim.d.ctrl .= u_sol[t] 113 | states[:, t] .= getstate(sim) 114 | end 115 | 116 | visualize(sim, trajectories=[states]); # ctrl + LEFT (access trajectory mode) 117 | 118 | -------------------------------------------------------------------------------- /examples/comparisons/acrobot/acrobot.xml: -------------------------------------------------------------------------------- 1 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | -------------------------------------------------------------------------------- /examples/comparisons/acrobot/acrobot_limits.xml: -------------------------------------------------------------------------------- 1 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | -------------------------------------------------------------------------------- /examples/comparisons/acrobot/mujoco_model.jl: -------------------------------------------------------------------------------- 1 | using FiniteDiff 2 | using BenchmarkTools 3 | 4 | struct MuJoCoModel{T,CX,CU} 5 | m::jlModel 6 | d::jlData 7 | x::Vector{T} 8 | fx::Matrix{T} 9 | fu::Matrix{T} 10 | nx::Int 11 | nu::Int 12 | idx_u::Vector{Int} 13 | idx_pos::Vector{Int} 14 | idx_vel::Vector{Int} 15 | idx_next::Vector{Int} 16 | cache_x::CX 17 | cache_u::CU 18 | end 19 | 20 | function MuJoCoModel(path) 21 | m = jlModel(path) 22 | d = jlData(m) 23 | 24 | nq = length(d.qpos) 25 | nv = length(d.qvel) 26 | nx = nq + nv 27 | nu = length(d.ctrl) 28 | 29 | x = zeros(nx) 30 | u = zeros(nu) 31 | fx = zeros(nx, nx) 32 | fu = zeros(nx, nu) 33 | 34 | idx_u = collect(1:nu) 35 | idx_pos = collect(1:nq) 36 | idx_vel = collect(nq .+ (1:nv)) 37 | idx_next = collect(1:nq) 38 | 39 | cache_x = FiniteDiff.JacobianCache(x, x) 40 | cache_u = FiniteDiff.JacobianCache(u, x) 41 | 42 | MuJoCoModel(m, d, x, 43 | fx, fu, 44 | nx, nu, 45 | idx_u, idx_pos, idx_vel, idx_next, 46 | cache_x, cache_u) 47 | end 48 | 49 | # dynamics methods 50 | function f_mujoco!(y, model::MuJoCoModel, x, u) 51 | model.d.ctrl .= @views u[model.idx_u] 52 | model.d.qpos[model.idx_next] .= @views x[model.idx_pos] 53 | model.d.qvel[model.idx_next] .= @views x[model.idx_vel] 54 | mj_step(model.m, model.d) 55 | y[model.idx_pos] .= @views model.d.qpos[model.idx_next] 56 | y[model.idx_vel] .= @views model.d.qvel[model.idx_next] 57 | return y 58 | end 59 | 60 | f_mujoco(model::MuJoCoModel, x, u) = f_mujoco!(model.x, model, x, u) 61 | 62 | function fx_mujoco!(dx, model::MuJoCoModel, x, u) 63 | FiniteDiff.finite_difference_jacobian!(dx, (a, b) -> f_mujoco!(a, model, b, u), x, model.cache_x) 64 | return dx 65 | end 66 | 67 | fx_mujoco(model::MuJoCoModel, x, u) = fx_mujoco!(model.fx, model, x, u) 68 | 69 | function fu_mujoco!(du, model::MuJoCoModel, x, u) 70 | FiniteDiff.finite_difference_jacobian!(du, (a, b) -> f_mujoco!(a, model, x, b), u, model.cache_u) 71 | return du 72 | end 73 | 74 | fu_mujoco(model::MuJoCoModel, x, u) = fu_mujoco!(model.fu, model, x, u) 75 | 76 | # x2 = zeros(acrobot.nx) 77 | # x1 = zeros(acrobot.nx) 78 | # u1 = zeros(acrobot.nu) 79 | # f(acrobot, x1, u1) 80 | # @benchmark f($acrobot, $x1, $u1) 81 | # f!(x2, acrobot, x1, u1) 82 | # @benchmark f!($x2, $acrobot, $x1, $u1) 83 | # fx(acrobot, x1, u1) 84 | # @benchmark fx($acrobot, $x1, $u1) 85 | # fu(acrobot, x1, u1) 86 | # @benchmark fu($acrobot, $x1, $u1) 87 | -------------------------------------------------------------------------------- /examples/comparisons/hopper.jl: -------------------------------------------------------------------------------- 1 | using RoboDojo 2 | using LinearAlgebra 3 | using DirectTrajectoryOptimization 4 | const DTO = DirectTrajectoryOptimization 5 | 6 | function hopper_dyn(mass_matrix, dynamics_bias, h, y, x, u, w) 7 | model = RoboDojo.hopper 8 | 9 | # dimensions 10 | nq = model.nq 11 | nu = model.nu 12 | 13 | # configurations 14 | 15 | q1⁻ = x[1:nq] 16 | q2⁻ = x[nq .+ (1:nq)] 17 | q2⁺ = y[1:nq] 18 | q3⁺ = y[nq .+ (1:nq)] 19 | 20 | # control 21 | u_control = u[1:nu] 22 | γ = u[nu .+ (1:4)] 23 | β = u[nu + 4 .+ (1:4)] 24 | 25 | E = [1.0 -1.0] # friction mapping 26 | J = RoboDojo.contact_jacobian(model, q2⁺) 27 | λ = transpose(J) * [[E * β[1:2]; γ[1]]; 28 | [E * β[3:4]; γ[2]]; 29 | γ[3:4]] 30 | λ[3] += (model.body_radius * E * β[1:2])[1] # friction on body creates a moment 31 | 32 | [ 33 | q2⁺ - q2⁻; 34 | RoboDojo.dynamics(model, mass_matrix, dynamics_bias, 35 | h, q1⁻, q2⁺, u_control, zeros(model.nw), λ, q3⁺) 36 | ] 37 | end 38 | 39 | function hopper_dyn1(mass_matrix, dynamics_bias, h, y, x, u, w) 40 | nx = 8 41 | [ 42 | hopper_dyn(mass_matrix, dynamics_bias, h, y, x, u, w); 43 | y[nx .+ (1:5)] - [u[2 .+ (1:4)]; u[end]]; 44 | y[nx + 5 .+ (1:nx)] - x 45 | ] 46 | end 47 | 48 | function hopper_dynt(mass_matrix, dynamics_bias, h, y, x, u, w) 49 | nx = 8 50 | [ 51 | hopper_dyn(mass_matrix, dynamics_bias, h, y, x, u, w); 52 | y[nx .+ (1:5)] - [u[2 .+ (1:4)]; u[end]]; 53 | y[nx + 5 .+ (1:nx)] - x[nx + 5 .+ (1:nx)] 54 | ] 55 | end 56 | 57 | function contact_constraints_inequality_1(h, x, u, w) 58 | model = RoboDojo.hopper 59 | 60 | nq = model.nq 61 | nu = model.nu 62 | nx = 2nq 63 | 64 | q2 = x[1:nq] 65 | q3 = x[nq .+ (1:nq)] 66 | 67 | u_control = u[1:nu] 68 | γ = u[nu .+ (1:4)] 69 | β = u[nu + 4 .+ (1:4)] 70 | ψ = u[nu + 4 + 4 .+ (1:2)] 71 | η = u[nu + 4 + 4 + 2 .+ (1:4)] 72 | sα = u[nu + 4 + 4 + 2 + 4 .+ (1:1)] 73 | 74 | ϕ = RoboDojo.signed_distance(model, q3) 75 | 76 | μ = [model.friction_body_world; model.friction_foot_world] 77 | fc = μ .* γ[1:2] - [sum(β[1:2]); sum(β[3:4])] 78 | 79 | [ 80 | -ϕ; 81 | -fc; 82 | β .* η .- sα; 83 | ψ .* fc .- sα; 84 | ] 85 | end 86 | 87 | function contact_constraints_inequality_t(h, x, u, w) 88 | model = RoboDojo.hopper 89 | 90 | nq = model.nq 91 | nu = model.nu 92 | nx = 2nq 93 | 94 | q2 = x[1:nq] 95 | q3 = x[nq .+ (1:nq)] 96 | 97 | γ = u[nu .+ (1:4)] 98 | β = u[nu + 4 .+ (1:4)] 99 | ψ = u[nu + 4 + 4 .+ (1:2)] 100 | η = u[nu + 4 + 4 + 2 .+ (1:4)] 101 | sα = u[nu + 4 + 4 + 2 + 4 .+ (1:1)] 102 | 103 | ϕ = RoboDojo.signed_distance(model, q3) 104 | γ⁻ = x[nx .+ (1:4)] 105 | sα⁻ = x[nx + 4 .+ (1:1)] 106 | 107 | μ = [model.friction_body_world; model.friction_foot_world] 108 | fc = μ .* γ[1:2] - [sum(β[1:2]); sum(β[3:4])] 109 | 110 | [ 111 | -ϕ; 112 | -fc; 113 | γ⁻ .* ϕ .- sα⁻; 114 | β .* η .- sα; 115 | ψ .* fc .- sα; 116 | ] 117 | end 118 | 119 | function contact_constraints_inequality_T(h, x, u, w) 120 | model = RoboDojo.hopper 121 | 122 | nq = model.nq 123 | nx = 2nq 124 | 125 | q2 = x[1:nq] 126 | q3 = x[nq .+ (1:nq)] 127 | 128 | ϕ = RoboDojo.signed_distance(model, q3) 129 | γ⁻ = x[nx .+ (1:4)] 130 | sα⁻ = x[nx + 4 .+ (1:1)] 131 | 132 | [ 133 | -ϕ; 134 | γ⁻ .* ϕ .- sα⁻; 135 | ] 136 | end 137 | 138 | function contact_constraints_equality(h, x, u, w) 139 | model = RoboDojo.hopper 140 | 141 | nq = model.nq 142 | nu = model.nu 143 | 144 | q2 = x[1:nq] 145 | q3 = x[nq .+ (1:nq)] 146 | 147 | γ = u[nu .+ (1:4)] 148 | β = u[nu + 4 .+ (1:4)] 149 | ψ = u[nu + 4 + 4 .+ (1:2)] 150 | η = u[nu + 4 + 4 + 2 .+ (1:4)] 151 | 152 | v = (q3 - q2) ./ h[1] 153 | vT_body = v[1] + model.body_radius * v[3] 154 | vT_foot = (RoboDojo.kinematics_foot_jacobian(model, q3) * v)[1] 155 | vT = [vT_body; -vT_body; vT_foot; -vT_foot] 156 | 157 | ψ_stack = [ψ[1] * ones(2); ψ[2] * ones(2)] 158 | 159 | [ 160 | η - vT - ψ_stack; 161 | ] 162 | end 163 | 164 | # ## horizon 165 | T = 21 166 | h = 0.05 167 | 168 | # ## hopper 169 | nx = 2 * RoboDojo.hopper.nq 170 | nu = RoboDojo.hopper.nu + 4 + 4 + 2 + 4 + 1 171 | 172 | # ## model 173 | mass_matrix, dynamics_bias = RoboDojo.codegen_dynamics(RoboDojo.hopper) 174 | d1 = DTO.Dynamics((y, x, u, w) -> hopper_dyn1(mass_matrix, dynamics_bias, [h], y, x, u, w), 2 * nx + 5, nx, nu) 175 | dt = DTO.Dynamics((y, x, u, w) -> hopper_dynt(mass_matrix, dynamics_bias, [h], y, x, u, w), 2 * nx + 5, 2 * nx + 5, nu) 176 | 177 | dyn = [d1, [dt for t = 2:T-1]...]; 178 | 179 | # ## initial conditions 180 | q1 = [0.0; 0.5 + RoboDojo.hopper.foot_radius; 0.0; 0.5] 181 | qM = [0.5; 0.5 + RoboDojo.hopper.foot_radius; 0.0; 0.5] 182 | qT = [1.0; 0.5 + RoboDojo.hopper.foot_radius; 0.0; 0.5] 183 | q_ref = [0.5; 0.75 + RoboDojo.hopper.foot_radius; 0.0; 0.25] 184 | 185 | x1 = [q1; q1] 186 | xM = [qM; qM] 187 | xT = [qT; qT] 188 | x_ref = [q_ref; q_ref] 189 | 190 | # ## gate 191 | GAIT = 1 192 | GAIT = 2 193 | GAIT = 3 194 | 195 | if GAIT == 1 196 | r_cost = 1.0e-1 197 | q_cost = 1.0e-1 198 | elseif GAIT == 2 199 | r_cost = 1.0 200 | q_cost = 1.0 201 | elseif GAIT == 3 202 | r_cost = 1.0e-3 203 | q_cost = 1.0e-1 204 | end 205 | 206 | # ## objective 207 | function obj1(x, u, w) 208 | J = 0.0 209 | J += 0.5 * transpose(x - x_ref) * Diagonal([1.0; 10.0; 1.0; 10.0; 1.0; 10.0; 1.0; 10.0]) * (x - x_ref) 210 | J += 0.5 * transpose(u) * Diagonal([r_cost * ones(RoboDojo.hopper.nu); zeros(nu - RoboDojo.hopper.nu)]) * u 211 | J += 1000.0 * u[nu] 212 | return J 213 | end 214 | 215 | function objt(x, u, w) 216 | J = 0.0 217 | J += 0.5 * transpose(x[1:nx] - x_ref) * Diagonal(q_cost * [1.0; 10.0; 1.0; 10.0; 1.0; 10.0; 1.0; 10.0]) * (x[1:nx] - x_ref) 218 | J += 0.5 * transpose(u) * Diagonal([r_cost * ones(RoboDojo.hopper.nu); zeros(nu - RoboDojo.hopper.nu)]) * u 219 | J += 1000.0 * u[nu] 220 | return J 221 | end 222 | 223 | function objT(x, u, w) 224 | J = 0.0 225 | J += 0.5 * transpose(x[1:nx] - x_ref) * Diagonal([1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0]) * (x[1:nx] - x_ref) 226 | return J 227 | end 228 | 229 | c1 = DTO.Cost(obj1, nx, nu) 230 | ct = DTO.Cost(objt, 2 * nx + 5, nu) 231 | cT = DTO.Cost(objT, 2 * nx + 5, 0) 232 | obj = [c1, [ct for t = 2:T-1]..., cT]; 233 | 234 | # ## constraints 235 | 236 | ql = [-Inf; 0; -Inf; 0.0] 237 | qu = [Inf; Inf; Inf; 1.0] 238 | xl1 = [q1; ql] 239 | xu1 = [q1; qu] 240 | xlt = [ql; ql; -Inf * ones(5); -Inf * ones(nx)] 241 | xut = [qu; qu; Inf * ones(5); Inf * ones(nx)] 242 | ul = [-10.0; -10.0; zeros(nu - 2)] 243 | uu = [10.0; 10.0; Inf * ones(nu - 2)] 244 | 245 | bnd1 = DTO.Bound(nx, nu, xl=xl1, xu=xu1, ul=ul, uu=uu) 246 | bndt = DTO.Bound(2 * nx + 5, nu, xl=xlt, xu=xut, ul=ul, uu=uu) 247 | bndT = DTO.Bound(2 * nx + 5, 0, xl=xlt, xu=xut) 248 | bnds = [bnd1, [bndt for t = 2:T-1]..., bndT]; 249 | 250 | function constraints_1(x, u, w) 251 | [ 252 | # equality (8) 253 | RoboDojo.kinematics_foot(RoboDojo.hopper, x[1:RoboDojo.hopper.nq]) - RoboDojo.kinematics_foot(RoboDojo.hopper, x1[1:RoboDojo.hopper.nq]); 254 | RoboDojo.kinematics_foot(RoboDojo.hopper, x[RoboDojo.hopper.nq .+ (1:RoboDojo.hopper.nq)]) - RoboDojo.kinematics_foot(RoboDojo.hopper, x1[RoboDojo.hopper.nq .+ (1:RoboDojo.hopper.nq)]); 255 | contact_constraints_equality(h, x, u, w); 256 | # inequality (12) 257 | contact_constraints_inequality_1(h, x, u, w); 258 | ] 259 | end 260 | 261 | function constraints_t(x, u, w) 262 | [ 263 | # equality (4) 264 | contact_constraints_equality(h, x, u, w); 265 | # inequality (16) 266 | contact_constraints_inequality_t(h, x, u, w); 267 | ] 268 | end 269 | 270 | function constraints_T(x, u, w) 271 | x_travel = 0.5 272 | θ = x[nx + 5 .+ (1:nx)] 273 | [ 274 | # equality (6) 275 | x[1:RoboDojo.hopper.nq][collect([2, 3, 4])] - θ[1:RoboDojo.hopper.nq][collect([2, 3, 4])]; 276 | x[RoboDojo.hopper.nq .+ (1:RoboDojo.hopper.nq)][collect([2, 3, 4])] - θ[RoboDojo.hopper.nq .+ (1:RoboDojo.hopper.nq)][collect([2, 3, 4])]; 277 | # inequality (10) 278 | x_travel - (x[1] - θ[1]) 279 | x_travel - (x[RoboDojo.hopper.nq + 1] - θ[RoboDojo.hopper.nq + 1]) 280 | contact_constraints_inequality_T(h, x, u, w); 281 | ] 282 | end 283 | 284 | con1 = DTO.Constraint(constraints_1, nx, nu, idx_ineq=collect(8 .+ (1:12))) 285 | cont = DTO.Constraint(constraints_t, 2nx + 5, nu, idx_ineq=collect(4 .+ (1:16))) 286 | conT = DTO.Constraint(constraints_T, 2nx + 5, nu, idx_ineq=collect(6 .+ (1:10))) 287 | cons = [con1, [cont for t = 2:T-1]..., conT]; 288 | 289 | # ## problem 290 | p = DTO.solver(dyn, obj, cons, bnds, 291 | options=DTO.Options( 292 | tol=1.0e-2, 293 | constr_viol_tol=1.0e-2, 294 | print_level=0)) 295 | 296 | # ## initialize 297 | x_interpolation = [x1, [[x1; zeros(5); x1] for t = 2:T]...] 298 | u_guess = [[0.0; RoboDojo.hopper.gravity * RoboDojo.hopper.mass_body * 0.5 * h[1]; 1.0e-1 * ones(nu - 2)] for t = 1:T-1] # may need to run more than once to get good trajectory 299 | DTO.initialize_states!(p, x_interpolation) 300 | DTO.initialize_controls!(p, u_guess); 301 | 302 | # ## solve 303 | @time DTO.solve!(p); 304 | 305 | # ## solution 306 | x_sol, u_sol = DTO.get_trajectory(p) 307 | @show x_sol[1] 308 | @show x_sol[T] 309 | sum([u[nu] for u in u_sol[1:end-1]]) 310 | x_sol[1] - x_sol[T][1:nx] 311 | 312 | # ## visualize 313 | vis = Visualizer() 314 | render(vis) 315 | q_sol = state_to_configuration([x[1:nx] for x in x_sol]) 316 | RoboDojo.visualize!(vis, RoboDojo.hopper, q_sol, Δt=h); 317 | 318 | # ## comparison 319 | function obj1_compare(x, u, w) 320 | x = x[1:8] 321 | u = u[1:2] 322 | 323 | J = 0.0 324 | J += 0.5 * transpose(x - x_ref) * Diagonal([1.0; 10.0; 1.0; 10.0; 1.0; 10.0; 1.0; 10.0]) * (x - x_ref) 325 | J += 0.5 * transpose(u) * Diagonal(r_cost * ones(hopper.nu)) * u 326 | return J 327 | end 328 | 329 | function objt_compare(x, u, w) 330 | x = x[1:8] 331 | u = u[1:2] 332 | J = 0.0 333 | J += 0.5 * transpose(x - x_ref) * q_cost * Diagonal([1.0; 10.0; 1.0; 10.0; 1.0; 10.0; 1.0; 10.0]) * (x - x_ref) 334 | J += 0.5 * transpose(u) * Diagonal(r_cost * ones(hopper.nu)) * u 335 | return J 336 | end 337 | 338 | function objT_compare(x, u, w) 339 | x = x[1:8] 340 | J = 0.0 341 | J += 0.5 * transpose(x - x_ref) * Diagonal([1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0]) * (x - x_ref) 342 | return J 343 | end 344 | 345 | function obj_compare(x, u, w) 346 | J = 0.0 347 | J += obj1_compare(x[1], u[1], w[1]) 348 | for t = 2:T-1 349 | J += objt_compare(x[t], u[t], w[t]) 350 | end 351 | J += objT_compare(x[T], nothing, nothing) 352 | 353 | return J 354 | end 355 | 356 | obj_compare(x_sol, u_sol, [zeros(0) for t = 1:T]) 357 | @show norm(constraints_T(x_sol[T], zeros(0), zeros(0))[1:6], Inf) -------------------------------------------------------------------------------- /examples/generate_notebooks.jl: -------------------------------------------------------------------------------- 1 | using Literate 2 | 3 | function preprocess(str) 4 | str = replace(str, "# PREAMBLE" => "") 5 | str = replace(str, "# PKG_SETUP" => "") 6 | return str 7 | end 8 | 9 | exampledir = @__DIR__ 10 | Literate.notebook(joinpath(exampledir, "acrobot.jl"), exampledir, execute=false, preprocess=preprocess) 11 | Literate.notebook(joinpath(exampledir, "cartpole.jl"), exampledir, execute=false, preprocess=preprocess) 12 | Literate.notebook(joinpath(exampledir, "hopper.jl"), exampledir, execute=false, preprocess=preprocess) 13 | Literate.notebook(joinpath(exampledir, "planar_push.jl"), exampledir, execute=false, preprocess=preprocess) 14 | Literate.notebook(joinpath(exampledir, "rocket.jl"), exampledir, execute=false, preprocess=preprocess) 15 | 16 | -------------------------------------------------------------------------------- /examples/hopper.jl: -------------------------------------------------------------------------------- 1 | using OptimizationDynamics 2 | const iLQR = OptimizationDynamics.IterativeLQR 3 | const RoboDojo = OptimizationDynamics.RoboDojo 4 | using LinearAlgebra 5 | using Random 6 | 7 | # ## visualize 8 | vis = Visualizer() 9 | render(vis); 10 | 11 | # ## state-space model 12 | T = 21 13 | h = 0.05 14 | hopper = RoboDojo.hopper 15 | 16 | struct ParameterOptInfo{T} 17 | idx_q1::Vector{Int} 18 | idx_q2::Vector{Int} 19 | idx_u1::Vector{Int} 20 | idx_uθ::Vector{Int} 21 | idx_uθ1::Vector{Int} 22 | idx_uθ2::Vector{Int} 23 | idx_xθ::Vector{Int} 24 | v1::Vector{T} 25 | end 26 | 27 | info = ParameterOptInfo( 28 | collect(1:hopper.nq), 29 | collect(hopper.nq .+ (1:hopper.nq)), 30 | collect(1:hopper.nu), 31 | collect(hopper.nu .+ (1:2 * hopper.nq)), 32 | collect(hopper.nu .+ (1:hopper.nq)), 33 | collect(hopper.nu + hopper.nq .+ (1:hopper.nq)), 34 | collect(2 * hopper.nq .+ (1:2 * hopper.nq)), 35 | zeros(hopper.nq) 36 | ) 37 | 38 | im_dyn1 = ImplicitDynamics(hopper, h, 39 | eval(RoboDojo.residual_expr(hopper)), 40 | eval(RoboDojo.jacobian_var_expr(hopper)), 41 | eval(RoboDojo.jacobian_data_expr(hopper)); 42 | r_tol=1.0e-8, κ_eval_tol=1.0e-4, κ_grad_tol=1.0e-3, 43 | n=(2 * hopper.nq), m=(hopper.nu + 2 * hopper.nq), nc=4, nb=2, info=info) 44 | 45 | im_dynt = ImplicitDynamics(hopper, h, 46 | eval(RoboDojo.residual_expr(hopper)), 47 | eval(RoboDojo.jacobian_var_expr(hopper)), 48 | eval(RoboDojo.jacobian_data_expr(hopper)); 49 | r_tol=1.0e-8, κ_eval_tol=1.0e-4, κ_grad_tol=1.0e-3, 50 | n=4 * hopper.nq, m=hopper.nu, nc=4, nb=2, info=info) 51 | 52 | function f1(d, model::ImplicitDynamics, x, u, w) 53 | 54 | θ = @views u[model.info.idx_uθ] 55 | q1 = @views u[model.info.idx_uθ1] 56 | q2 = @views u[model.info.idx_uθ2] 57 | u1 = @views u[model.info.idx_u1] 58 | 59 | model.info.v1 .= q2 60 | model.info.v1 .-= q1 61 | model.info.v1 ./= model.eval_sim.h 62 | 63 | q3 = RoboDojo.step!(model.eval_sim, q2, model.info.v1, u1, 1) 64 | 65 | d[model.info.idx_q1] = q2 66 | d[model.info.idx_q2] = q3 67 | d[model.info.idx_xθ] = θ 68 | 69 | return d 70 | end 71 | 72 | function f1x(dx, model::ImplicitDynamics, x, u, w) 73 | dx .= 0.0 74 | return dx 75 | end 76 | 77 | function f1u(du, model::ImplicitDynamics, x, u, w) 78 | nq = model.grad_sim.model.nq 79 | 80 | θ = @views u[model.info.idx_uθ] 81 | q1 = @views u[model.info.idx_uθ1] 82 | q2 = @views u[model.info.idx_uθ2] 83 | u1 = @views u[model.info.idx_u1] 84 | 85 | model.info.v1 .= q2 86 | model.info.v1 .-= q1 87 | model.info.v1 ./= model.grad_sim.h 88 | 89 | RoboDojo.step!(model.grad_sim, q2, model.info.v1, u1, 1) 90 | 91 | for i = 1:nq 92 | du[model.info.idx_q1[i], model.info.idx_uθ[i]] = 1.0 93 | end 94 | du[model.info.idx_q2, model.info.idx_u1] = model.grad_sim.grad.∂q3∂u1[1] 95 | du[model.info.idx_q2, model.info.idx_uθ1] = model.grad_sim.grad.∂q3∂q1[1] 96 | du[model.info.idx_q2, model.info.idx_uθ2] = model.grad_sim.grad.∂q3∂q2[1] 97 | 98 | return du 99 | end 100 | 101 | function ft(d, model::ImplicitDynamics, x, u, w) 102 | 103 | θ = @views x[model.info.idx_xθ] 104 | q1 = @views x[model.info.idx_q1] 105 | q2 = @views x[model.info.idx_q2] 106 | u1 = u 107 | 108 | model.info.v1 .= q2 109 | model.info.v1 .-= q1 110 | model.info.v1 ./= model.eval_sim.h 111 | 112 | q3 = RoboDojo.step!(model.eval_sim, q2, model.info.v1, u1, 1) 113 | 114 | d[model.info.idx_q1] = q2 115 | d[model.info.idx_q2] = q3 116 | d[model.info.idx_xθ] = θ 117 | 118 | return d 119 | end 120 | 121 | function ftx(dx, model::ImplicitDynamics, x, u, w) 122 | nq = model.grad_sim.model.nq 123 | 124 | θ = @views x[model.info.idx_xθ] 125 | q1 = @views x[model.info.idx_q1] 126 | q2 = @views x[model.info.idx_q2] 127 | u1 = u 128 | 129 | model.info.v1 .= q2 130 | model.info.v1 .-= q1 131 | model.info.v1 ./= model.grad_sim.h 132 | 133 | q3 = RoboDojo.step!(model.grad_sim, q2, model.info.v1, u1, 1) 134 | 135 | for i = 1:nq 136 | dx[model.info.idx_q1[i], model.info.idx_q2[i]] = 1.0 137 | end 138 | dx[model.info.idx_q2, model.info.idx_q1] = model.grad_sim.grad.∂q3∂q1[1] 139 | dx[model.info.idx_q2, model.info.idx_q2] = model.grad_sim.grad.∂q3∂q2[1] 140 | for i in model.info.idx_xθ 141 | dx[i, i] = 1.0 142 | end 143 | 144 | return dx 145 | end 146 | 147 | function ftu(du, model::ImplicitDynamics, x, u, w) 148 | θ = @views x[model.info.idx_xθ] 149 | q1 = @views x[model.info.idx_q1] 150 | q2 = @views x[model.info.idx_q2] 151 | u1 = u 152 | 153 | model.info.v1 .= q2 154 | model.info.v1 .-= q1 155 | model.info.v1 ./= model.grad_sim.h 156 | 157 | q3 = RoboDojo.step!(model.grad_sim, q2, model.info.v1, u1, 1) 158 | 159 | du[model.info.idx_q2, model.info.idx_u1] = model.grad_sim.grad.∂q3∂u1[1] 160 | 161 | return du 162 | end 163 | 164 | # ## iLQR model 165 | ilqr_dyn1 = iLQR.Dynamics((d, x, u, w) -> f1(d, im_dyn1, x, u, w), 166 | (dx, x, u, w) -> f1x(dx, im_dyn1, x, u, w), 167 | (du, x, u, w) -> f1u(du, im_dyn1, x, u, w), 168 | 4 * hopper.nq, 2 * hopper.nq, hopper.nu + 2 * hopper.nq) 169 | 170 | ilqr_dynt = iLQR.Dynamics((d, x, u, w) -> ft(d, im_dynt, x, u, w), 171 | (dx, x, u, w) -> ftx(dx, im_dynt, x, u, w), 172 | (du, x, u, w) -> ftu(du, im_dynt, x, u, w), 173 | 4 * hopper.nq, 4 * hopper.nq, hopper.nu) 174 | 175 | model = [ilqr_dyn1, [ilqr_dynt for t = 2:T-1]...]; 176 | 177 | # ## initial conditions 178 | q1 = [0.0; 0.5 + hopper.foot_radius; 0.0; 0.5] 179 | qM = [0.5; 0.5 + hopper.foot_radius; 0.0; 0.5] 180 | qT = [1.0; 0.5 + hopper.foot_radius; 0.0; 0.5] 181 | q_ref = [0.5; 0.75 + hopper.foot_radius; 0.0; 0.25] 182 | 183 | x1 = [q1; q1] 184 | xM = [qM; qM] 185 | xT = [qT; qT] 186 | x_ref = [q_ref; q_ref] 187 | 188 | # ## objective 189 | 190 | GAIT = 1 191 | ## GAIT = 2 192 | ## GAIT = 3 193 | 194 | if GAIT == 1 195 | r_cost = 1.0e-1 196 | q_cost = 1.0e-1 197 | elseif GAIT == 2 198 | r_cost = 1.0 199 | q_cost = 1.0 200 | elseif GAIT == 3 201 | r_cost = 1.0e-3 202 | q_cost = 1.0e-1 203 | end 204 | 205 | function obj1(x, u, w) 206 | J = 0.0 207 | J += 0.5 * transpose(x - x_ref) * Diagonal([1.0; 10.0; 1.0; 10.0; 1.0; 10.0; 1.0; 10.0]) * (x - x_ref) 208 | J += 0.5 * transpose(u) * Diagonal([r_cost * ones(hopper.nu); 1.0e-1 * ones(hopper.nq); 1.0e-5 * ones(hopper.nq)]) * u 209 | return J 210 | end 211 | 212 | function objt(x, u, w) 213 | J = 0.0 214 | J += 0.5 * transpose(x - [x_ref; zeros(2 * hopper.nq)]) * q_cost * Diagonal([1.0; 10.0; 1.0; 10.0; 1.0; 10.0; 1.0; 10.0; zeros(2 * hopper.nq)]) * (x - [x_ref; zeros(2 * hopper.nq)]) 215 | J += 0.5 * transpose(u) * Diagonal(r_cost * ones(hopper.nu)) * u 216 | return J 217 | end 218 | 219 | function objT(x, u, w) 220 | J = 0.0 221 | J += 0.5 * transpose(x - [x_ref; zeros(2 * hopper.nq)]) * Diagonal([1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0; zeros(2 * hopper.nq)]) * (x - [x_ref; zeros(2 * hopper.nq)]) 222 | return J 223 | end 224 | 225 | c1 = iLQR.Cost(obj1, 2 * hopper.nq, hopper.nu + 2 * hopper.nq) 226 | ct = iLQR.Cost(objt, 4 * hopper.nq, hopper.nu) 227 | cT = iLQR.Cost(objT, 4 * hopper.nq, 0) 228 | obj = [c1, [ct for t = 2:T-1]..., cT]; 229 | 230 | # ## constraints 231 | ul = [-10.0; -10.0] 232 | uu = [10.0; 10.0] 233 | 234 | function stage1_con(x, u, w) 235 | [ 236 | ul - u[1:hopper.nu]; # control limit (lower) 237 | u[1:hopper.nu] - uu; # control limit (upper) 238 | 239 | u[hopper.nu .+ (1:hopper.nq)] - x1[1:hopper.nq]; 240 | 241 | RoboDojo.kinematics_foot(hopper, u[hopper.nu .+ (1:hopper.nq)]) - RoboDojo.kinematics_foot(hopper, x1[1:hopper.nq]); 242 | RoboDojo.kinematics_foot(hopper, u[hopper.nu + hopper.nq .+ (1:hopper.nq)]) - RoboDojo.kinematics_foot(hopper, x1[hopper.nq .+ (1:hopper.nq)]) 243 | ] 244 | end 245 | 246 | function staget_con(x, u, w) 247 | [ 248 | ul - u[collect(1:hopper.nu)]; # control limit (lower) 249 | u[collect(1:hopper.nu)] - uu; # control limit (upper) 250 | ] 251 | end 252 | 253 | function terminal_con(x, u, w) 254 | x_travel = 0.5 255 | θ = x[2 * hopper.nq .+ (1:(2 * hopper.nq))] 256 | [ 257 | x_travel - (x[1] - θ[1]) 258 | x_travel - (x[hopper.nq + 1] - θ[hopper.nq + 1]) 259 | x[1:hopper.nq][collect([2, 3, 4])] - θ[1:hopper.nq][collect([2, 3, 4])] 260 | x[hopper.nq .+ (1:hopper.nq)][collect([2, 3, 4])] - θ[hopper.nq .+ (1:hopper.nq)][collect([2, 3, 4])] 261 | ] 262 | end 263 | 264 | con1 = iLQR.Constraint(stage1_con, 2 * hopper.nq, hopper.nu + 2 * hopper.nq, idx_ineq=collect(1:4)) 265 | cont = iLQR.Constraint(staget_con, 4 * hopper.nq, hopper.nu, idx_ineq=collect(1:4)) 266 | conT = iLQR.Constraint(terminal_con, 4 * hopper.nq, 0, idx_ineq=collect(1:2)) 267 | cons = [con1, [cont for t = 2:T-1]..., conT]; 268 | 269 | # ## rollout 270 | ū_stand = [t == 1 ? [0.0; hopper.gravity * hopper.mass_body * 0.5 * h; x1] : [0.0; hopper.gravity * hopper.mass_body * 0.5 * h] for t = 1:T-1] 271 | x̄ = iLQR.rollout(model, x1, ū_stand) 272 | q̄ = state_to_configuration(x̄) 273 | RoboDojo.visualize!(vis, hopper, x̄, Δt=h); 274 | 275 | # ## solver 276 | solver = iLQR.solver(model, obj, cons, 277 | opts=iLQR.Options(linesearch = :armijo, 278 | α_min=1.0e-5, 279 | obj_tol=1.0e-3, 280 | grad_tol=1.0e-3, 281 | max_iter=10, 282 | max_al_iter=15, 283 | con_tol=0.001, 284 | ρ_init=1.0, 285 | ρ_scale=10.0, 286 | verbose=false)) 287 | iLQR.initialize_controls!(solver, ū_stand) 288 | iLQR.initialize_states!(solver, x̄); 289 | 290 | # ## solve 291 | iLQR.reset!(solver.s_data) 292 | @time iLQR.solve!(solver); 293 | 294 | @show iLQR.eval_obj(solver.m_data.obj.costs, solver.m_data.x, solver.m_data.u, solver.m_data.w) 295 | @show solver.s_data.iter[1] 296 | @show norm(terminal_con(solver.m_data.x[T], zeros(0), zeros(0))[3:4], Inf) 297 | @show solver.s_data.obj[1] # augmented Lagrangian cost 298 | 299 | # ## solution 300 | x_sol, u_sol = iLQR.get_trajectory(solver) 301 | q_sol = state_to_configuration(x_sol) 302 | RoboDojo.visualize!(vis, hopper, q_sol, Δt=h); 303 | 304 | # ## benchmark (NOTE: gate 3 seems to break @benchmark, just run @time instead...) 305 | solver.options.verbose = false 306 | @benchmark iLQR.solve!($solver, $x̄, $ū_stand); 307 | 308 | 309 | -------------------------------------------------------------------------------- /examples/planar_push.jl: -------------------------------------------------------------------------------- 1 | using OptimizationDynamics 2 | const iLQR = OptimizationDynamics.IterativeLQR 3 | using LinearAlgebra 4 | using Random 5 | 6 | # ## visualization 7 | vis = Visualizer() 8 | render(vis); 9 | 10 | # ## mode 11 | MODE = :translate 12 | MODE = :rotate 13 | 14 | # ## gradient bundle 15 | GB = false 16 | 17 | # ## state-space model 18 | h = 0.1 19 | T = 26 20 | 21 | im_dyn = ImplicitDynamics(planarpush, h, eval(r_pp_func), eval(rz_pp_func), eval(rθ_pp_func); 22 | r_tol=1.0e-8, κ_eval_tol=1.0e-4, κ_grad_tol=1.0e-2, nc=1, nb=9, info=(GB ? GradientBundle(planarpush, N=50, ϵ=1.0e-4) : nothing)) 23 | 24 | nx = 2 * planarpush.nq 25 | nu = planarpush.nu 26 | 27 | # ## iLQR model 28 | ilqr_dyn = iLQR.Dynamics((d, x, u, w) -> f(d, im_dyn, x, u, w), 29 | (dx, x, u, w) -> GB ? fx_gb(dx, im_dyn, x, u, w) : fx(dx, im_dyn, x, u, w), 30 | (du, x, u, w) -> GB ? fu_gb(du, im_dyn, x, u, w) : fu(du, im_dyn, x, u, w), 31 | nx, nx, nu) 32 | 33 | model = [ilqr_dyn for t = 1:T-1]; 34 | 35 | # ## initial conditions and goal 36 | r_dim = 0.1 37 | if MODE == :translate 38 | q0 = [0.0, 0.0, 0.0, -r_dim - 1.0e-8, 0.0] 39 | q1 = [0.0, 0.0, 0.0, -r_dim - 1.0e-8, 0.0] 40 | x_goal = 1.0 41 | y_goal = 0.0 42 | θ_goal = 0.0 * π 43 | qT = [x_goal, y_goal, θ_goal, x_goal - r_dim, y_goal - r_dim] 44 | xT = [qT; qT] 45 | elseif MODE == :rotate 46 | q0 = [0.0, 0.0, 0.0, -r_dim - 1.0e-8, -0.01] 47 | q1 = [0.0, 0.0, 0.0, -r_dim - 1.0e-8, -0.01] 48 | x1 = [q1; q1] 49 | x_goal = 0.5 50 | y_goal = 0.5 51 | θ_goal = 0.5 * π 52 | qT = [x_goal, y_goal, θ_goal, x_goal-r_dim, y_goal-r_dim] 53 | xT = [qT; qT] 54 | end 55 | 56 | # ## objective 57 | function objt(x, u, w) 58 | J = 0.0 59 | 60 | q1 = x[1:planarpush.nq] 61 | q2 = x[planarpush.nq .+ (1:planarpush.nq)] 62 | v1 = (q2 - q1) ./ h 63 | 64 | J += 0.5 * transpose(v1) * Diagonal([1.0, 1.0, 1.0, 0.1, 0.1]) * v1 65 | J += 0.5 * transpose(x - xT) * Diagonal([1.0, 1.0, 1.0, 0.1, 0.1, 1.0, 1.0, 1.0, 0.1, 0.1]) * (x - xT) 66 | J += 0.5 * (MODE == :translate ? 1.0e-1 : 1.0e-2) * transpose(u) * u 67 | 68 | return J 69 | end 70 | 71 | function objT(x, u, w) 72 | J = 0.0 73 | 74 | q1 = x[1:planarpush.nq] 75 | q2 = x[planarpush.nq .+ (1:planarpush.nq)] 76 | v1 = (q2 - q1) ./ h 77 | 78 | J += 0.5 * transpose(v1) * Diagonal([1.0, 1.0, 1.0, 0.1, 0.1]) * v1 79 | J += 0.5 * transpose(x - xT) * Diagonal([1.0, 1.0, 1.0, 0.1, 0.1, 1.0, 1.0, 1.0, 0.1, 0.1]) * (x - xT) 80 | 81 | return J 82 | end 83 | 84 | ct = iLQR.Cost(objt, nx, nu) 85 | cT = iLQR.Cost(objT, nx, 0) 86 | obj = [[ct for t = 1:T-1]..., cT]; 87 | 88 | # ## constraints 89 | ul = [-5.0; -5.0] 90 | uu = [5.0; 5.0] 91 | 92 | function stage_con(x, u, w) 93 | [ 94 | ul - u; # control limit (lower) 95 | u - uu; # control limit (upper) 96 | ] 97 | end 98 | 99 | function terminal_con(x, u, w) 100 | [ 101 | (x - xT)[collect([(1:3)..., (6:8)...])]; # goal 102 | ] 103 | end 104 | 105 | cont = iLQR.Constraint(stage_con, nx, nu, idx_ineq=collect(1:(2 * nu))) 106 | conT = iLQR.Constraint(terminal_con, nx, 0) 107 | cons = [[cont for t = 1:T-1]..., conT]; 108 | 109 | # ## rollout 110 | x1 = [q0; q1] 111 | ū = MODE == :translate ? [t < 5 ? [1.0; 0.0] : [0.0; 0.0] for t = 1:T-1] : [t < 5 ? [1.0; 0.0] : t < 10 ? [0.5; 0.0] : [0.0; 0.0] for t = 1:T-1] 112 | x̄ = iLQR.rollout(model, x1, ū) 113 | q̄ = state_to_configuration(x̄) 114 | visualize!(vis, planarpush, q̄, Δt=h); 115 | 116 | # ## solver 117 | solver = iLQR.solver(model, obj, cons, 118 | opts=iLQR.Options( 119 | linesearch = :armijo, 120 | α_min=1.0e-5, 121 | obj_tol=1.0e-3, 122 | grad_tol=1.0e-3, 123 | max_iter=10, 124 | max_al_iter=10, 125 | con_tol=0.005, 126 | ρ_init=1.0, 127 | ρ_scale=10.0, 128 | verbose=false)) 129 | iLQR.initialize_controls!(solver, ū) 130 | iLQR.initialize_states!(solver, x̄); 131 | 132 | # ## solve 133 | iLQR.reset!(solver.s_data) 134 | @time iLQR.solve!(solver); 135 | 136 | @show iLQR.eval_obj(solver.m_data.obj.costs, solver.m_data.x, solver.m_data.u, solver.m_data.w) 137 | @show solver.s_data.iter[1] 138 | @show norm(terminal_con(solver.m_data.x[T], zeros(0), zeros(0)), Inf) 139 | @show solver.s_data.obj[1] # augmented Lagrangian cost 140 | 141 | # ## solution 142 | x_sol, u_sol = iLQR.get_trajectory(solver) 143 | q_sol = state_to_configuration(x_sol) 144 | visualize!(vis, planarpush, q_sol, Δt=h); 145 | 146 | # ## benchmark 147 | solver.options.verbose = false 148 | @benchmark iLQR.solve!($solver, x̄, ū) setup=(x̄=deepcopy(x̄), ū=deepcopy(ū)); 149 | -------------------------------------------------------------------------------- /examples/rocket.jl: -------------------------------------------------------------------------------- 1 | using OptimizationDynamics 2 | const iLQR = OptimizationDynamics.IterativeLQR 3 | const Rotations = OptimizationDynamics.Rotations 4 | using LinearAlgebra 5 | using Random 6 | 7 | # ## visualization 8 | vis = Visualizer() 9 | render(vis) 10 | 11 | # ## mode 12 | MODE = :projection 13 | MODE = :nominal 14 | 15 | # ## thrust max 16 | u_max = 12.5 17 | 18 | # ## state-space model 19 | h = 0.05 20 | T = 61 21 | info = RocketInfo(rocket, u_max, h, 22 | eval(r_rocket_func), eval(rz_rocket_func), eval(rθ_rocket_func), 23 | eval(r_proj_func), eval(rz_proj_func), eval(rθ_proj_func)) 24 | 25 | nx = rocket.nq 26 | nu = rocket.nu 27 | 28 | # ## iLQR model 29 | if MODE == :projection 30 | ilqr_dyn = iLQR.Dynamics((d, x, u, w) -> f_rocket_proj(d, info, x, u, w), 31 | (dx, x, u, w) -> fx_rocket_proj(dx, info, x, u, w), 32 | (du, x, u, w) -> fu_rocket_proj(du, info, x, u, w), 33 | nx, nx, nu) 34 | else 35 | ilqr_dyn = iLQR.Dynamics((d, x, u, w) -> f_rocket(d, info, x, u, w), 36 | (dx, x, u, w) -> fx_rocket(dx, info, x, u, w), 37 | (du, x, u, w) -> fu_rocket(du, info, x, u, w), 38 | nx, nx, nu) 39 | end 40 | 41 | model = [ilqr_dyn for t = 1:T-1]; 42 | 43 | # ## initial conditions 44 | x1 = zeros(rocket.nq) 45 | x1[1] = 2.5 46 | x1[2] = 2.5 47 | x1[3] = 10.0 48 | mrp = Rotations.MRP(Rotations.RotZ(0.25 * π) * Rotations.RotY(-0.5 * π)) 49 | x1[4:6] = [mrp.x; mrp.y; mrp.z] 50 | x1[9] = -1.0 51 | 52 | xT = zeros(rocket.nq) 53 | xT[3] = rocket.length 54 | mrpT = Rotations.MRP(Rotations.RotZ(0.25 * π) * Rotations.RotY(0.0)) 55 | xT[4:6] = [mrpT.x; mrpT.y; mrpT.z] 56 | 57 | # ## objective 58 | function objt(x, u, w) 59 | J = 0.0 60 | 61 | J += 0.5 * transpose(x - xT) * Diagonal(h * [1.0e-1 * ones(3); 1.0e-5 * ones(3); 1.0e-1 * ones(3); 1000.0 * ones(3)]) * (x - xT) 62 | J += 0.5 * transpose(u) * Diagonal(h * [1000.0; 1000.0; 100.0]) * u 63 | 64 | return J 65 | end 66 | 67 | function objT(x, u, w) 68 | J = 0.0 69 | 70 | J += 0.5 * transpose(x - xT) * Diagonal(h * 1000.0 * ones(rocket.nq)) * (x - xT) 71 | 72 | return J 73 | end 74 | 75 | ct = iLQR.Cost(objt, nx, nu) 76 | cT = iLQR.Cost(objT, nx, 0) 77 | obj = [[ct for t = 1:T-1]..., cT]; 78 | 79 | # ## constraints 80 | x_con = [-0.5; 0.5] 81 | y_con = [-0.75; 0.75] 82 | 83 | function stage_con(x, u, w) 84 | if MODE == :projection 85 | [ 86 | rocket.length - x[3]; 87 | ] 88 | else 89 | [ 90 | -1.0 - u[1]; 91 | u[1] - 1.0; 92 | -1.0 - u[2]; 93 | u[2] - 1.0; 94 | 0.0 - u[3]; # control limit (lower) 95 | u[3] - u_max; # control limit (upper) 96 | rocket.length - x[3]; 97 | ] 98 | end 99 | end 100 | 101 | function terminal_con(x, u, w) 102 | [ 103 | x_con[1] - x[1]; 104 | x[1] - x_con[2]; 105 | y_con[1] - x[2]; 106 | x[2] - y_con[2]; 107 | (x - xT)[collect([3, 4, 5, 6, 7, 8, 9, 10, 11, 12])] 108 | ] 109 | end 110 | 111 | cont = iLQR.Constraint(stage_con, nx, nu, idx_ineq=collect(1:(1 + (MODE == :projection ? 0 : 6)))) 112 | conT = iLQR.Constraint(terminal_con, nx, 0, idx_ineq=collect(1:4)) 113 | cons = [[cont for t = 1:T-1]..., conT]; 114 | 115 | # ## rollout 116 | Random.seed!(1) 117 | ū = [1.0e-3 * randn(nu) for t = 1:T-1] 118 | x̄ = iLQR.rollout(model, x1, ū) 119 | visualize!(vis, rocket, x̄, Δt=h) 120 | 121 | # ## solver 122 | solver = iLQR.solver(model, obj, cons, 123 | opts=iLQR.Options( 124 | linesearch = :armijo, 125 | α_min=1.0e-5, 126 | obj_tol=1.0e-3, 127 | grad_tol=1.0e-3, 128 | max_iter=100, 129 | max_al_iter=15, 130 | con_tol=0.005, 131 | ρ_init=1.0, 132 | ρ_scale=10.0, 133 | verbose=false)) 134 | iLQR.initialize_controls!(solver, ū) 135 | iLQR.initialize_states!(solver, x̄); 136 | 137 | # ## solve 138 | iLQR.reset!(solver.s_data) 139 | @time iLQR.solve!(solver); 140 | 141 | @show iLQR.eval_obj(solver.m_data.obj.costs, solver.m_data.x, solver.m_data.u, solver.m_data.w) 142 | @show solver.s_data.iter[1] 143 | @show norm(terminal_con(solver.m_data.x[T], zeros(0), zeros(0))[4 .+ (1:10)], Inf) 144 | @show solver.s_data.obj[1] # augmented Lagrangian cost 145 | 146 | # ## solution 147 | x_sol, u_sol = iLQR.get_trajectory(solver) 148 | visualize!(vis, rocket, x_sol, Δt=h) 149 | open(vis) 150 | # ## test thrust cone constraint 151 | all([norm(u[1:2]) <= u[3] for u in u_sol]) 152 | 153 | # ## benchmark 154 | solver.options.verbose = false 155 | @benchmark iLQR.solve!($solver, x̄, ū) setup=(x̄=deepcopy(x̄), ū=deepcopy(ū)); 156 | 157 | -------------------------------------------------------------------------------- /src/OptimizationDynamics.jl: -------------------------------------------------------------------------------- 1 | module OptimizationDynamics 2 | 3 | using LinearAlgebra 4 | using BenchmarkTools 5 | using Symbolics 6 | using IfElse 7 | using JLD2 8 | import JLD2: load 9 | using DirectTrajectoryOptimization 10 | using IterativeLQR 11 | using MeshCat 12 | using Colors 13 | using CoordinateTransformations 14 | using GeometryBasics 15 | using Rotations 16 | using RoboDojo 17 | import RoboDojo: LinearSolver, LUSolver, Model, ResidualMethods, Space, Disturbances, IndicesZ, InteriorPoint, EmptySolver, Policy, Trajectory, GradientTrajectory, InteriorPointOptions, IndicesOptimization, interior_point, interior_point_solve!, bilinear_violation, residual_violation, general_correction_term!, r!, rz!, rθ!, linear_solve!, lu_solver, empty_policy, empty_disturbances, friction_coefficients, SimulatorStatistics, SimulatorOptions, indices_θ, num_data, initialize_z!, initialize_θ!, indices_z, indices_θ, simulate!, policy, process!, Simulator, cone_product, lagrangian_derivatives, Indicesθ 18 | using Scratch 19 | 20 | export LinearSolver, LUSolver, Model, ResidualMethods, Space, Disturbances, IndicesZ, InteriorPoint, EmptySolver, Policy, Trajectory, GradientTrajectory, InteriorPointOptions, IndicesOptimization, interior_point, interior_point_solve!, bilinear_violation, residual_violation, general_correction_term!, r!, rz!, rθ!, linear_solve!, lu_solver, empty_policy, empty_disturbances, friction_coefficients, SimulatorStatistics, SimulatorOptions, indices_θ, num_data, initialize_z!, initialize_θ!, indices_z, indices_θ, simulate!, policy, process!, Simulator, cone_product, lagrangian_derivatives, Indicesθ 21 | 22 | export 23 | load 24 | 25 | export 26 | Visualizer, render, open, visualize! 27 | 28 | include("dynamics.jl") 29 | include("ls.jl") 30 | include("gradient_bundle.jl") 31 | 32 | export 33 | ImplicitDynamics, f, fx, fu, state_to_configuration, 34 | f_gb, fx_gb, fu_gb 35 | 36 | # acrobot 37 | include("../src/models/acrobot/model.jl") 38 | include("../src/models/acrobot/simulator_impact.jl") 39 | include("../src/models/acrobot/simulator_nominal.jl") 40 | include("../src/models/acrobot/visuals.jl") 41 | path_acrobot = @get_scratch!("acrobot") 42 | @load joinpath(path_acrobot, "impact.jld2") r_acrobot_impact_func rz_acrobot_impact_func rθ_acrobot_impact_func rz_acrobot_impact_array rθ_acrobot_impact_array 43 | @load joinpath(path_acrobot, "nominal.jld2") r_acrobot_nominal_func rz_acrobot_nominal_func rθ_acrobot_nominal_func rz_acrobot_nominal_array rθ_acrobot_nominal_array 44 | 45 | # cartpole 46 | include("../src/models/cartpole/model.jl") 47 | include("../src/models/cartpole/simulator_friction.jl") 48 | include("../src/models/cartpole/simulator_frictionless.jl") 49 | include("../src/models/cartpole/visuals.jl") 50 | path_cartpole = @get_scratch!("cartpole") 51 | @load joinpath(path_cartpole, "friction.jld2") r_cartpole_friction_func rz_cartpole_friction_func rθ_cartpole_friction_func rz_cartpole_friction_array rθ_cartpole_friction_array 52 | @load joinpath(path_cartpole, "frictionless.jld2") r_cartpole_frictionless_func rz_cartpole_frictionless_func rθ_cartpole_frictionless_func rz_cartpole_frictionless_array rθ_cartpole_frictionless_array 53 | 54 | 55 | # hopper from RoboDojo.jl 56 | 57 | # planar push 58 | include("../src/models/planar_push/model.jl") 59 | include("../src/models/planar_push/simulator.jl") 60 | include("../src/models/planar_push/visuals.jl") 61 | path_planarpush = @get_scratch!("planarpush") 62 | @load joinpath(path_planarpush, "residual.jld2") r_pp_func rz_pp_func rθ_pp_func rz_pp_array rθ_pp_array 63 | 64 | # rocket 65 | include("../src/models/rocket/model.jl") 66 | include("../src/models/rocket/simulator.jl") 67 | include("../src/models/rocket/dynamics.jl") 68 | include("../src/models/rocket/visuals.jl") 69 | path_rocket = @get_scratch!("rocket") 70 | @load joinpath(path_rocket, "residual.jld2") r_rocket_func rz_rocket_func rθ_rocket_func rz_rocket_array rθ_rocket_array 71 | @load joinpath(path_rocket, "projection.jld2") r_proj_func rz_proj_func rθ_proj_func rz_proj_array rθ_proj_array 72 | 73 | include("../src/models/visualize.jl") 74 | 75 | export 76 | acrobot_impact, acrobot_nominal, 77 | cartpole_friction, cartpole_frictionless, 78 | planarpush, 79 | rocket, RocketInfo, f_rocket_proj, fx_rocket_proj, fu_rocket_proj, f_rocket, fx_rocket, fu_rocket 80 | 81 | export 82 | r_acrobot_impact_func, rz_acrobot_impact_func, rθ_acrobot_impact_func, rz_acrobot_impact_array, rθ_acrobot_impact_array, 83 | r_acrobot_nominal_func, rz_acrobot_nominal_func, rθ_acrobot_nominal_func, rz_acrobot_nominal_array, rθ_acrobot_nominal_array, 84 | r_cartpole_friction_func, rz_cartpole_friction_func, rθ_cartpole_friction_func, rz_cartpole_friction_array, rθ_cartpole_friction_array, 85 | r_cartpole_frictionless_func, rz_cartpole_frictionless_func, rθ_cartpole_frictionless_func, rz_cartpole_frictionless_array, rθ_cartpole_frictionless_array, 86 | r_pp_func, rz_pp_func, rθ_pp_func, rz_pp_array, rθ_pp_array, 87 | r_rocket_func, rz_rocket_func, rθ_rocket_func, rz_rocket_array, rθ_rocket_array, 88 | r_proj_func, rz_proj_func, rθ_proj_func, rz_proj_array, rθ_proj_array 89 | 90 | end # module 91 | -------------------------------------------------------------------------------- /src/dynamics.jl: -------------------------------------------------------------------------------- 1 | struct ImplicitDynamics{T,R,RZ,Rθ,M<:RoboDojo.Model{T},P<:RoboDojo.Policy{T},D<:RoboDojo.Disturbances{T},I} <: Model{T} 2 | n::Int 3 | m::Int 4 | d::Int 5 | eval_sim::Simulator{T,R,RZ,Rθ,M,P,D} 6 | grad_sim::Simulator{T,R,RZ,Rθ,M,P,D} 7 | q1::Vector{T} 8 | q2::Vector{T} 9 | v1::Vector{T} 10 | idx_q1::Vector{Int} 11 | idx_q2::Vector{Int} 12 | idx_u1::Vector{Int} 13 | info::I 14 | end 15 | 16 | function get_simulator(model, h, r_func, rz_func, rθ_func; 17 | T=1, r_tol=1.0e-8, κ_eval_tol=1.0e-4, nc=model.nc, nb=model.nc, diff_sol=true) 18 | 19 | sim = Simulator(model, T; 20 | h=h, 21 | residual=r_func, 22 | jacobian_z=rz_func, 23 | jacobian_θ=rθ_func, 24 | diff_sol=diff_sol, 25 | solver_opts=InteriorPointOptions( 26 | undercut=Inf, 27 | γ_reg=0.1, 28 | r_tol=r_tol, 29 | κ_tol=κ_eval_tol, 30 | max_ls=25, 31 | ϵ_min=0.25, 32 | diff_sol=diff_sol, 33 | verbose=false)) 34 | 35 | # set trajectory sizes 36 | sim.traj.γ .= [zeros(nc) for t = 1:T] 37 | sim.traj.b .= [zeros(nb) for t = 1:T] 38 | 39 | sim.grad.∂γ1∂q1 .= [zeros(nc, model.nq) for t = 1:T] 40 | sim.grad.∂γ1∂q2 .= [zeros(nc, model.nq) for t = 1:T] 41 | sim.grad.∂γ1∂v1 .= [zeros(nc, model.nq) for t = 1:T] 42 | sim.grad.∂γ1∂u1 .= [zeros(nc, model.nu) for t = 1:T] 43 | sim.grad.∂b1∂q1 .= [zeros(nb, model.nq) for t = 1:T] 44 | sim.grad.∂b1∂q2 .= [zeros(nb, model.nq) for t = 1:T] 45 | sim.grad.∂b1∂v1 .= [zeros(nb, model.nq) for t = 1:T] 46 | sim.grad.∂b1∂u1 .= [zeros(nb, model.nu) for t = 1:T] 47 | 48 | return sim 49 | end 50 | 51 | function ImplicitDynamics(model, h, r_func, rz_func, rθ_func; 52 | T=1, r_tol=1.0e-8, κ_eval_tol=1.0e-6, κ_grad_tol=1.0e-6, 53 | no_impact=false, no_friction=false, 54 | n=(2 * model.nq), m=model.nu, d=model.nw, nc=model.nc, nb=model.nc, 55 | info=nothing) 56 | 57 | # set trajectory sizes 58 | no_impact && (nc = 0) 59 | no_friction && (nb = 0) 60 | 61 | eval_sim = get_simulator(model, h, r_func, rz_func, rθ_func; 62 | T=T, r_tol=r_tol, κ_eval_tol=κ_eval_tol, nc=nc, nb=nb, diff_sol=false) 63 | 64 | grad_sim = get_simulator(model, h, r_func, rz_func, rθ_func; 65 | T=T, r_tol=r_tol, κ_eval_tol=κ_grad_tol, nc=nc, nb=nb, diff_sol=true) 66 | 67 | q1 = zeros(model.nq) 68 | q2 = zeros(model.nq) 69 | v1 = zeros(model.nq) 70 | 71 | idx_q1 = collect(1:model.nq) 72 | idx_q2 = collect(model.nq .+ (1:model.nq)) 73 | idx_u1 = collect(1:model.nu) 74 | 75 | ImplicitDynamics(n, m, d, 76 | eval_sim, grad_sim, 77 | q1, q2, v1, 78 | idx_q1, idx_q2, idx_u1, info) 79 | end 80 | 81 | function f(d, model::ImplicitDynamics, x, u, w) 82 | q1 = @views x[model.idx_q1] 83 | q2 = @views x[model.idx_q2] 84 | model.v1 .= q2 85 | model.v1 .-= q1 86 | model.v1 ./= model.eval_sim.h 87 | 88 | q3 = RoboDojo.step!(model.eval_sim, q2, model.v1, u, 1) 89 | 90 | d[model.idx_q1] .= q2 91 | d[model.idx_q2] .= q3 92 | 93 | return d 94 | end 95 | 96 | function fx(dx, model::ImplicitDynamics, x, u, w) 97 | q1 = @views x[model.idx_q1] 98 | q2 = @views x[model.idx_q2] 99 | model.v1 .= q2 100 | model.v1 .-= q1 101 | model.v1 ./= model.grad_sim.h 102 | 103 | RoboDojo.step!(model.grad_sim, q2, model.v1, u, 1) 104 | 105 | nq = model.grad_sim.model.nq 106 | for i = 1:nq 107 | dx[model.idx_q1[i], model.idx_q2[i]] = 1.0 108 | end 109 | 110 | dx[model.idx_q2, model.idx_q1] .= model.grad_sim.grad.∂q3∂q1[1] 111 | dx[model.idx_q2, model.idx_q2] .= model.grad_sim.grad.∂q3∂q2[1] 112 | 113 | return dx 114 | end 115 | 116 | function fu(du, model::ImplicitDynamics, x, u, w) 117 | q1 = @views x[model.idx_q1] 118 | q2 = @views x[model.idx_q2] 119 | model.v1 .= q2 120 | model.v1 .-= q1 121 | model.v1 ./= model.grad_sim.h 122 | 123 | RoboDojo.step!(model.grad_sim, q2, model.v1, u, 1) 124 | 125 | du[model.idx_q2, :] .= model.grad_sim.grad.∂q3∂u1[1] 126 | 127 | return du 128 | end 129 | 130 | 131 | function state_to_configuration(x::Vector{Vector{T}}) where T 132 | H = length(x) 133 | n = length(x[1]) 134 | nq = convert(Int, floor(length(x[1]) / 2)) 135 | q = Vector{T}[] 136 | 137 | for t = 1:H 138 | if t == 1 139 | push!(q, x[t][1:nq]) 140 | end 141 | push!(q, x[t][nq .+ (1:nq)]) 142 | end 143 | 144 | return q 145 | end 146 | 147 | # using BenchmarkTools 148 | # using InteractiveUtils 149 | # x = x̄[1] 150 | # u = ū[1] 151 | # w_ = w[1] 152 | # d = zeros(nx) 153 | # dx = zeros(nx, nx) 154 | # du = zeros(nx, nu) 155 | 156 | # f(d, im_dyn, x, u, w_) 157 | # @benchmark f($d, $im_dyn, $x, $u, $w_) 158 | # @code_warntype f(d, im_dyn, x, u, w_) 159 | 160 | # fx(dx, im_dyn, x, u, w_) 161 | # @benchmark fx($dx, $im_dyn, $x, $u, $w_) 162 | # @code_warntype fx(dx, im_dyn, x, u, w_) 163 | 164 | # fu(du, im_dyn, x, u, w_) 165 | # @benchmark fu($du, $im_dyn, $x, $u, $w_) 166 | # @code_warntype fu(du, im_dyn, x, u, w_) 167 | -------------------------------------------------------------------------------- /src/gradient_bundle.jl: -------------------------------------------------------------------------------- 1 | struct MInfo{T} 2 | idx_q1::Vector{Int} 3 | idx_q2::Vector{Int} 4 | idx_u1::Vector{Int} 5 | v1::Vector{T} 6 | end 7 | 8 | function _step(sim::Simulator, info::MInfo, q1, q2, u1) 9 | info.v1 .= q2 10 | info.v1 .-= q1 11 | info.v1 ./= sim.h 12 | RoboDojo.step!(sim, q2, info.v1, u1, 1) 13 | end 14 | 15 | struct GradientBundle{T} 16 | ls::LeastSquares{T} 17 | q1η::Vector{T} 18 | q2η::Vector{T} 19 | u1η::Vector{T} 20 | dz::Matrix{T} 21 | ny::Int 22 | nz::Int 23 | info::MInfo{T} 24 | end 25 | 26 | function GradientBundle(model; N=100, ϵ=1.0e-4) 27 | nx = 2 * model.nq 28 | ny = model.nq 29 | nu = model.nu 30 | nz = nx + nu 31 | nθ = ny * (nz) 32 | 33 | @variables fz[1:ny] fη[1:ny] η[1:nz] θ[1:nθ] 34 | 35 | function cost(fz, fη, η, θ) 36 | M = reshape(θ, ny, nz) 37 | r = fη - fz - M * η 38 | return [transpose(r) * r] 39 | end 40 | 41 | c = cost(fz, fη, η, θ) 42 | cθ = Symbolics.gradient(c[1], θ) 43 | cθθ = Symbolics.hessian(c[1], θ) 44 | 45 | c_func_ls = eval(Symbolics.build_function(c, fz, fη, η, θ)[2]) 46 | cθ_func_ls = eval(Symbolics.build_function(cθ, fz, fη, η, θ)[2]) 47 | cθθ_func_ls = eval(Symbolics.build_function(cθθ, fz, fη, η, θ)[2]) 48 | 49 | η = [zeros(nz) for i = 1:N] 50 | 51 | for i = 1:N 52 | w = ϵ * randn(1)[1] 53 | η[i][rand([j for j = 1:nz], 1)[1]] = w 54 | end 55 | 56 | # for i = 1:nz 57 | # η[i][i] = ϵ 58 | # η[i+nz][i] = -ϵ 59 | # end 60 | 61 | info = MInfo( 62 | collect(1:model.nq), 63 | collect(model.nq .+ (1:model.nq)), 64 | collect(2 * model.nq .+ (1:model.nu)), 65 | zeros(model.nq)) 66 | 67 | fz = zeros(model.nq) 68 | fη = [zeros(model.nq) for i = 1:N] 69 | 70 | ls = LeastSquares(N, 71 | fz, fη, η, 72 | c_func_ls, cθ_func_ls, cθθ_func_ls, 73 | [0.0], zeros(nθ), zeros(nθ, nθ), 74 | [0.0], zeros(nθ), zeros(nθ, nθ), 75 | zeros(nθ), zeros(nθ), 76 | lu_solver(zeros(nθ, nθ))) 77 | 78 | GradientBundle(ls, 79 | zeros(nq), 80 | zeros(nq), 81 | zeros(nu), 82 | zeros(ny, nz), 83 | ny, nz, 84 | info) 85 | end 86 | 87 | function gradient!(sim::Simulator, gb::GradientBundle, q1, q2, u1; verbose=false) 88 | gb.ls.fz .= _step(sim, gb.info, q1, q2, u1) 89 | for i = 1:gb.ls.N 90 | ηq1 = @views gb.ls.η[i][gb.info.idx_q1] 91 | ηq2 = @views gb.ls.η[i][gb.info.idx_q2] 92 | ηu1 = @views gb.ls.η[i][gb.info.idx_u1] 93 | gb.q1η .= q1 94 | gb.q1η .+= ηq1 95 | gb.q2η .= q2 96 | gb.q2η .+= ηq2 97 | gb.u1η .= u1 98 | gb.u1η .+= ηu1 99 | gb.ls.fη[i] .= _step(sim, gb.info, gb.q1η, gb.q2η, gb.u1η) 100 | end 101 | update!(gb.ls, verbose=verbose) 102 | gb.dz .= Base.ReshapedArray(gb.ls.θ, (gb.ny, gb.nz), ()) 103 | return gb.dz 104 | end 105 | 106 | # gradient!(im_dyn.eval_sim, gb, q1, q2, u1) 107 | # @benchmark gradient!($im_dyn.eval_sim, $gb, $q1, $q2, $u1) 108 | 109 | function fx_gb(dx, model::ImplicitDynamics, x, u, w) 110 | q1 = @views x[model.idx_q1] 111 | q2 = @views x[model.idx_q2] 112 | model.v1 .= q2 113 | model.v1 .-= q1 114 | model.v1 ./= model.eval_sim.h 115 | RoboDojo.step!(model.eval_sim, q2, model.v1, u, 1) 116 | nq = model.eval_sim.model.nq 117 | for i = 1:nq 118 | dx[model.idx_q1[i], model.idx_q2[i]] = 1.0 119 | end 120 | gradient!(model.eval_sim, model.info, q1, q2, u) 121 | ∂q3∂q1 = @views model.info.dz[model.idx_q1, model.idx_q1] 122 | ∂q3∂q2 = @views model.info.dz[model.idx_q1, model.idx_q2] 123 | dx[model.idx_q2, model.idx_q1] = ∂q3∂q1 124 | dx[model.idx_q2, model.idx_q2] = ∂q3∂q2 125 | return dx 126 | end 127 | 128 | # dx = zeros(2 * planarpush.nq, 2 * planarpush.nq) 129 | # x = rand(2 * planarpush.nq) 130 | # u = rand(planarpush.nu) 131 | # w = zeros(0) 132 | # fx_gb(dx, im_dyn, x, u, w) 133 | # @benchmark fx_gb($dx, $im_dyn, $x, $u, $w) 134 | # @code_warntype fx_gb(dx, im_dyn, x, u, w) 135 | 136 | function fu_gb(du, model::ImplicitDynamics, x, u, w) 137 | q1 = @views x[model.idx_q1] 138 | q2 = @views x[model.idx_q2] 139 | model.v1 .= q2 140 | model.v1 .-= q1 141 | model.v1 ./= model.eval_sim.h 142 | RoboDojo.step!(model.eval_sim, q2, model.v1, u, 1) 143 | gradient!(model.eval_sim, model.info, q1, q2, u) 144 | ∂q3∂u1 = @views model.info.dz[model.idx_q1, model.info.info.idx_u1] 145 | du[model.idx_q2, :] = ∂q3∂u1 146 | return du 147 | end 148 | 149 | # du = zeros(2 * planarpush.nq, planarpush.nu) 150 | # x = rand(2 * planarpush.nq) 151 | # u = rand(planarpush.nu) 152 | # w = zeros(0) 153 | # fu_gb(du, im_dyn, x, u, w) 154 | # @benchmark fu_gb($du, $im_dyn, $x, $u, $w) 155 | # @code_warntype fu_gb(du, im_dyn, x, u, w) 156 | 157 | 158 | 159 | 160 | 161 | 162 | 163 | 164 | 165 | 166 | 167 | -------------------------------------------------------------------------------- /src/ls.jl: -------------------------------------------------------------------------------- 1 | struct LeastSquares{T} 2 | N::Int 3 | fz::Vector{T} 4 | fη::Vector{Vector{T}} 5 | η::Vector{Vector{T}} 6 | cost::Function 7 | grad::Function 8 | hess::Function 9 | val_cost::Vector{T} 10 | val_grad::Vector{T} 11 | val_hess::Matrix{T} 12 | cache_cost::Vector{T} 13 | cache_grad::Vector{T} 14 | cache_hess::Matrix{T} 15 | θ::Vector{T} 16 | Δ::Vector{T} 17 | solver::RoboDojo.LinearSolver 18 | end 19 | 20 | function eval_cost!(ls::LeastSquares) 21 | fill!(ls.val_cost, 0.0) 22 | for i = 1:ls.N 23 | ls.cost(ls.cache_cost, ls.fz, ls.fη[i], ls.η[i], ls.θ) 24 | ls.val_cost .+= ls.cache_cost 25 | end 26 | end 27 | 28 | function eval_grad!(ls::LeastSquares) 29 | fill!(ls.val_grad, 0.0) 30 | for i = 1:ls.N 31 | ls.grad(ls.cache_grad, ls.fz, ls.fη[i], ls.η[i], ls.θ) 32 | ls.val_grad .+= ls.cache_grad 33 | end 34 | end 35 | 36 | function eval_hess!(ls::LeastSquares) 37 | fill!(ls.val_hess, 0.0) 38 | for i = 1:ls.N 39 | ls.hess(ls.cache_hess, ls.fz, ls.fη[i], ls.η[i], ls.θ) 40 | ls.val_hess .+= ls.cache_hess 41 | end 42 | end 43 | 44 | function update!(ls::LeastSquares; tol=1.0e-8, verbose=true) 45 | iter = 0 46 | eval_grad!(ls) 47 | res = norm(ls.val_grad, Inf) 48 | verbose && (@show res) 49 | while res > tol && iter < 100 50 | eval_hess!(ls) 51 | fill!(ls.Δ, 0.0) 52 | RoboDojo.linear_solve!(ls.solver, ls.Δ, ls.val_hess, ls.val_grad) 53 | ls.θ .-= ls.Δ 54 | eval_grad!(ls) 55 | res = norm(ls.val_grad, Inf) 56 | verbose && (@show res) 57 | verbose && (@show ls.θ) 58 | iter += 1 59 | end 60 | end 61 | 62 | # ## test 63 | # function f(z) 64 | # x = z[1:2] 65 | # u = z[2 .+ (1:1)] 66 | # A = [1.0 1.0; 0.0 1.0] 67 | # B = [0.0; 1.0] 68 | 69 | # A * x + B * u[1] 70 | # end 71 | 72 | # nx = 2 73 | # nu = 1 74 | # nz = nx + nu 75 | # nθ = nx * (nz) 76 | 77 | # @variables fz[1:nx] fη[1:nx] η[1:nz] θ[1:nθ] 78 | 79 | # function cost(fz, fη, η, θ) 80 | # M = reshape(θ, nx, nz) 81 | # r = fη - fz - M * η 82 | # return [transpose(r) * r] 83 | # return [transpose(θ) * θ] 84 | # end 85 | 86 | # c = cost(fz, fη, η, θ) 87 | # cθ = Symbolics.gradient(c[1], θ) 88 | # cθθ = Symbolics.hessian(c[1], θ) 89 | 90 | # c_func = eval(Symbolics.build_function(c, fz, fη, η, θ)[2]) 91 | # cθ_func = eval(Symbolics.build_function(cθ, fz, fη, η, θ)[2]) 92 | # cθθ_func = eval(Symbolics.build_function(cθθ, fz, fη, η, θ)[2]) 93 | 94 | # fz0 = rand(nx) 95 | # fη0 = rand(nx) 96 | # η0 = rand(nz) 97 | # θ0 = rand(nθ) 98 | # θθ0 = rand(nθ, nθ) 99 | 100 | # c0 = zeros(1) 101 | # cθ0 = zeros(nθ) 102 | # cθθ0 = zeros(nθ, nθ) 103 | 104 | # c_func(c0, fz0, fη0, η0, θ0) 105 | # @benchmark c_func($c0, $fz0, $fη0, $η0, $θ0) 106 | 107 | # cθ_func(cθ0, fz0, fη0, η0, θ0) 108 | # @benchmark cθ_func($cθ0, $fz0, $fη0, $η0, $θ0) 109 | 110 | # cθθ_func(cθθ0, fz0, fη0, η0, θ0) 111 | # @benchmark cθ_func($cθθ0, $fz0, $fη0, $η0, $θθ0) 112 | 113 | # N = 2 * nz 114 | # η = [zeros(nz) for i = 1:N] 115 | # ϵ = 0.1 116 | # for i = 1:nz 117 | # η[i][i] = ϵ 118 | # η[i+nz][i] = -ϵ 119 | # end 120 | 121 | # z0 = rand(nz) 122 | # fz = f(z0) 123 | # fη = [f(z0 + η[i]) for i = 1:N] 124 | 125 | # ls = LeastSquares(N, 126 | # fz, fη, η, 127 | # c_func, cθ_func, cθθ_func, 128 | # [0.0], zeros(nθ), zeros(nθ, nθ), 129 | # [0.0], zeros(nθ), zeros(nθ, nθ), 130 | # zeros(nθ), zeros(nθ), 131 | # lu_solver(zeros(nθ, nθ))) 132 | 133 | # # eval_cost!(ls, z0) 134 | # # @benchmark eval_cost!($ls) 135 | 136 | # # eval_grad!(ls, z0) 137 | # # @benchmark eval_grad!($ls) 138 | 139 | # # eval_hess!(ls, z0) 140 | # # @benchmark eval_hess!($ls) 141 | 142 | # update!(ls) 143 | # reshape(ls.θ, nx, nz) 144 | # # @benchmark update!($ls, z) setup=(z=copy(z0)) 145 | 146 | 147 | 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | -------------------------------------------------------------------------------- /src/models/acrobot/codegen.jl: -------------------------------------------------------------------------------- 1 | path = @get_scratch!("acrobot") 2 | 3 | # generate residual methods 4 | nq = acrobot_impact.nq 5 | nu = acrobot_impact.nu 6 | nc = acrobot_impact.nc 7 | nz = nq + nc + nc 8 | nz_nominal = nq 9 | nθ = nq + nq + nu + 1 10 | 11 | # Declare variables 12 | @variables z[1:nz] 13 | @variables z_nominal[1:nz_nominal] 14 | @variables θ[1:nθ] 15 | @variables κ[1:1] 16 | 17 | # Residual 18 | r_impact = residual(acrobot_impact, z, θ, κ) 19 | r_impact = Symbolics.simplify.(r_impact) 20 | rz_impact = Symbolics.jacobian(r_impact, z, simplify=true) 21 | rθ_impact = Symbolics.jacobian(r_impact, θ, simplify=true) 22 | 23 | # Build function 24 | r_acrobot_impact_func = build_function(r_impact, z, θ, κ)[2] 25 | rz_acrobot_impact_func = build_function(rz_impact, z, θ)[2] 26 | rθ_acrobot_impact_func = build_function(rθ_impact, z, θ)[2] 27 | 28 | rz_acrobot_impact_array = similar(rz_impact, Float64) 29 | rθ_acrobot_impact_array = similar(rθ_impact, Float64) 30 | 31 | @save joinpath(path, "impact.jld2") r_acrobot_impact_func rz_acrobot_impact_func rθ_acrobot_impact_func rz_acrobot_impact_array rθ_acrobot_impact_array 32 | @load joinpath(path, "impact.jld2") r_acrobot_impact_func rz_acrobot_impact_func rθ_acrobot_impact_func rz_acrobot_impact_array rθ_acrobot_impact_array 33 | 34 | # Residual 35 | r_nominal = residual(acrobot_nominal, z_nominal, θ, κ) 36 | r_nominal = Symbolics.simplify.(r_nominal) 37 | rz_nominal = Symbolics.jacobian(r_nominal, z_nominal, simplify=true) 38 | rθ_nominal = Symbolics.jacobian(r_nominal, θ, simplify=true) 39 | 40 | # Build function 41 | r_acrobot_nominal_func = build_function(r_nominal, z_nominal, θ, κ)[2] 42 | rz_acrobot_nominal_func = build_function(rz_nominal, z_nominal, θ)[2] 43 | rθ_acrobot_nominal_func = build_function(rθ_nominal, z_nominal, θ)[2] 44 | 45 | rz_acrobot_nominal_array = similar(rz_nominal, Float64) 46 | rθ_acrobot_nominal_array = similar(rθ_nominal, Float64) 47 | 48 | path = @get_scratch!("acrobot") 49 | @save joinpath(path, "nominal.jld2") r_acrobot_nominal_func rz_acrobot_nominal_func rθ_acrobot_nominal_func rz_acrobot_nominal_array rθ_acrobot_nominal_array 50 | @load joinpath(path, "nominal.jld2") r_acrobot_nominal_func rz_acrobot_nominal_func rθ_acrobot_nominal_func rz_acrobot_nominal_array rθ_acrobot_nominal_array 51 | -------------------------------------------------------------------------------- /src/models/acrobot/model.jl: -------------------------------------------------------------------------------- 1 | """ 2 | Double pendulum 3 | """ 4 | struct Impact end 5 | struct Nominal end 6 | 7 | struct DoublePendulum{T,X} <: Model{T} 8 | nq::Int 9 | nu::Int 10 | nw::Int 11 | nc::Int 12 | 13 | m1::T # mass link 1 14 | J1::T # inertia link 1 15 | l1::T # length link 1 16 | lc1::T # length to COM link 1 17 | 18 | m2::T # mass link 2 19 | J2::T # inertia link 2 20 | l2::T # length link 2 21 | lc2::T # length to COM link 2 22 | 23 | g::T # gravity 24 | 25 | b1::T # joint friction 26 | b2::T 27 | end 28 | 29 | lagrangian(model::DoublePendulum, q, q̇) = 0.0 30 | 31 | function kinematics(model::DoublePendulum, x) 32 | [model.l1 * sin(x[1]) + model.l2 * sin(x[1] + x[2]), 33 | -1.0 * model.l1 * cos(x[1]) - model.l2 * cos(x[1] + x[2])] 34 | end 35 | 36 | function kinematics_elbow(model::DoublePendulum, x) 37 | [model.l1 * sin(x[1]), 38 | -1.0 * model.l1 * cos(x[1])] 39 | end 40 | 41 | function M_func(model::DoublePendulum, x) 42 | a = (model.J1 + model.J2 + model.m2 * model.l1 * model.l1 43 | + 2.0 * model.m2 * model.l1 * model.lc2 * cos(x[2])) 44 | 45 | b = model.J2 + model.m2 * model.l1 * model.lc2 * cos(x[2]) 46 | 47 | c = model.J2 48 | 49 | [a b; 50 | b c] 51 | end 52 | 53 | function τ_func(model::DoublePendulum, x) 54 | a = (-1.0 * model.m1 * model.g * model.lc1 * sin(x[1]) 55 | - model.m2 * model.g * (model.l1 * sin(x[1]) 56 | + model.lc2 * sin(x[1] + x[2]))) 57 | 58 | b = -1.0 * model.m2 * model.g * model.lc2 * sin(x[1] + x[2]) 59 | 60 | [a, b] 61 | end 62 | 63 | function c_func(model::DoublePendulum, q, q̇) 64 | a = -2.0 * model.m2 * model.l1 * model.lc2 * sin(q[2]) * q̇[2] 65 | b = -1.0 * model.m2 * model.l1 * model.lc2 * sin(q[2]) * q̇[2] 66 | c = model.m2 * model.l1 * model.lc2 * sin(q[2]) * q̇[1] 67 | d = 0.0 68 | 69 | [a b; 70 | c d] 71 | end 72 | 73 | function B_func(model::DoublePendulum, x) 74 | [0.0; 1.0] 75 | end 76 | 77 | function C_func(model::DoublePendulum, q, q̇) 78 | c_func(model, q, q̇) * q̇ - τ_func(model, q) 79 | end 80 | 81 | function ϕ_func(model, q) 82 | [0.5 * π - q[2], q[2] + 0.5 * π] 83 | end 84 | 85 | function P_func(model, q) 86 | ϕ = ϕ_func(model, q) 87 | Symbolics.jacobian(ϕ, q) 88 | end 89 | 90 | function dynamics(model::DoublePendulum{T,Impact}, mass_matrix, dynamics_bias, h, q0, q1, u1, w1, λ1, q2) where T 91 | # evalutate at midpoint 92 | qm1 = 0.5 * (q0 + q1) 93 | vm1 = (q1 - q0) / h[1] 94 | qm2 = 0.5 * (q1 + q2) 95 | vm2 = (q2 - q1) / h[1] 96 | 97 | D1L1, D2L1 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm1, vm1) 98 | D1L2, D2L2 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm2, vm2) 99 | 100 | return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2 101 | + B_func(model, qm2) * u1[1] 102 | + transpose(P_func(model, q2)) * λ1 103 | - h[1] * 0.5 .* vm2) # damping 104 | end 105 | 106 | function dynamics(model::DoublePendulum{T,Nominal}, mass_matrix, dynamics_bias, h, q0, q1, u1, w1, λ1, q2) where T 107 | # evalutate at midpoint 108 | qm1 = 0.5 * (q0 + q1) 109 | vm1 = (q1 - q0) / h[1] 110 | qm2 = 0.5 * (q1 + q2) 111 | vm2 = (q2 - q1) / h[1] 112 | 113 | D1L1, D2L1 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm1, vm1) 114 | D1L2, D2L2 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm2, vm2) 115 | 116 | return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2 117 | + B_func(model, qm2) * u1[1] 118 | - h[1] * 0.5 .* vm2) # damping 119 | end 120 | 121 | function residual(model::DoublePendulum{T,Impact}, z, θ, κ) where T 122 | nq = model.nq 123 | nu = model.nu 124 | nc = model.nc 125 | 126 | q0 = θ[1:nq] 127 | q1 = θ[nq .+ (1:nq)] 128 | u1 = θ[2nq .+ (1:nu)] 129 | h = θ[2nq + nu .+ (1:1)] 130 | 131 | q2 = z[1:nq] 132 | λ1 = z[nq .+ (1:nc)] 133 | s1 = z[nq + nc .+ (1:nc)] 134 | 135 | [ 136 | dynamics(model, a -> M_func(model, a), (a, b) -> C_func(model, a, b), 137 | h, q0, q1, u1, zeros(model.nw), λ1, q2); 138 | s1 .- ϕ_func(model, q2); 139 | λ1 .* s1 .- κ; 140 | ] 141 | 142 | end 143 | 144 | function residual(model::DoublePendulum{T,Nominal}, z, θ, κ) where T 145 | nq = model.nq 146 | nu = model.nu 147 | nc = model.nc 148 | 149 | q0 = θ[1:nq] 150 | q1 = θ[nq .+ (1:nq)] 151 | u1 = θ[2nq .+ (1:nu)] 152 | h = θ[2nq + nu .+ (1:1)] 153 | 154 | q2 = z[1:nq] 155 | return dynamics(model, a -> M_func(model, a), (a, b) -> C_func(model, a, b), 156 | h, q0, q1, u1, zeros(model.nw), zeros(model.nc), q2) 157 | end 158 | 159 | acrobot_impact = DoublePendulum{Float64,Impact}(2, 1, 0, 2, 160 | 1.0, 0.333, 1.0, 0.5, 1.0, 0.333, 1.0, 0.5, 9.81, 0.0, 0.0) 161 | 162 | acrobot_nominal = DoublePendulum{Float64,Nominal}(2, 1, 0, 0, 163 | 1.0, 0.333, 1.0, 0.5, 1.0, 0.333, 1.0, 0.5, 9.81, 0.0, 0.0) 164 | 165 | -------------------------------------------------------------------------------- /src/models/acrobot/simulator_impact.jl: -------------------------------------------------------------------------------- 1 | function RoboDojo.indices_z(model::DoublePendulum{T,Impact}) where T 2 | nq = model.nq 3 | nc = model.nc 4 | q = collect(1:nq) 5 | γ = collect(nq .+ (1:2)) 6 | sγ = collect(nq + 2 .+ (1:2)) 7 | ψ = collect(1:0) 8 | b = collect(1:0) 9 | sψ = collect(1:0) 10 | sb = collect(1:0) 11 | IndicesZ(q, γ, sγ, ψ, b, sψ, sb) 12 | end 13 | 14 | RoboDojo.num_var(model::DoublePendulum{T,Impact}) where T = model.nq + 2 * model.nc 15 | 16 | function RoboDojo.indices_optimization(model::DoublePendulum{T,Impact}) where T 17 | nq = model.nq 18 | nc = model.nc 19 | nz = nq + 2 * nc 20 | RoboDojo.IndicesOptimization( 21 | nz, 22 | nz, 23 | [collect(nq .+ (1:2)), collect(nq + 2 .+ (1:2))], 24 | [collect(nq .+ (1:2)), collect(nq + 2 .+ (1:2))], 25 | Vector{Vector{Vector{Int}}}(), 26 | Vector{Vector{Vector{Int}}}(), 27 | collect(1:(nq + nc)), 28 | collect(nq + nc .+ (1:nc)), 29 | collect(1:0), 30 | Vector{Vector{Int}}(), 31 | collect(nq + nc .+ (1:nc))) 32 | end 33 | 34 | function RoboDojo.initialize_z!(z, model::DoublePendulum{T,Impact}, idx::RoboDojo.IndicesZ, q) where T 35 | z[idx.q] .= q 36 | z[idx.γ] .= 1.0 37 | z[idx.sγ] .= 1.0 38 | end 39 | 40 | RoboDojo.nominal_configuration(model::DoublePendulum) = zeros(model.nq) 41 | friction_coefficients(model::DoublePendulum{T}) where T = T[] 42 | -------------------------------------------------------------------------------- /src/models/acrobot/simulator_nominal.jl: -------------------------------------------------------------------------------- 1 | function RoboDojo.indices_z(model::DoublePendulum{T,Nominal}) where T 2 | nq = model.nq 3 | nc = model.nc 4 | q = collect(1:nq) 5 | γ = collect(1:0) 6 | sγ = collect(1:0) 7 | ψ = collect(1:0) 8 | b = collect(1:0) 9 | sψ = collect(1:0) 10 | sb = collect(1:0) 11 | IndicesZ(q, γ, sγ, ψ, b, sψ, sb) 12 | end 13 | 14 | RoboDojo.num_var(model::DoublePendulum{T,Nominal}) where T = model.nq 15 | 16 | function RoboDojo.indices_optimization(model::DoublePendulum{T,Nominal}) where T 17 | nq = model.nq 18 | nz = nq 19 | IndicesOptimization( 20 | nz, 21 | nz, 22 | [collect(1:0), collect(1:0)], 23 | [collect(1:0), collect(1:0)], 24 | Vector{Vector{Vector{Int}}}(), 25 | Vector{Vector{Vector{Int}}}(), 26 | collect(1:nq), 27 | collect(1:0), 28 | collect(1:0), 29 | Vector{Vector{Int}}(), 30 | collect(1:0)) 31 | end 32 | 33 | function RoboDojo.initialize_z!(z, model::DoublePendulum{T,Nominal}, idx::IndicesZ, q) where T 34 | z[idx.q] .= q 35 | end 36 | -------------------------------------------------------------------------------- /src/models/acrobot/visuals.jl: -------------------------------------------------------------------------------- 1 | # visualization 2 | function visualize!(vis, model::DoublePendulum, x; 3 | color=Colors.RGBA(0.0, 0.0, 0.0, 1.0), 4 | r = 0.1, Δt = 0.1) 5 | default_background!(vis) 6 | 7 | i = 1 8 | l1 = Cylinder(Point3f0(0.0, 0.0, 0.0), Point3f0(0.0, 0.0, model.l1), 9 | convert(Float32, 0.025)) 10 | setobject!(vis["l1$i"], l1, MeshPhongMaterial(color = color)) 11 | l2 = Cylinder(Point3f0(0.0,0.0,0.0), Point3f0(0.0, 0.0, model.l2), 12 | convert(Float32, 0.025)) 13 | setobject!(vis["l2$i"], l2, MeshPhongMaterial(color = color)) 14 | 15 | setobject!(vis["elbow$i"], Sphere(Point3f0(0.0), 16 | convert(Float32, 0.05)), 17 | MeshPhongMaterial(color = color)) 18 | setobject!(vis["ee$i"], Sphere(Point3f0(0.0), 19 | convert(Float32, 0.05)), 20 | MeshPhongMaterial(color = color)) 21 | 22 | anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt))) 23 | 24 | T = length(x) 25 | for t = 1:T 26 | 27 | MeshCat.atframe(anim,t) do 28 | p_mid = [kinematics_elbow(model, x[t])[1], 0.0, kinematics_elbow(model, x[t])[2]] 29 | p_ee = [kinematics(model, x[t])[1], 0.0, kinematics(model, x[t])[2]] 30 | 31 | settransform!(vis["l1$i"], cable_transform(zeros(3), p_mid)) 32 | settransform!(vis["l2$i"], cable_transform(p_mid, p_ee)) 33 | 34 | settransform!(vis["elbow$i"], Translation(p_mid)) 35 | settransform!(vis["ee$i"], Translation(p_ee)) 36 | end 37 | end 38 | 39 | settransform!(vis["/Cameras/default"], 40 | compose(Translation(0.0, 0.0, -1.0), LinearMap(RotZ(- pi / 2)))) 41 | setvisible!(vis["/Grid"], true) 42 | 43 | MeshCat.setanimation!(vis, anim) 44 | end 45 | 46 | # visualization 47 | function _create_acrobot!(vis, model::DoublePendulum; 48 | tl = 1.0, 49 | limit_color = Colors.RGBA(0.0, 1.0, 0.0, tl), 50 | color = Colors.RGBA(0.0, 0.0, 0.0, tl), 51 | i = 0, 52 | r = 0.1) 53 | 54 | l1 = Cylinder(Point3f0(0.0, 0.0, 0.0), Point3f0(0.0, 0.0, model.l1), 55 | convert(Float32, 0.025)) 56 | setobject!(vis["l1_$i"], l1, MeshPhongMaterial(color = color)) 57 | l2 = Cylinder(Point3f0(0.0,0.0,0.0), Point3f0(0.0, 0.0, model.l2), 58 | convert(Float32, 0.025)) 59 | setobject!(vis["l2_$i"], l2, MeshPhongMaterial(color = color)) 60 | 61 | setobject!(vis["elbow_nominal_$i"], Sphere(Point3f0(0.0), 62 | convert(Float32, 0.05)), 63 | MeshPhongMaterial(color = Colors.RGBA(0.0, 0.0, 0.0, tl))) 64 | setobject!(vis["elbow_limit_$i"], Sphere(Point3f0(0.0), 65 | convert(Float32, 0.05)), 66 | MeshPhongMaterial(color = limit_color)) 67 | setobject!(vis["ee_$i"], Sphere(Point3f0(0.0), 68 | convert(Float32, 0.05)), 69 | MeshPhongMaterial(color = Colors.RGBA(0.0, 0.0, 0.0, tl))) 70 | end 71 | 72 | function _set_acrobot!(vis, model::DoublePendulum, x; 73 | i = 0, ϵ = 1.0e-1) 74 | 75 | p_mid = [kinematics_elbow(model, x)[1], 0.0, kinematics_elbow(model, x)[2]] 76 | p_ee = [kinematics(model, x)[1], 0.0, kinematics(model, x)[2]] 77 | 78 | settransform!(vis["l1_$i"], cable_transform(zeros(3), p_mid)) 79 | settransform!(vis["l2_$i"], cable_transform(p_mid, p_ee)) 80 | 81 | settransform!(vis["elbow_nominal_$i"], Translation(p_mid)) 82 | settransform!(vis["elbow_limit_$i"], Translation(p_mid)) 83 | settransform!(vis["ee_$i"], Translation(p_ee)) 84 | 85 | if x[2] <= -0.5 * π + ϵ || x[2] >= 0.5 * π - ϵ 86 | setvisible!(vis["elbow_nominal_$i"], false) 87 | setvisible!(vis["elbow_limit_$i"], true) 88 | else 89 | setvisible!(vis["elbow_nominal_$i"], true) 90 | setvisible!(vis["elbow_limit_$i"], false) 91 | end 92 | end 93 | 94 | # visualization 95 | function visualize_elbow!(vis, model::DoublePendulum, x; 96 | tl = 1.0, 97 | i = 0, 98 | color = Colors.RGBA(0.0, 0.0, 0.0, 1.0), 99 | limit_color = Colors.RGBA(0.0, 1.0, 0.0, tl), 100 | r = 0.1, Δt = 0.1, 101 | ϵ = 1.0e-1) 102 | 103 | _create_acrobot!(vis, model, 104 | tl = tl, 105 | color = color, 106 | limit_color = limit_color, 107 | i = i, 108 | r = r) 109 | 110 | anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt))) 111 | 112 | T = length(x) 113 | for t = 1:T 114 | MeshCat.atframe(anim,t) do 115 | _set_acrobot!(vis, model, x[t], i = i, ϵ = ϵ) 116 | end 117 | end 118 | 119 | # settransform!(vis["/Cameras/default"], 120 | # compose(Translation(0.0 , 0.0 , 0.0), LinearMap(RotZ(pi / 2.0)))) 121 | 122 | MeshCat.setanimation!(vis, anim) 123 | end 124 | -------------------------------------------------------------------------------- /src/models/cartpole/codegen.jl: -------------------------------------------------------------------------------- 1 | path = @get_scratch!("cartpole") 2 | 3 | nq = cartpole_friction.nq 4 | nu = cartpole_friction.nu 5 | nc = cartpole_friction.nc 6 | nz = nq + 4 * nc 7 | nθ = nq + nq + nu + 2 + 1 8 | nz_nf = nq 9 | nθ_nf = nq + nq + nu + 1 10 | 11 | # Declare variables 12 | @variables z[1:nz] 13 | @variables θ[1:nθ] 14 | @variables κ[1:1] 15 | @variables z_nf[1:nz_nf] 16 | @variables θ_nf[1:nθ_nf] 17 | @variables κ_nf[1:1] 18 | 19 | # Residual 20 | r_friction= residual(cartpole_friction, z, θ, κ) 21 | r_friction = Symbolics.simplify.(r_friction) 22 | rz_friction = Symbolics.jacobian(r_friction, z, simplify=true) 23 | rθ_friction = Symbolics.jacobian(r_friction, θ, simplify=true) 24 | 25 | # Build function 26 | r_cartpole_friction_func = build_function(r_friction, z, θ, κ)[2] 27 | rz_cartpole_friction_func = build_function(rz_friction, z, θ)[2] 28 | rθ_cartpole_friction_func = build_function(rθ_friction, z, θ)[2] 29 | 30 | rz_cartpole_friction_array = similar(rz_friction, Float64) 31 | rθ_cartpole_friction_array = similar(rθ_friction, Float64) 32 | 33 | @save joinpath(path, "friction.jld2") r_cartpole_friction_func rz_cartpole_friction_func rθ_cartpole_friction_func rz_cartpole_friction_array rθ_cartpole_friction_array 34 | @load joinpath(path, "friction.jld2") r_cartpole_friction_func rz_cartpole_friction_func rθ_cartpole_friction_func rz_cartpole_friction_array rθ_cartpole_friction_array 35 | 36 | # Residual 37 | r_frictionless = residual(cartpole_frictionless, z_nf, θ_nf, κ_nf) 38 | r_frictionless = Symbolics.simplify.(r_frictionless) 39 | rz_frictionless = Symbolics.jacobian(r_frictionless, z_nf, simplify=true) 40 | rθ_frictionless = Symbolics.jacobian(r_frictionless, θ_nf, simplify=true) 41 | 42 | # Build function 43 | r_cartpole_frictionless_func = build_function(r_frictionless, z_nf, θ_nf, κ_nf)[2] 44 | rz_cartpole_frictionless_func = build_function(rz_frictionless, z_nf, θ_nf)[2] 45 | rθ_cartpole_frictionless_func = build_function(rθ_frictionless, z_nf, θ_nf)[2] 46 | 47 | rz_cartpole_frictionless_array = similar(rz_frictionless, Float64) 48 | rθ_cartpole_frictionless_array = similar(rθ_frictionless, Float64) 49 | 50 | @save joinpath(path, "frictionless.jld2") r_cartpole_frictionless_func rz_cartpole_frictionless_func rθ_cartpole_frictionless_func rz_cartpole_frictionless_array rθ_cartpole_frictionless_array 51 | @load joinpath(path, "frictionless.jld2") r_cartpole_frictionless_func rz_cartpole_frictionless_func rθ_cartpole_frictionless_func rz_cartpole_frictionless_array rθ_cartpole_frictionless_array -------------------------------------------------------------------------------- /src/models/cartpole/model.jl: -------------------------------------------------------------------------------- 1 | """ 2 | Cartpole w/ internal friction 3 | """ 4 | struct Friction end 5 | struct Frictionless end 6 | 7 | struct Cartpole{T,X} <: Model{T} 8 | # dimensions 9 | nq::Int # generalized coordinates 10 | nu::Int # controls 11 | nw::Int # parameters 12 | nc::Int # contact points 13 | 14 | mc::T # mass of the cart in kg 15 | mp::T # mass of the pole (point mass at the end) in kg 16 | l::T # length of the pole in m 17 | g::T # gravity m/s^2 18 | 19 | friction::Vector{T} # friction coefficients for slider and arm joints 20 | end 21 | 22 | function kinematics(model::Cartpole, q) 23 | [q[1] + model.l * sin(q[2]); -model.l * cos(q[2])] 24 | end 25 | 26 | lagrangian(model::Cartpole, q, q̇) = 0.0 27 | 28 | function M_func(model::Cartpole, x) 29 | H = [model.mc + model.mp model.mp * model.l * cos(x[2]); 30 | model.mp * model.l * cos(x[2]) model.mp * model.l^2.0] 31 | return H 32 | end 33 | 34 | function B_func(model::Cartpole, x) 35 | [1.0; 0.0] 36 | end 37 | 38 | function P_func(model::Cartpole, x) 39 | [1.0 0.0; 40 | 0.0 1.0] 41 | end 42 | 43 | function C_func(model::Cartpole, q, q̇) 44 | C = [0.0 -1.0 * model.mp * q̇[2] * model.l * sin(q[2]); 45 | 0.0 0.0] 46 | G = [0.0, 47 | model.mp * model.g * model.l * sin(q[2])] 48 | return -C * q̇ + G 49 | end 50 | 51 | function dynamics(model::Cartpole{T,Friction}, mass_matrix, dynamics_bias, h, q0, q1, u1, w1, λ1, q2) where T 52 | # evalutate at midpoint 53 | qm1 = 0.5 * (q0 + q1) 54 | vm1 = (q1 - q0) / h[1] 55 | qm2 = 0.5 * (q1 + q2) 56 | vm2 = (q2 - q1) / h[1] 57 | 58 | D1L1, D2L1 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm1, vm1) 59 | D1L2, D2L2 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm2, vm2) 60 | 61 | d = 0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2 # variational integrator (midpoint) 62 | 63 | return d + B_func(model, qm2) * u1[1] + transpose(P_func(model, q2)) * λ1 64 | end 65 | 66 | function dynamics(model::Cartpole{T,Frictionless}, mass_matrix, dynamics_bias, h, q0, q1, u1, w1, λ1, q2) where T 67 | # evalutate at midpoint 68 | qm1 = 0.5 * (q0 + q1) 69 | vm1 = (q1 - q0) / h[1] 70 | qm2 = 0.5 * (q1 + q2) 71 | vm2 = (q2 - q1) / h[1] 72 | 73 | D1L1, D2L1 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm1, vm1) 74 | D1L2, D2L2 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm2, vm2) 75 | 76 | d = 0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2 # variational integrator (midpoint) 77 | 78 | return d + B_func(model, qm2) * u1[1] 79 | end 80 | 81 | function residual(model::Cartpole{T,Friction}, z, θ, κ) where T 82 | nq = model.nq 83 | nu = model.nu 84 | nc = model.nc 85 | 86 | q0 = θ[1:nq] 87 | q1 = θ[nq .+ (1:nq)] 88 | u1 = θ[2nq .+ (1:nu)] 89 | μ_slider = θ[2nq + nu .+ (1:1)] 90 | μ_angle = θ[2nq + nu + 1 .+ (1:1)] 91 | h = θ[2nq + nu + 1 + 1 .+ (1:1)] 92 | 93 | q2 = z[1:nq] 94 | ψ = z[nq .+ (1:nc)] 95 | b = z[nq + nc .+ (1:nc)] 96 | sψ = z[nq + nc + nc .+ (1:nc)] 97 | sb = z[nq + nc + nc + nc .+ (1:nc)] 98 | 99 | vT1 = (q2[1] - q1[1]) / h 100 | vT2 = (q2[2] - q1[2]) / h 101 | 102 | λ1 = b 103 | 104 | [ 105 | dynamics(model, a -> M_func(model, a), (a, b) -> C_func(model, a, b), 106 | h, q0, q1, u1, zeros(model.nw), λ1, q2); 107 | sb[1] - vT1[1]; 108 | ψ[1] .- μ_slider[1] * (model.mp + model.mc) * model.g * h[1]; 109 | sb[2] - vT2[1]; 110 | ψ[2] .- μ_angle[1] * (model.mp * model.g * model.l) * h[1]; 111 | cone_product([ψ[1]; b[1]], [sψ[1]; sb[1]]) - [κ[1]; 0.0]; 112 | cone_product([ψ[2]; b[2]], [sψ[2]; sb[2]]) - [κ[1]; 0.0]; 113 | ] 114 | end 115 | 116 | function residual(model::Cartpole{T,Frictionless}, z, θ, κ) where T 117 | nq = model.nq 118 | nu = model.nu 119 | 120 | q0 = θ[1:nq] 121 | q1 = θ[nq .+ (1:nq)] 122 | u1 = θ[2nq .+ (1:nu)] 123 | h = θ[2nq + nu .+ (1:1)] 124 | 125 | q2 = z[1:nq] 126 | 127 | return dynamics(model, a -> M_func(model, a), (a, b) -> C_func(model, a, b), 128 | h, q0, q1, u1, zeros(model.nw), zeros(0), q2); 129 | end 130 | 131 | # models 132 | cartpole_friction = Cartpole{Float64,Friction}(2, 1, 0, 2, 1.0, 0.2, 0.5, 9.81, [0.1; 0.1]) 133 | cartpole_frictionless = Cartpole{Float64,Frictionless}(2, 1, 0, 2, 1.0, 0.2, 0.5, 9.81, [0.0; 0.0]) 134 | 135 | -------------------------------------------------------------------------------- /src/models/cartpole/simulator_friction.jl: -------------------------------------------------------------------------------- 1 | friction_coefficients(model::Cartpole{T,Friction}) where T = model.friction 2 | 3 | function RoboDojo.indices_z(model::Cartpole{T,Friction}) where T 4 | nq = model.nq 5 | nc = model.nc 6 | q = collect(1:nq) 7 | γ = collect(nq .+ (1:0)) 8 | sγ = collect(nq .+ (1:0)) 9 | ψ = collect(nq .+ (1:nc)) 10 | b = collect(nq + nc .+ (1:nc)) 11 | sψ = collect(nq + nc + nc .+ (1:nc)) 12 | sb = collect(nq + nc + nc + nc .+ (1:nc)) 13 | IndicesZ(q, γ, sγ, ψ, b, sψ, sb) 14 | end 15 | 16 | RoboDojo.num_var(model::Cartpole{T,Friction}) where T = model.nq + 4 * model.nc 17 | 18 | function RoboDojo.indices_optimization(model::Cartpole{T,Friction}) where T 19 | nq = model.nq 20 | nc = model.nc 21 | nz = nq + 4 * nc 22 | IndicesOptimization( 23 | nz, 24 | nz, 25 | [collect(1:0), collect(1:0)], 26 | [collect(1:0), collect(1:0)], 27 | [[collect([nq + i, nq + nc + i]), collect([nq + nc + nc + i, nq + nc + nc + nc + i])] for i = 1:nc], 28 | [[collect([nq + i, nq + nc + i]), collect([nq + nc + nc + i, nq + nc + nc + nc + i])] for i = 1:nc], 29 | collect(1:(nq + 2 * nc)), 30 | collect(1:0), 31 | collect(nq + 2 * nc .+ (1:(2 * nc))), 32 | [collect(nq + 2 * nc + (i - 1) * 2 .+ (1:2)) for i = 1:nc], 33 | collect(nq + 2 * nc .+ (1:(2 * nc)))) 34 | end 35 | 36 | function RoboDojo.initialize_z!(z, model::Cartpole{T,Friction}, idx::IndicesZ, q) where T 37 | z[idx.q] .= q 38 | z[idx.ψ] .= 1.0 39 | z[idx.b] .= 0.1 40 | z[idx.sψ] .= 1.0 41 | z[idx.sb] .= 0.1 42 | end 43 | 44 | RoboDojo.nominal_configuration(model::Cartpole) = zeros(model.nq) 45 | -------------------------------------------------------------------------------- /src/models/cartpole/simulator_frictionless.jl: -------------------------------------------------------------------------------- 1 | friction_coefficients(model::Cartpole{T,Frictionless}) where T = T[] 2 | 3 | function RoboDojo.indices_z(model::Cartpole{T,Frictionless}) where T 4 | nq = model.nq 5 | nc = model.nc 6 | q = collect(1:nq) 7 | γ = collect(1:0) 8 | sγ = collect(1:0) 9 | ψ = collect(1:0) 10 | b = collect(1:0) 11 | sψ = collect(1:0) 12 | sb = collect(1:0) 13 | IndicesZ(q, γ, sγ, ψ, b, sψ, sb) 14 | end 15 | 16 | RoboDojo.num_var(model::Cartpole{T,Frictionless}) where T = model.nq 17 | 18 | function RoboDojo.indices_optimization(model::Cartpole{T,Frictionless}) where T 19 | nq = model.nq 20 | nz = nq 21 | IndicesOptimization( 22 | nz, 23 | nz, 24 | [collect(1:0), collect(1:0)], 25 | [collect(1:0), collect(1:0)], 26 | Vector{Vector{Vector{Int}}}(), 27 | Vector{Vector{Vector{Int}}}(), 28 | collect(1:(nq)), 29 | collect(1:0), 30 | collect(1:0), 31 | Vector{Vector{Int}}(), 32 | collect(1:0)) 33 | end 34 | 35 | function RoboDojo.initialize_z!(z, model::Cartpole{T,Frictionless}, idx::IndicesZ, q) where T 36 | z[idx.q] .= q 37 | end -------------------------------------------------------------------------------- /src/models/cartpole/visuals.jl: -------------------------------------------------------------------------------- 1 | function _create_cartpole!(vis, model; 2 | i = 0, 3 | tl = 1.0, 4 | color = Colors.RGBA(0, 0, 0, tl)) 5 | 6 | l2 = Cylinder(Point3f0(-model.l * 10.0, 0.0, 0.0), 7 | Point3f0(model.l * 10.0, 0.0, 0.0), 8 | convert(Float32, 0.0125)) 9 | 10 | setobject!(vis["slider_$i"], l2, MeshPhongMaterial(color = Colors.RGBA(0.0, 0.0, 0.0, tl))) 11 | 12 | l1 = Cylinder(Point3f0(0.0, 0.0, 0.0), 13 | Point3f0(0.0, 0.0, model.l), 14 | convert(Float32, 0.025)) 15 | 16 | setobject!(vis["arm_$i"], l1, 17 | MeshPhongMaterial(color = Colors.RGBA(0.0, 0.0, 0.0, tl))) 18 | 19 | setobject!(vis["base_$i"], HyperSphere(Point3f0(0.0), 20 | convert(Float32, 0.1)), 21 | MeshPhongMaterial(color = color)) 22 | 23 | setobject!(vis["ee_$i"], HyperSphere(Point3f0(0.0), 24 | convert(Float32, 0.05)), 25 | MeshPhongMaterial(color = color)) 26 | end 27 | 28 | function _set_cartpole!(vis, model, x; 29 | i = 0) 30 | 31 | px = x[1] + model.l * sin(x[2]) 32 | pz = -model.l * cos(x[2]) 33 | settransform!(vis["arm_$i"], cable_transform([x[1]; 0;0], [px; 0.0; pz])) 34 | settransform!(vis["base_$i"], Translation([x[1]; 0.0; 0.0])) 35 | settransform!(vis["ee_$i"], Translation([px; 0.0; pz])) 36 | end 37 | 38 | function visualize!(vis, model, q; 39 | i = 0, 40 | tl = 1.0, 41 | Δt = 0.1, 42 | color = Colors.RGBA(0,0,0,1.0)) 43 | 44 | default_background!(vis) 45 | _create_cartpole!(vis, model, i = i, color = color, tl = tl) 46 | 47 | anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt))) 48 | 49 | for t = 1:length(q) 50 | MeshCat.atframe(anim,t) do 51 | _set_cartpole!(vis, model, q[t], i = i) 52 | end 53 | end 54 | 55 | settransform!(vis["/Cameras/default"], 56 | compose(Translation(0.0, 0.0, -1.0), LinearMap(RotZ(- pi / 2)))) 57 | setvisible!(vis["/Grid"], false) 58 | 59 | MeshCat.setanimation!(vis,anim) 60 | end 61 | -------------------------------------------------------------------------------- /src/models/planar_push/codegen.jl: -------------------------------------------------------------------------------- 1 | path = @get_scratch!("planarpush") 2 | 3 | nq = planarpush.nq 4 | nu = planarpush.nu 5 | nc = planarpush.nc 6 | nz = num_var(planarpush) 7 | nθ = 2 * planarpush.nq + planarpush.nu + planarpush.nw + 1 8 | 9 | # Declare variables 10 | @variables z[1:nz] 11 | @variables θ[1:nθ] 12 | @variables κ[1:1] 13 | 14 | # Residual 15 | r_pp = residual(planarpush, z, θ, κ) 16 | rz_pp = Symbolics.jacobian(r_pp, z) 17 | rθ_pp = Symbolics.jacobian(r_pp, θ) 18 | 19 | # Build function 20 | r_pp_func = build_function(r_pp, z, θ, κ)[2] 21 | rz_pp_func = build_function(rz_pp, z, θ)[2] 22 | rθ_pp_func = build_function(rθ_pp, z, θ)[2] 23 | rz_pp_array = similar(rz_pp, Float64) 24 | rθ_pp_array = similar(rθ_pp, Float64) 25 | 26 | @save joinpath(path, "residual.jld2") r_pp_func rz_pp_func rθ_pp_func rz_pp_array rθ_pp_array 27 | @load joinpath(path, "residual.jld2") r_pp_func rz_pp_func rθ_pp_func rz_pp_array rθ_pp_array -------------------------------------------------------------------------------- /src/models/planar_push/model.jl: -------------------------------------------------------------------------------- 1 | """ 2 | planar push block 3 | particle with contacts at each corner 4 | """ 5 | struct PlanarPush{T} <: Model{T} 6 | # dimensions 7 | nq::Int # generalized coordinates 8 | nu::Int # controls 9 | nw::Int # parameters 10 | nc::Int # contact points 11 | 12 | mass_block::T 13 | mass_pusher::T 14 | 15 | inertia::T 16 | μ_surface::Vector{T} 17 | μ_pusher::T 18 | gravity::T 19 | 20 | contact_corner_offset::Vector{Vector{T}} 21 | end 22 | 23 | # Kinematics 24 | r_dim = 0.1 25 | 26 | function sd_2d_box(p, pose) 27 | x, y, θ = pose 28 | Δ = rotation_matrix(-θ) * (p - pose[1:2]) 29 | s = 10 30 | sum(Δ.^s)^(1/s) - r_dim 31 | end 32 | 33 | # contact corner 34 | cc1 = [r_dim, r_dim] 35 | cc2 = [-r_dim, r_dim] 36 | cc3 = [r_dim, -r_dim] 37 | cc4 = [-r_dim, -r_dim] 38 | 39 | contact_corner_offset = [cc1, cc2, cc3, cc4] 40 | 41 | # Parameters 42 | μ_surface = 0.5 # coefficient of friction 43 | μ_pusher = 0.5 44 | gravity = 9.81 45 | mass_block = 1.0 # mass 46 | mass_pusher = 10.0 47 | inertia = 1.0 / 12.0 * mass_block * ((2.0 * r_dim)^2 + (2.0 * r_dim)^2) 48 | 49 | # rnd = 0.01 50 | # dim = [r_dim, r_dim] 51 | # dim_rnd = [r_dim - rnd, r_dim - rnd] 52 | 53 | # Methods 54 | M_func(model::PlanarPush, q) = Diagonal([model.mass_block, model.mass_block, 55 | model.inertia, model.mass_pusher, model.mass_pusher]) 56 | 57 | function C_func(model::PlanarPush, q, q̇) 58 | [0.0, 0.0, 0.0, 0.0, 0.0] 59 | end 60 | 61 | function rotation_matrix(x) 62 | [cos(x) -sin(x); sin(x) cos(x)] 63 | end 64 | 65 | function ϕ_func(model::PlanarPush, q) 66 | p_block = q[1:3] 67 | p_pusher = q[4:5] 68 | 69 | sdf = sd_2d_box(p_pusher, p_block) 70 | 71 | [sdf] 72 | end 73 | 74 | function B_func(model::PlanarPush, q) 75 | [0.0 0.0; 76 | 0.0 0.0; 77 | 0.0 0.0; 78 | 1.0 0.0; 79 | 0.0 1.0] 80 | end 81 | 82 | function N_func(model::PlanarPush, q) 83 | ϕ = ϕ_func(model, q) 84 | vec(Symbolics.jacobian(ϕ, q)) 85 | end 86 | 87 | function p_func(model, x) 88 | pos = x[1:2] 89 | θ = x[3] 90 | R = rotation_matrix(θ) 91 | 92 | [(pos + R * model.contact_corner_offset[1])[1:2]; 93 | (pos + R * model.contact_corner_offset[2])[1:2]; 94 | (pos + R * model.contact_corner_offset[3])[1:2]; 95 | (pos + R * model.contact_corner_offset[4])[1:2]] 96 | end 97 | 98 | function P_func(model::PlanarPush, q) 99 | pf = p_func(model, q) 100 | P_block = Symbolics.jacobian(pf, q) 101 | 102 | p_block = q[1:3] 103 | p_pusher = q[4:5] 104 | 105 | ϕ = ϕ_func(model, q) 106 | N = vec(Symbolics.jacobian(ϕ, q)) 107 | 108 | N_pusher = N[4:5] 109 | 110 | n_dir = N_pusher ./ sqrt(N_pusher[1]^2.0 + N_pusher[2]^2.0) 111 | t_dir = [-n_dir[2]; n_dir[1]] 112 | 113 | r = p_pusher - p_block[1:2] 114 | m = r[1] * t_dir[2] - r[2] * t_dir[1] 115 | 116 | P = [t_dir[1]; t_dir[2]; m; -t_dir[1]; -t_dir[2]] 117 | 118 | return [P_block; transpose(P)] 119 | end 120 | 121 | function residual(model, z, θ, κ) 122 | nq = model.nq 123 | nu = model.nu 124 | nc = model.nc 125 | nc_impact = 1 126 | 127 | nb = 3 * 4 + 2 * 1 128 | 129 | q0 = θ[1:nq] 130 | q1 = θ[nq .+ (1:nq)] 131 | u1 = θ[2nq .+ (1:nu)] 132 | h = θ[2nq + nu .+ (1:1)] 133 | 134 | q2 = z[1:nq] 135 | γ1 = z[nq .+ (1:nc_impact)] 136 | s1 = z[nq + nc_impact .+ (1:nc_impact)] 137 | 138 | ψ1 = z[nq + 2 * nc_impact .+ (1:5)] 139 | b1 = z[nq + 2 * nc_impact + 5 .+ (1:9)] 140 | sψ1 = z[nq + 2 * nc_impact + 5 + 9 .+ (1:5)] 141 | sb1 = z[nq + 2 * nc_impact + 5 + 9 + 5 .+ (1:9)] 142 | 143 | ϕ = ϕ_func(model, q2) 144 | N = vec(Symbolics.jacobian(ϕ, q2)) 145 | 146 | # λ1 = [b1; γ1] 147 | P = P_func(model, q2) 148 | vT = P * (q2 - q1) ./ h[1] 149 | 150 | qm1 = 0.5 * (q0 + q1) 151 | vm1 = (q1 - q0) / h[1] 152 | qm2 = 0.5 * (q1 + q2) 153 | vm2 = (q2 - q1) / h[1] 154 | 155 | D1L1, D2L1 = lagrangian_derivatives(a -> M_func(model, a), (a, b) -> C_func(model, a, b), qm1, vm1) 156 | D1L2, D2L2 = lagrangian_derivatives(a -> M_func(model, a), (a, b) -> C_func(model, a, b), qm2, vm2) 157 | 158 | d = (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2# 159 | + B_func(model, qm2) * u1 160 | + N * γ1[1] 161 | + transpose(P) * b1) 162 | 163 | [ 164 | d; 165 | 166 | s1 .- ϕ; 167 | 168 | ψ1[1] .- model.μ_surface[1] * model.mass_block * model.gravity * h[1] * 0.25; 169 | 170 | ψ1[2] .- model.μ_surface[2] * model.mass_block * model.gravity * h[1] * 0.25; 171 | 172 | ψ1[3] .- model.μ_surface[3] * model.mass_block * model.gravity * h[1] * 0.25; 173 | 174 | ψ1[4] .- model.μ_surface[4] * model.mass_block * model.gravity * h[1] * 0.25; 175 | 176 | ψ1[5] .- model.μ_pusher * γ1[1]; 177 | 178 | vT - sb1; 179 | 180 | γ1 .* s1 .- κ[1]; 181 | cone_product([ψ1[1]; b1[1:2]], [sψ1[1]; sb1[1:2]]) - [κ[1]; 0.0; 0.0]; 182 | cone_product([ψ1[2]; b1[2 .+ (1:2)]], [sψ1[2]; sb1[2 .+ (1:2)]]) - [κ[1]; 0.0; 0.0]; 183 | cone_product([ψ1[3]; b1[4 .+ (1:2)]], [sψ1[3]; sb1[4 .+ (1:2)]]) - [κ[1]; 0.0; 0.0]; 184 | cone_product([ψ1[4]; b1[6 .+ (1:2)]], [sψ1[4]; sb1[6 .+ (1:2)]]) - [κ[1]; 0.0; 0.0]; 185 | cone_product([ψ1[5]; b1[8 .+ (1:1)]], [sψ1[5]; sb1[8 .+ (1:1)]]) - [κ[1]; 0.0]; 186 | ] 187 | end 188 | 189 | # Dimensions 190 | nq = 5 # configuration dimension 191 | nu = 2 # control dimension 192 | nc = 5 # number of contact points 193 | nc_impact = 1 194 | nf = 3 # number of faces for friction cone pyramid 195 | nb = (nc - nc_impact) * nf + (nf - 1) * nc_impact 196 | 197 | planarpush = PlanarPush(nq, nu, 0, nc, 198 | mass_block, mass_pusher, 199 | inertia, [μ_surface for i = 1:nc], μ_pusher, gravity, 200 | contact_corner_offset) 201 | -------------------------------------------------------------------------------- /src/models/planar_push/simulator.jl: -------------------------------------------------------------------------------- 1 | function RoboDojo.indices_z(model::PlanarPush) 2 | nq = model.nq 3 | nc = model.nc 4 | q = collect(1:nq) 5 | γ = collect(nq .+ (1:1)) 6 | sγ = collect(nq + 1 .+ (1:1)) 7 | ψ = collect(nq + 2 .+ (1:5)) 8 | b = collect(nq + 2 + 5 .+ (1:9)) 9 | sψ = collect(nq + 2 + 5 + 9 .+ (1:5)) 10 | sb = collect(nq + 2 + 5 + 9 + 5 .+ (1:9)) 11 | IndicesZ(q, γ, sγ, ψ, b, sψ, sb) 12 | end 13 | 14 | RoboDojo.nominal_configuration(model::PlanarPush) = [-1.0; 0.0; 0.0; 0.0; 0.0] 15 | 16 | function RoboDojo.indices_optimization(model::PlanarPush) 17 | nq = model.nq 18 | nz = num_var(model) 19 | IndicesOptimization( 20 | nz, 21 | nz, 22 | [collect(nq .+ (1:1)), collect(nq + 1 .+ (1:1))], 23 | [collect(nq .+ (1:1)), collect(nq + 1 .+ (1:1))], 24 | [ 25 | [collect([collect(nq + 2 .+ (1:1)); collect(nq + 2 + 5 .+ (1:2))]), collect([collect(nq + 2 + 5 + 9 .+ (1:1)); collect(nq + 2 + 5 + 9 + 5 .+ (1:2))])], 26 | [collect([collect(nq + 2 + 1 .+ (1:1)); collect(nq + 2 + 5 + 2 .+ (1:2))]), collect([collect(nq + 2 + 5 + 9 + 1 .+ (1:1)); collect(nq + 2 + 5 + 9 + 5 + 2 .+ (1:2))])], 27 | [collect([collect(nq + 2 + 2 .+ (1:1)); collect(nq + 2 + 5 + 4 .+ (1:2))]), collect([collect(nq + 2 + 5 + 9 + 2 .+ (1:1)); collect(nq + 2 + 5 + 9 + 5 + 4 .+ (1:2))])], 28 | [collect([collect(nq + 2 + 3 .+ (1:1)); collect(nq + 2 + 5 + 6 .+ (1:2))]), collect([collect(nq + 2 + 5 + 9 + 3 .+ (1:1)); collect(nq + 2 + 5 + 9 + 5 + 6 .+ (1:2))])], 29 | [collect([collect(nq + 2 + 4 .+ (1:1)); collect(nq + 2 + 5 + 8 .+ (1:1))]), collect([collect(nq + 2 + 5 + 9 + 4 .+ (1:1)); collect(nq + 2 + 5 + 9 + 5 + 8 .+ (1:1))])], 30 | ], 31 | [ 32 | [collect([collect(nq + 2 .+ (1:1)); collect(nq + 2 + 5 .+ (1:2))]), collect([collect(nq + 2 + 5 + 9 .+ (1:1)); collect(nq + 2 + 5 + 9 + 5 .+ (1:2))])], 33 | [collect([collect(nq + 2 + 1 .+ (1:1)); collect(nq + 2 + 5 + 2 .+ (1:2))]), collect([collect(nq + 2 + 5 + 9 + 1 .+ (1:1)); collect(nq + 2 + 5 + 9 + 5 + 2 .+ (1:2))])], 34 | [collect([collect(nq + 2 + 2 .+ (1:1)); collect(nq + 2 + 5 + 4 .+ (1:2))]), collect([collect(nq + 2 + 5 + 9 + 2 .+ (1:1)); collect(nq + 2 + 5 + 9 + 5 + 4 .+ (1:2))])], 35 | [collect([collect(nq + 2 + 3 .+ (1:1)); collect(nq + 2 + 5 + 6 .+ (1:2))]), collect([collect(nq + 2 + 5 + 9 + 3 .+ (1:1)); collect(nq + 2 + 5 + 9 + 5 + 6 .+ (1:2))])], 36 | [collect([collect(nq + 2 + 4 .+ (1:1)); collect(nq + 2 + 5 + 8 .+ (1:1))]), collect([collect(nq + 2 + 5 + 9 + 4 .+ (1:1)); collect(nq + 2 + 5 + 9 + 5 + 8 .+ (1:1))])], 37 | ], 38 | collect(1:(nq + 15)), 39 | collect(nq + 15 .+ (1:1)), 40 | collect(nq + 15 + 1 .+ (1:14)), 41 | [ 42 | collect(nq + 15 + 1 .+ (1:3)), 43 | collect(nq + 15 + 4 .+ (1:3)), 44 | collect(nq + 15 + 7 .+ (1:3)), 45 | collect(nq + 15 + 10 .+ (1:3)), 46 | collect(nq + 15 + 13 .+ (1:2)), 47 | ], 48 | collect(nq + 15 .+ (1:15)) 49 | ) 50 | end 51 | 52 | function RoboDojo.initialize_z!(z, model::PlanarPush, idx::IndicesZ, q) 53 | z[idx.q] .= q 54 | z[idx.γ] .= 1.0 55 | z[idx.sγ] .= 1.0 56 | z[idx.ψ] .= 1.0 57 | z[idx.b] .= 0.1 58 | z[idx.sψ] .= 1.0 59 | z[idx.sb] .= 0.1 60 | end 61 | 62 | num_var(model::PlanarPush) = model.nq + 2 * 1 + 2 * 14 63 | friction_coefficients(model::PlanarPush{T}) where T = T[] 64 | 65 | 66 | -------------------------------------------------------------------------------- /src/models/planar_push/visuals.jl: -------------------------------------------------------------------------------- 1 | function _create_planar_push!(vis, model::PlanarPush; 2 | i = 1, 3 | r = 0.1, 4 | r_pusher = 0.025, 5 | tl = 1.0, 6 | box_color = Colors.RGBA(0.0, 0.0, 0.0, tl), 7 | pusher_color = Colors.RGBA(0.5, 0.5, 0.5, tl)) 8 | 9 | r_box = r - r_pusher 10 | 11 | setobject!(vis["box_$i"], GeometryBasics.Rect(Vec(-1.0 * r_box, 12 | -1.0 * r_box, 13 | -1.0 * r_box), 14 | Vec(2.0 * r_box, 2.0 * r_box, 2.0 * r_box)), 15 | MeshPhongMaterial(color = box_color)) 16 | 17 | setobject!(vis["pusher_$i"], 18 | Cylinder(Point(0.0, 0.0, 0.0), Point(0.0, 0.0, r_box), r_pusher), 19 | MeshPhongMaterial(color = pusher_color)) 20 | end 21 | 22 | function _set_planar_push!(vis, model::PlanarPush, q; 23 | i = 1) 24 | settransform!(vis["box_$i"], 25 | compose(Translation(q[1], q[2], 0.01 * i), LinearMap(RotZ(q[3])))) 26 | settransform!(vis["pusher_$i"], Translation(q[4], q[5], 0.01 * i)) 27 | end 28 | 29 | function visualize!(vis, model::PlanarPush, q; 30 | i = 1, 31 | r = 0.1, 32 | r_pusher = 0.025, 33 | tl = 1.0, 34 | box_color = Colors.RGBA(0.0, 0.0, 0.0, tl), 35 | pusher_color = Colors.RGBA(0.5, 0.5, 0.5, tl), 36 | Δt = 0.1) 37 | 38 | default_background!(vis) 39 | 40 | _create_planar_push!(vis, model, 41 | i = i, 42 | r = r, 43 | r_pusher = r_pusher, 44 | tl = tl, 45 | box_color = box_color, 46 | pusher_color = pusher_color) 47 | 48 | anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt))) 49 | 50 | T = length(q) 51 | for t = 1:T-1 52 | MeshCat.atframe(anim, t) do 53 | _set_planar_push!(vis, model, q[t]) 54 | end 55 | end 56 | 57 | settransform!(vis["/Cameras/default"], 58 | compose(Translation(0.0, 0.0, 50.0), LinearMap(RotZ(0.5 * pi) * RotY(-pi/2.5)))) 59 | setprop!(vis["/Cameras/default/rotated/"], "zoom", 50) 60 | 61 | 62 | MeshCat.setanimation!(vis, anim) 63 | end 64 | -------------------------------------------------------------------------------- /src/models/rocket/codegen.jl: -------------------------------------------------------------------------------- 1 | path = @get_scratch!("rocket") 2 | 3 | # rocket 4 | nx = rocket.nq 5 | nu = rocket.nu 6 | nw = rocket.nw 7 | 8 | nz = nx 9 | nθ = nx + nu + 1 10 | 11 | @variables z[1:nz] θ[1:nθ] κ[1:1] 12 | 13 | # Residual 14 | function residual(rocket, z, θ, κ) 15 | y = z[1:nx] 16 | x = θ[1:nx] 17 | u = θ[nx .+ (1:nu)] 18 | h = θ[nx + nu .+ (1:1)] 19 | w = zeros(rocket.nw) 20 | 21 | return y - (x + h[1] * f(rocket, 0.5 * (x + y), u, w)) 22 | end 23 | 24 | r_rocket = residual(rocket, z, θ, κ) 25 | r_rocket = Symbolics.simplify.(r_rocket) 26 | rz_rocket = Symbolics.jacobian(r_rocket, z, simplify=true) 27 | rθ_rocket = Symbolics.jacobian(r_rocket, θ, simplify=true) 28 | 29 | # Build function 30 | r_rocket_func = build_function(r_rocket, z, θ, κ)[2] 31 | rz_rocket_func = build_function(rz_rocket, z, θ)[2] 32 | rθ_rocket_func = build_function(rθ_rocket, z, θ)[2] 33 | 34 | rz_rocket_array = similar(rz_rocket, Float64) 35 | rθ_rocket_array = similar(rθ_rocket, Float64) 36 | 37 | @save joinpath(path, "residual.jld2") r_rocket_func rz_rocket_func rθ_rocket_func rz_rocket_array rθ_rocket_array 38 | @load joinpath(path, "residual.jld2") r_rocket_func rz_rocket_func rθ_rocket_func rz_rocket_array rθ_rocket_array 39 | 40 | 41 | # projection 42 | nz = 3 + 1 + 1 + 1 + 1 + 3 43 | nθ = 3 + 1 44 | 45 | function residual_projection(z, θ, κ) 46 | u = z[1:3] 47 | p = z[4:4] 48 | s = z[5:5] 49 | w = z[6:6] 50 | y = z[7:7] 51 | v = z[8:10] 52 | 53 | ū = θ[1:3] 54 | uu = θ[3 .+ (1:1)] 55 | idx = [3; 1; 2] 56 | [ 57 | u - ū - v - [0.0; 0.0; y[1] + p[1]]; 58 | uu[1] - u[3] - s[1]; 59 | -y[1] - w[1]; 60 | w[1] * s[1] - κ[1]; 61 | p[1] * u[3] - κ[1]; 62 | cone_product(u[idx], v[idx]) - [κ[1]; 0.0; 0.0] 63 | ] 64 | end 65 | 66 | @variables z[1:nz], θ[1:nθ], κ[1:1] 67 | 68 | r_proj = residual_projection(z, θ, κ) 69 | r_proj .= simplify(r_proj) 70 | r_proj_func = Symbolics.build_function(r_proj, z, θ, κ)[2] 71 | 72 | rz_proj = Symbolics.jacobian(r_proj, z) 73 | rz_proj = simplify.(rz_proj) 74 | rz_proj_func = Symbolics.build_function(rz_proj, z, θ)[2] 75 | rθ_proj = Symbolics.jacobian(r_proj, θ) 76 | rθ_proj = simplify.(rθ_proj) 77 | rθ_proj_func = Symbolics.build_function(rθ_proj, z, θ)[2] 78 | 79 | rz_proj_array = similar(rz_proj, Float64) 80 | rθ_proj_array = similar(rθ_proj, Float64) 81 | 82 | @save joinpath(path, "projection.jld2") r_proj_func rz_proj_func rθ_proj_func rz_proj_array rθ_proj_array 83 | @load joinpath(path, "projection.jld2") r_proj_func rz_proj_func rθ_proj_func rz_proj_array rθ_proj_array -------------------------------------------------------------------------------- /src/models/rocket/dynamics.jl: -------------------------------------------------------------------------------- 1 | struct RocketInfo1{T,R,RZ,Rθ} 2 | ip_dyn::InteriorPoint{T,R,RZ,Rθ} 3 | ip_proj::InteriorPoint{T,R,RZ,Rθ} 4 | idx_z::IndicesZ 5 | idx_θ::Indicesθ 6 | idx3::Vector{Int} 7 | h::T 8 | u_max::T 9 | du_dyn_cache::Matrix{T} 10 | du_proj_cache::Matrix{T} 11 | end 12 | 13 | function RocketInfo(rocket, u_max, h, r_func, rz_func, rθ_func, r_func_proj, rz_func_proj, rθ_func_proj) 14 | nx = rocket.nq 15 | nu = rocket.nu 16 | nw = rocket.nw 17 | 18 | nz = nx 19 | nθ = nx + nu + 1 20 | 21 | solver_opts=InteriorPointOptions( 22 | r_tol=1e-8, 23 | κ_tol=1.0, 24 | max_ls=25, 25 | ϵ_min=0.25, 26 | diff_sol=true, 27 | verbose=false) 28 | 29 | z0 = zeros(nx) 30 | θ0 = zeros(nθ) 31 | 32 | idx_opt = RoboDojo.indices_optimization(rocket) 33 | 34 | ip_dyn = interior_point( 35 | z0, 36 | θ0, 37 | idx=idx_opt, 38 | r! = r_func, 39 | rz! = rz_func, 40 | rθ! = rθ_func, 41 | rz=zeros(nz, nz), 42 | rθ=zeros(nz, nθ), 43 | opts=solver_opts) 44 | 45 | idx_z = RoboDojo.indices_z(rocket) 46 | idx_θ = RoboDojo.indices_θ(rocket) 47 | 48 | # projection 49 | nz_proj = 10 50 | nθ_proj = 4 51 | 52 | idx_opt_proj = IndicesOptimization( 53 | 10, 54 | 10, 55 | [[5, 3], [6, 4]], 56 | [[5, 3], [6, 4]], 57 | [[collect([3, 1, 2]), collect([10, 8, 9])]], 58 | [[collect([3, 1, 2]), collect([10, 8, 9])]], 59 | collect(1:5), 60 | collect(6:7), 61 | collect(8:10), 62 | [collect(8:10)], 63 | collect(6:10)) 64 | 65 | z0_proj = zeros(nz_proj) 66 | θ0_proj = zeros(nθ_proj) 67 | 68 | ip_proj = interior_point( 69 | z0_proj, 70 | θ0_proj, 71 | idx=idx_opt_proj, 72 | r! = r_func_proj, 73 | rz! = rz_func_proj, 74 | rθ! = rθ_func_proj, 75 | rz=zeros(nz_proj, nz_proj), 76 | rθ=zeros(nz_proj, nθ_proj), 77 | opts=InteriorPointOptions( 78 | r_tol=1e-8, 79 | κ_tol=1.0e-4, 80 | max_ls=25, 81 | ϵ_min=0.0, 82 | undercut=Inf, 83 | γ_reg=0.0, 84 | κ_reg=0.0, 85 | diff_sol=false, 86 | verbose=false)) 87 | 88 | RocketInfo1( 89 | ip_dyn, 90 | ip_proj, 91 | idx_z, 92 | idx_θ, 93 | collect(1:3), 94 | h, 95 | u_max, 96 | zeros(nx, nu), 97 | zeros(nu, nu) 98 | ) 99 | end 100 | 101 | function f_rocket(d, info::RocketInfo1, x, u, w) 102 | # initialize 103 | info.ip_dyn.z[info.idx_z.q] .= x 104 | info.ip_dyn.θ[info.idx_θ.q1] .= x 105 | info.ip_dyn.θ[info.idx_θ.u] .= u 106 | info.ip_dyn.θ[info.idx_θ.h] .= info.h 107 | # solve 108 | info.ip_dyn.opts.diff_sol = false 109 | interior_point_solve!(info.ip_dyn) 110 | # solution 111 | sol = @views info.ip_dyn.z[info.idx_z.q] 112 | d .= sol 113 | return d 114 | end 115 | 116 | # d = zeros(nx) 117 | # dx = zeros(nx, nx) 118 | # du = zeros(nx, nu) 119 | 120 | # x = rand(nx) 121 | # u = rand(nu) 122 | # w = rand(nw) 123 | 124 | # using BenchmarkTools 125 | # f_rocket(d, info, x, u, w) 126 | # @benchmark f_rocket($d, $info, $x, $u, $w) 127 | 128 | # fx_rocket(dx, info, x, u, w) 129 | # @benchmark fx_rocket($dx, $info, $x, $u, $w) 130 | 131 | # fu_rocket(du, info, x, u, w) 132 | # @benchmark fu_rocket($du, $info, $x, $u, $w) 133 | 134 | function fx_rocket(dx, info::RocketInfo1, x, u, w) 135 | # initialize 136 | info.ip_dyn.z[info.idx_z.q] .= x 137 | info.ip_dyn.θ[info.idx_θ.q1] .= x 138 | info.ip_dyn.θ[info.idx_θ.u] .= u 139 | info.ip_dyn.θ[info.idx_θ.h] .= info.h 140 | # solve 141 | info.ip_dyn.opts.diff_sol = true 142 | interior_point_solve!(info.ip_dyn) 143 | # solution 144 | sol = @views info.ip_dyn.δz[info.idx_z.q, info.idx_θ.q1] 145 | dx .= sol 146 | return dx 147 | end 148 | 149 | function fu_rocket(du, info::RocketInfo1, x, u, w) 150 | # initialize 151 | info.ip_dyn.z[info.idx_z.q] .= x 152 | info.ip_dyn.θ[info.idx_θ.q1] .= x 153 | info.ip_dyn.θ[info.idx_θ.u] .= u 154 | info.ip_dyn.θ[info.idx_θ.h] .= info.h 155 | # solve 156 | info.ip_dyn.opts.diff_sol = true 157 | interior_point_solve!(info.ip_dyn) 158 | 159 | # solution 160 | sol = @views info.ip_dyn.δz[info.idx_z.q, info.idx_θ.u] 161 | du .= sol 162 | return du 163 | end 164 | 165 | # info = RocketInfo(rocket, eval(r_func), eval(rz_func), eval(rθ_func), eval(r_func_proj), eval(rz_func_proj), eval(rθ_func_proj)) 166 | 167 | 168 | function soc_projection(x, info::RocketInfo1) 169 | info.ip_proj.z .= 0.1 170 | info.ip_proj.z[3] += 1.0 171 | info.ip_proj.z[10] += 1.0 172 | info.ip_proj.z[7] = 0.0 173 | 174 | info.ip_proj.θ[info.idx3] .= x 175 | info.ip_proj.θ[4] = info.u_max 176 | 177 | info.ip_proj.opts.diff_sol = false 178 | status = interior_point_solve!(info.ip_proj) 179 | 180 | # !status && (@warn "projection failure (res norm: $(norm(info.ip_proj.r, Inf))) \n 181 | # z = $(info.ip_proj.z), \n 182 | # θ = $(info.ip_proj.θ)") 183 | 184 | sol = @views info.ip_proj.z[info.idx3] 185 | return sol 186 | end 187 | # soc_projection(u, info) 188 | # @benchmark soc_projection($u, $info) 189 | 190 | function soc_projection_gradient(x, info::RocketInfo1) 191 | info.ip_proj.z .= 0.1 192 | info.ip_proj.z[3] += 1.0 193 | info.ip_proj.z[10] += 1.0 194 | info.ip_proj.z[7] = 0.0 195 | 196 | # info.ip_proj.θ .= [x; uu] 197 | info.ip_proj.θ[info.idx3] .= x 198 | info.ip_proj.θ[4] = info.u_max 199 | 200 | info.ip_proj.opts.diff_sol = true 201 | status = interior_point_solve!(info.ip_proj) 202 | 203 | # !status && (@warn "projection failure (res norm: $(norm(info.ip_proj.r, Inf))) \n 204 | # z = $(info.ip_proj.z), \n 205 | # θ = $(info.ip_proj.θ)") 206 | 207 | sol = @views info.ip_proj.δz[info.idx3, info.idx3] 208 | 209 | return sol 210 | end 211 | 212 | # soc_projection_gradient(u, info) 213 | # @benchmark soc_projection_gradient($u, $info) 214 | 215 | function f_rocket_proj(d, info::RocketInfo1, x, u, w) 216 | # initialize 217 | info.ip_dyn.z[info.idx_z.q] .= x 218 | info.ip_dyn.θ[info.idx_θ.q1] .= x 219 | info.ip_dyn.θ[info.idx_θ.u] .= soc_projection(u, info) 220 | info.ip_dyn.θ[info.idx_θ.h] .= info.h 221 | # solve 222 | info.ip_dyn.opts.diff_sol = false 223 | interior_point_solve!(info.ip_dyn) 224 | # solution 225 | sol = @views info.ip_dyn.z[info.idx_z.q] 226 | d .= sol 227 | return d 228 | end 229 | 230 | # f_rocket_proj(d, info, x, u, w) 231 | # @benchmark f_rocket_proj($d, $info, $x, $u, $w) 232 | 233 | # fx_rocket_proj(dx, info, x, u, w) 234 | # @benchmark fx_rocket_proj($dx, $info, $x, $u, $w) 235 | 236 | # fu_rocket_proj(du, info, x, u, w) 237 | # @benchmark fu_rocket_proj($du, $info, $x, $u, $w) 238 | 239 | function fx_rocket_proj(dx, info::RocketInfo1, x, u, w) 240 | # initialize 241 | info.ip_dyn.z[info.idx_z.q] .= x 242 | info.ip_dyn.θ[info.idx_θ.q1] .= x 243 | info.ip_dyn.θ[info.idx_θ.u] .= soc_projection(u, info) 244 | info.ip_dyn.θ[info.idx_θ.h] .= info.h 245 | # solve 246 | info.ip_dyn.opts.diff_sol = true 247 | interior_point_solve!(info.ip_dyn) 248 | # solution 249 | sol = @views info.ip_dyn.δz[info.idx_z.q, info.idx_θ.q1] 250 | dx .= sol 251 | return dx 252 | end 253 | 254 | function fu_rocket_proj(du, info::RocketInfo1, x, u, w) 255 | # initialize 256 | info.ip_dyn.z[info.idx_z.q] .= x 257 | info.ip_dyn.θ[info.idx_θ.q1] .= x 258 | info.ip_dyn.θ[info.idx_θ.u] .= soc_projection(u, info) 259 | info.ip_dyn.θ[info.idx_θ.h] .= info.h 260 | # solve 261 | info.ip_dyn.opts.diff_sol = true 262 | interior_point_solve!(info.ip_dyn) 263 | # solution 264 | info.du_dyn_cache .= @views info.ip_dyn.δz[info.idx_z.q, info.idx_θ.u] 265 | info.du_proj_cache .= soc_projection_gradient(u, info) 266 | # du .= sol_dyn * sol_proj 267 | mul!(du, info.du_dyn_cache, info.du_proj_cache) 268 | return du 269 | end -------------------------------------------------------------------------------- /src/models/rocket/model.jl: -------------------------------------------------------------------------------- 1 | struct Rocket{T} <: Model{T} 2 | nq::Int 3 | nu::Int 4 | nw::Int 5 | nc::Int 6 | 7 | mass::T # mass 8 | inertia::Diagonal{T,Vector{T}} # inertia matrix 9 | inertia_inv::Diagonal{T,Vector{T}} # inertia matrix inverse 10 | gravity::Vector{T} # gravity 11 | length::T # length (com to thruster) 12 | end 13 | 14 | function f(model::Rocket, z, u, w) 15 | # states 16 | x = z[1:3] 17 | r = z[4:6] 18 | v = z[7:9] 19 | ω = z[10:12] 20 | 21 | # force in body frame 22 | F = u[1:3] 23 | 24 | # torque in body frame 25 | τ = [model.length * u[2], 26 | -model.length * u[1], 27 | 0.0] 28 | 29 | [v; 30 | 0.25 * ((1.0 - r' * r) * ω - 2.0 * cross(ω, r) + 2.0*(ω' * r) * r); 31 | model.gravity + (1.0 / model.mass) * MRP(r[1], r[2], r[3]) * F; 32 | model.inertia_inv * (τ - cross(ω, model.inertia * ω))] 33 | end 34 | 35 | nq, nu, nw = 12, 3, 0 36 | 37 | mass = 1.0 38 | len = 1.0 39 | 40 | inertia = Diagonal([1.0 / 12.0 * mass * len^2.0, 1.0 / 12.0 * mass * len^2.0, 1.0e-5]) 41 | inertia_inv = Diagonal([1.0 / (1.0 / 12.0 * mass * len^2.0), 1.0 / (1.0 / 12.0 * mass * len^2.0), 1.0 / (1.0e-5)]) 42 | 43 | rocket = Rocket(nq, nu, nw, 0, 44 | mass, 45 | inertia, 46 | inertia_inv, 47 | [0.0, 0.0, -9.81], 48 | len) 49 | 50 | -------------------------------------------------------------------------------- /src/models/rocket/simulator.jl: -------------------------------------------------------------------------------- 1 | friction_coefficients(model::Rocket{T}) where T = T[] 2 | 3 | function RoboDojo.indices_z(model::Rocket) 4 | nq = model.nq 5 | q = collect(1:nq) 6 | γ = collect(1:0) 7 | sγ = collect(1:0) 8 | ψ = collect(1:0) 9 | b = collect(1:0) 10 | sψ = collect(1:0) 11 | sb = collect(1:0) 12 | IndicesZ(q, γ, sγ, ψ, b, sψ, sb) 13 | end 14 | 15 | function RoboDojo.indices_θ(model::Rocket; nf=0) 16 | nq = model.nq 17 | nu = model.nu 18 | nw = model.nw 19 | 20 | q1 = collect(1:nq) 21 | q2 = collect(1:0) 22 | u = collect(nq .+ (1:nu)) 23 | w = collect(1:0) 24 | f = collect(1:0) 25 | h = collect(nq + nu .+ (1:1)) 26 | 27 | Indicesθ(q1, q2, u, w, f, h) 28 | end 29 | 30 | RoboDojo.num_var(model::Rocket) = model.nq 31 | 32 | RoboDojo.nominal_configuration(model::Rocket) = zeros(model.nq) 33 | 34 | function RoboDojo.indices_optimization(model::Rocket) 35 | nq = model.nq 36 | nz = nq 37 | IndicesOptimization( 38 | nz, 39 | nz, 40 | [collect(1:0), collect(1:0)], 41 | [collect(1:0), collect(1:0)], 42 | Vector{Vector{Vector{Int}}}(), 43 | Vector{Vector{Vector{Int}}}(), 44 | collect(1:(nq)), 45 | collect(1:0), 46 | collect(1:0), 47 | Vector{Vector{Int}}(), 48 | collect(1:0)) 49 | end 50 | 51 | function RoboDojo.initialize_z!(z, model::Rocket, idx::IndicesZ, q) 52 | z[idx.q] .= q 53 | end -------------------------------------------------------------------------------- /src/models/rocket/visuals.jl: -------------------------------------------------------------------------------- 1 | function visualize!(vis, p::Rocket, q; 2 | Δt=0.1, mesh=false, T_off=length(q)) 3 | default_background!(vis) 4 | 5 | if mesh 6 | obj_rocket = joinpath(pwd(), "models/starship/Starship.obj") 7 | mtl_rocket = joinpath(pwd(), "models/starship/Starship.mtl") 8 | ctm = ModifiedMeshFileObject(obj_rocket, mtl_rocket, scale=1.0) 9 | setobject!(vis["rocket"]["starship"], ctm) 10 | 11 | settransform!(vis["rocket"]["starship"], 12 | compose(Translation(0.0, 0.0, -p.length), 13 | LinearMap(0.25 * RotY(0.0) * RotZ(0.5 * π) * RotX(0.5 * π)))) 14 | 15 | body = Cylinder(Point3f0(0.0, 0.0, -1.25), 16 | Point3f0(0.0, 0.0, 0.5), 17 | convert(Float32, 0.125)) 18 | 19 | setobject!(vis["rocket"]["body"], body, 20 | MeshPhongMaterial(color = RGBA(1.0, 0.0, 0.0, 1.0))) 21 | else 22 | body = Cylinder(Point3f0(0.0, 0.0, -1.0 * p.length), 23 | Point3f0(0.0, 0.0, 1.0 * p.length), 24 | convert(Float32, 0.15)) 25 | 26 | setobject!(vis["rocket"], body, 27 | MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0))) 28 | anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt))) 29 | end 30 | 31 | anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt))) 32 | 33 | for t = 1:length(q) 34 | MeshCat.atframe(anim, t) do 35 | if t >= T_off 36 | setvisible!(vis["rocket"]["body"], false) 37 | else 38 | setvisible!(vis["rocket"]["body"], true) 39 | end 40 | settransform!(vis["rocket"], 41 | compose(Translation(q[t][1:3]), 42 | LinearMap(MRP(q[t][4:6]...)))) 43 | end 44 | end 45 | 46 | MeshCat.setanimation!(vis, anim) 47 | end -------------------------------------------------------------------------------- /src/models/visualize.jl: -------------------------------------------------------------------------------- 1 | function cable_transform(y, z) 2 | v1 = [0.0, 0.0, 1.0] 3 | v2 = y[1:3,1] - z[1:3,1] 4 | normalize!(v2) 5 | ax = cross(v1, v2) 6 | ang = acos(v1'*v2) 7 | R = AngleAxis(ang, ax...) 8 | 9 | if any(isnan.(R)) 10 | R = I 11 | else 12 | nothing 13 | end 14 | 15 | compose(Translation(z), LinearMap(R)) 16 | end 17 | 18 | function default_background!(vis) 19 | setvisible!(vis["/Background"], true) 20 | setprop!(vis["/Background"], "top_color", RGBA(1.0, 1.0, 1.0, 1.0)) 21 | setprop!(vis["/Background"], "bottom_color", RGBA(1.0, 1.0, 1.0, 1.0)) 22 | setvisible!(vis["/Axes"], false) 23 | end 24 | --------------------------------------------------------------------------------