1# 2# This file is part of CasADi. 3# 4# CasADi -- A symbolic framework for dynamic optimization. 5# Copyright (C) 2010-2014 Joel Andersson, Joris Gillis, Moritz Diehl, 6# K.U. Leuven. All rights reserved. 7# Copyright (C) 2011-2014 Greg Horn 8# 9# CasADi is free software; you can redistribute it and/or 10# modify it under the terms of the GNU Lesser General Public 11# License as published by the Free Software Foundation; either 12# version 3 of the License, or (at your option) any later version. 13# 14# CasADi is distributed in the hope that it will be useful, 15# but WITHOUT ANY WARRANTY; without even the implied warranty of 16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU 17# Lesser General Public License for more details. 18# 19# You should have received a copy of the GNU Lesser General Public 20# License along with CasADi; if not, write to the Free Software 21# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA 22# 23# 24 25from pylab import * 26 27# End time 28T = 10. 29 30# Number of control intervals 31N = 20 32 33# Number of Runge-Kutta 4 steps per interval and step size 34NK = 20 35DT = T/(N*NK) 36 37# Number of discrete control values 38NU = 101 39 40# Number of discrete state values 41NX = 101 42 43# System dynamics, can be called with matricex 44def f(x1,x2,u): 45 x1_dot = (1 - x2*x2)*x1 - x2 + u 46 x2_dot = x1 47 q_dot = x1*x1 + x2*x2 + u*u 48 return (x1_dot, x2_dot, q_dot) 49 50# Control enumeration 51U = linspace(-1,1,NU) 52 53# State space enumeration 54x1 = linspace(-1,1,NX) 55x2 = linspace(-1,1,NX) 56X1,X2 = meshgrid(x1,x2) 57 58# For each control action and state, precalculate next state and stage cost 59stage_J = [] 60next_x1 = [] 61next_x2 = [] 62for u in U: 63 # Take number of integration steps 64 X1_k = copy(X1) 65 X2_k = copy(X2) 66 Q_k = zeros(X1.shape) 67 for k in range(NK): 68 # RK4 integration for x1, x2 and q 69 k1_x1, k1_x2, k1_q = f(X1_k, X2_k, u) 70 k2_x1, k2_x2, k2_q = f(X1_k + DT/2 * k1_x1, X2_k + DT/2 * k1_x2, u) 71 k3_x1, k3_x2, k3_q = f(X1_k + DT/2 * k2_x1, X2_k + DT/2 * k2_x2, u) 72 k4_x1, k4_x2, k4_q = f(X1_k + DT * k3_x1, X2_k + DT * k3_x2, u) 73 X1_k += DT/6*(k1_x1 + 2*k2_x1 + 2*k3_x1 + k4_x1) 74 X2_k += DT/6*(k1_x2 + 2*k2_x2 + 2*k3_x2 + k4_x2) 75 Q_k += DT/6*(k1_q + 2*k2_q + 2*k3_q + k4_q ) 76 77 # Find out which state comes next (index) 78 X1_k = matrix.round((X1_k+1)/2*(NX-1)).astype(int) 79 X2_k = matrix.round((X2_k+1)/2*(NX-1)).astype(int) 80 81 # Infinite cost if state gets out-of-bounds 82 I = X1_k < 0; Q_k[I]=inf; X1_k[I]=0 83 I = X2_k < 0; Q_k[I]=inf; X2_k[I]=0 84 I = X1_k >= NX; Q_k[I]=inf; X1_k[I]=0 85 I = X2_k >= NX; Q_k[I]=inf; X2_k[I]=0 86 87 # Save the stage cost and next state 88 next_x1.append(X1_k) 89 next_x2.append(X2_k) 90 stage_J.append(Q_k) 91 92# Calculate cost-to-go (no end cost) and optimal control 93J = zeros(X1.shape) 94U_opt = [] 95for k in reversed(list(range(N))): 96 # Cost to go for the previous step, optimal control action 97 J_prev = inf*ones(X1.shape) 98 u_prev = -ones(X1.shape,dtype=int) 99 100 # Test all control actions 101 for uind in range(NU): 102 J_prev_test = J[next_x2[uind],next_x1[uind]]+stage_J[uind] 103 better = J_prev_test<J_prev 104 u_prev[better] = uind 105 J_prev[better] = J_prev_test[better] 106 107 # Update cost-to-go and save optimal control 108 J = J_prev 109 U_opt.append(u_prev) 110 111# Reorder U_opt by stage 112U_opt.reverse() 113 114# Find optimal control starting at x1=0, x2=1 115i1 = NX//2 116i2 = NX-1 117u_opt = [] 118x1_opt = [x1[i1]] 119x2_opt = [x2[i2]] 120cost = 0 121for k in range(N): 122 # Get the optimal control and go to next step 123 u_ind = U_opt[k][i2,i1] 124 cost += stage_J[u_ind][i2,i1] 125 i1, i2 = next_x1[u_ind][i2,i1], next_x2[u_ind][i2,i1] 126 127 # Save the trajectories 128 u_opt.append(U[u_ind]) 129 x1_opt.append(x1[i1]) 130 x2_opt.append(x2[i2]) 131 132# Optimal cost 133print("Minimal cost: ", cost) 134assert abs(cost-J[NX-1,NX//2])<1e-8 # Consistency check 135 136# Plot 137figure(1) 138clf() 139 140# Plot optimal cost-to-go 141subplot(121) 142contourf(X1,X2,J) 143colorbar() 144xlabel('x1') 145ylabel('x2') 146title('Cost-to-go') 147 148subplot(122) 149plot(linspace(0,T,N+1),x1_opt,'--') 150plot(linspace(0,T,N+1),x2_opt,'-.') 151step(linspace(0,T,N),u_opt,'-') 152plt.title("Dynamic programming solution") 153plt.xlabel('time') 154plt.legend(['x1 trajectory','x2 trajectory','u trajectory']) 155grid(True) 156show() 157