Begin adding timeit's

RoboticExplorationLab · Jul 26, 2019 · 73b9d34 · 73b9d34
1 parent 1b99400
commit 73b9d34
Show file tree

Hide file tree

Showing 6 changed files with 28 additions and 16 deletions.
diff --git a/src/TrajectoryOptimization.jl b/src/TrajectoryOptimization.jl
@@ -24,6 +24,7 @@ using Parameters
 using Rotations
 using BlockArrays
 using MathOptInterface
+using TimerOutputs
 
 export
     Dynamics,

diff --git a/src/solvers/augmented_lagrangian/augmented_lagrangian_methods.jl b/src/solvers/augmented_lagrangian/augmented_lagrangian_methods.jl
@@ -52,14 +52,15 @@ end
 "Augmented Lagrangian step"
 function step!(prob::Problem{T}, solver::AugmentedLagrangianSolver{T},
         unconstrained_solver::AbstractSolver) where T
+    to = solver.stats[:timer]
 
     # Solve the unconstrained problem
-    solve!(prob, unconstrained_solver)
-    J = cost(prob)
+    @timeit to "uncon solve" solve!(prob, unconstrained_solver)
+    @timeit to "cost" J = cost(prob)
 
     # Outer loop update
-    dual_update!(prob, solver)
-    penalty_update!(prob, solver)
+    @timeit to "dual update" dual_update!(prob, solver)
+    @timeit to "penalty update" penalty_update!(prob, solver)
     copyto!(solver.C_prev,solver.C)
 
     return J

diff --git a/src/solvers/augmented_lagrangian/augmented_lagrangian_solver.jl b/src/solvers/augmented_lagrangian/augmented_lagrangian_solver.jl
@@ -121,8 +121,11 @@ function AbstractSolver(prob::Problem{T,D}, opts::AugmentedLagrangianSolverOptio
     # check for conflicting convergence criteria between unconstrained solver and AL: warn
 
     # Init solver statistics
+    to = TimerOutput()
+    reset_timer!(to)
     stats = Dict{Symbol,Any}(:iterations=>0,:iterations_total=>0,
-        :iterations_inner=>Int[],:cost=>T[],:c_max=>T[],:penalty_max=>[])
+        :iterations_inner=>Int[],:cost=>T[],:c_max=>T[],:penalty_max=>[],
+        :timer=>to)
     stats_uncon = Dict{Symbol,Any}[]
 
     # Init solver results
@@ -131,6 +134,7 @@ function AbstractSolver(prob::Problem{T,D}, opts::AugmentedLagrangianSolverOptio
     C,∇C,λ,μ,active_set = init_constraint_trajectories(prob.constraints,n,m,N)
 
     solver_uncon = AbstractSolver(prob,opts.opts_uncon)
+    solver_uncon.stats[:timer] = to
 
     AugmentedLagrangianSolver{T}(opts,stats,stats_uncon,C,copy(C),∇C,λ,μ,active_set,solver_uncon)
 end

diff --git a/src/solvers/ilqr/forward_pass.jl b/src/solvers/ilqr/forward_pass.jl
@@ -13,6 +13,7 @@ function forwardpass!(prob::Problem, solver::iLQRSolver, ΔV::Array, J_prev::Flo
     expected = 0.
 
     logger = current_logger()
+    to = solver.stats[:timer]
     # @logmsg InnerIters :iter value=0
     # @logmsg InnerIters :cost value=J_prev
     while (z ≤ solver.opts.line_search_lower_bound || z > solver.opts.line_search_upper_bound) && J >= J_prev
@@ -36,7 +37,7 @@ function forwardpass!(prob::Problem, solver::iLQRSolver, ΔV::Array, J_prev::Flo
         end
 
         # Otherwise, rollout a new trajectory for current alpha
-        flag = rollout!(prob,solver,alpha)
+        @timeit to "rollout" flag = rollout!(prob,solver,alpha)
 
         # Check if rollout completed
         if ~flag
@@ -48,7 +49,7 @@ function forwardpass!(prob::Problem, solver::iLQRSolver, ΔV::Array, J_prev::Flo
         end
 
         # Calcuate cost
-        J = cost(prob.obj, X̄, Ū, get_dt_traj(prob,Ū))   # Unconstrained cost
+        @timeit to "cost" J = cost(prob.obj, X̄, Ū, get_dt_traj(prob,Ū))   # Unconstrained cost
 
         expected = -alpha*(ΔV[1] + alpha*ΔV[2])
         if expected > 0

diff --git a/src/solvers/ilqr/ilqr_methods.jl b/src/solvers/ilqr/ilqr_methods.jl
@@ -2,6 +2,7 @@
 "iLQR solve method"
 function solve!(prob::Problem{T,Discrete}, solver::iLQRSolver{T}) where T<:AbstractFloat
     reset!(solver)
+    to = solver.stats[:timer]
 
     n,m,N = size(prob)
     J = Inf
@@ -13,9 +14,10 @@ function solve!(prob::Problem{T,Discrete}, solver::iLQRSolver{T}) where T<:Abstr
     live_plotting(prob,solver)
 
     J_prev = cost(prob.obj, prob.X, prob.U, get_dt_traj(prob))
-    record_iteration!(prob, solver, J_prev, Inf)
+
 
     with_logger(logger) do
+        record_iteration!(prob, solver, J_prev, Inf)
         for i = 1:solver.opts.iterations
             J = step!(prob, solver, J_prev)
 
@@ -25,26 +27,29 @@ function solve!(prob::Problem{T,Discrete}, solver::iLQRSolver{T}) where T<:Abstr
                 return solver
             end
 
-            copyto!(prob.X, solver.X̄)
-            copyto!(prob.U, solver.Ū)
+            @timeit to "copy" begin
+                copyto!(prob.X, solver.X̄)
+                copyto!(prob.U, solver.Ū)
+            end
 
             dJ = abs(J - J_prev)
             J_prev = copy(J)
             record_iteration!(prob, solver, J, dJ)
             live_plotting(prob,solver)
 
             println(logger, InnerLoop)
-            evaluate_convergence(solver) ? break : nothing
+            @timeit to "convergence" evaluate_convergence(solver) ? break : nothing
         end
     end
     return solver
 end
 
 function step!(prob::Problem{T}, solver::iLQRSolver{T}, J::T) where T
-    jacobian!(prob,solver)
-    cost_expansion!(prob,solver)
-    ΔV = backwardpass!(prob,solver)
-    forwardpass!(prob,solver,ΔV,J)
+    to = solver.stats[:timer]
+    @timeit to "jacobians" jacobian!(prob,solver)
+    @timeit to "cost expansion" cost_expansion!(prob,solver)
+    @timeit to "backward pass" ΔV = backwardpass!(prob,solver)
+    @timeit to "forward pass" forwardpass!(prob,solver,ΔV,J)
 end
 
 function cost_expansion!(prob::Problem{T},solver::iLQRSolver{T}) where T

diff --git a/src/solvers/ilqr/ilqr_solver.jl b/src/solvers/ilqr/ilqr_solver.jl
@@ -117,7 +117,7 @@ end
 
 function AbstractSolver(prob::Problem{T,D}, opts::iLQRSolverOptions{T}) where {T<:AbstractFloat,D<:DynamicsType}
     # Init solver statistics
-    stats = Dict{Symbol,Any}()
+    stats = Dict{Symbol,Any}(:timer=>TimerOutput())
 
     # Init solver results
     n = prob.model.n; m = prob.model.m; N = prob.N