-
-
Notifications
You must be signed in to change notification settings - Fork 303
/
lqcontrol.jl
318 lines (241 loc) · 8.82 KB
/
lqcontrol.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
#=
Provides a type called LQ for solving linear quadratic control
problems.
@author : Spencer Lyon <[email protected]>
@author : Zac Cranko <[email protected]>
@date : 2014-07-05
References
----------
https://lectures.quantecon.org/jl/lqcontrol.html
=#
@doc doc"""
Linear quadratic optimal control of either infinite or finite horizon
The infinite horizon problem can be written
```math
\min \mathbb{E} \sum_{t=0}^{\infty} \beta^t r(x_t, u_t)
```
with
```math
r(x_t, u_t) := x_t' R x_t + u_t' Q u_t + 2 u_t' N x_t
```
The finite horizon form is
```math
\min \mathbb{E} \sum_{t=0}^{T-1} \beta^t r(x_t, u_t) + \beta^T x_T' R_f x_T
```
Both are minimized subject to the law of motion
```math
x_{t+1} = A x_t + B u_t + C w_{t+1}
```
Here ``x`` is n x 1, ``u`` is k x 1, ``w`` is j x 1 and the matrices are conformable for
these dimensions. The sequence ``{w_t}`` is assumed to be white noise, with zero
mean and ``\mathbb{E} w_t w_t' = I``, the j x j identity.
For this model, the time ``t`` value (i.e., cost-to-go) function ``V_t`` takes the form
```math
x' P_T x + d_T
```
and the optimal policy is of the form ``u_T = -F_T x_T``. In the infinite horizon
case, ``V, P, d`` and ``F`` are all stationary.
##### Fields
- `Q::ScalarOrArray` : `k x k` payoff coefficient for control variable u. Must be
symmetric and nonnegative definite
- `R::ScalarOrArray` : `n x n` payoff coefficient matrix for state variable x.
Must be symmetric and nonnegative definite
- `A::ScalarOrArray` : `n x n` coefficient on state in state transition
- `B::ScalarOrArray` : `n x k` coefficient on control in state transition
- `C::ScalarOrArray` : `n x j` coefficient on random shock in state transition
- `N::ScalarOrArray` : `k x n` cross product in payoff equation
- `bet::Real` : Discount factor in `[0, 1]`
- `capT::Union{Int, Void}` : Terminal period in finite horizon problem
- `rf::ScalarOrArray` : `n x n` terminal payoff in finite horizon problem. Must be symmetric and nonnegative definite
- `P::ScalarOrArray` : `n x n` matrix in value function representation ``V(x) = x'Px + d``
- `d::Real` : Constant in value function representation
- `F::ScalarOrArray` : Policy rule that specifies optimal control in each period
"""
mutable struct LQ
Q::ScalarOrArray
R::ScalarOrArray
A::ScalarOrArray
B::ScalarOrArray
C::ScalarOrArray
N::ScalarOrArray
bet::Real
capT::Union{Int, Nothing} # terminal period
rf::ScalarOrArray
P::ScalarOrArray
d::Real
F::ScalarOrArray # policy rule
end
"""
Main constructor for LQ type
Specifies default argumets for all fields not part of the payoff function or
transition equation.
##### Arguments
- `Q::ScalarOrArray` : `k x k` payoff coefficient for control variable u. Must be
symmetric and nonnegative definite
- `R::ScalarOrArray` : `n x n` payoff coefficient matrix for state variable x.
Must be symmetric and nonnegative definite
- `A::ScalarOrArray` : `n x n` coefficient on state in state transition
- `B::ScalarOrArray` : `n x k` coefficient on control in state transition
- `;C::ScalarOrArray{zero(size(R}(1)))` : `n x j` coefficient on random shock in
state transition
- `;N::ScalarOrArray{zero(size(B,1)}(size(A, 2)))` : `k x n` cross product in
payoff equation
- `;bet::Real(1.0)` : Discount factor in `[0, 1]`
- `capT::Union{Int, Void}(Void)` : Terminal period in finite horizon
problem
- `rf::ScalarOrArray{fill(NaN}(size(R)...))` : `n x n` terminal payoff in finite
horizon problem. Must be symmetric and nonnegative definite.
"""
function LQ(Q::ScalarOrArray,
R::ScalarOrArray,
A::ScalarOrArray,
B::ScalarOrArray,
C::ScalarOrArray=fill(zero(eltype(R)), size(R, 1)),
N::ScalarOrArray=zero(B'A);
bet::ScalarOrArray=1.0,
capT::Union{Int,Nothing}=nothing,
rf::ScalarOrArray=fill(NaN, size(R)...))
k = size(Q, 1)
n = size(R, 1)
F = k==n==1 ? zero(Float64) : fill(zero(Float64), k, n)
P = copy(rf)
d = 0.0
LQ(Q, R, A, B, C, N, bet, capT, rf, P, d, F)
end
@doc doc"""
Update `P` and `d` from the value function representation in finite horizon case
##### Arguments
- `lq::LQ` : instance of `LQ` type
##### Returns
- `P::ScalarOrArray` : `n x n` matrix in value function representation
``V(x) = x'Px + d``
- `d::Real` : Constant in value function representation
##### Notes
This function updates the `P` and `d` fields on the `lq` instance in addition to
returning them
"""
function update_values!(lq::LQ)
# Simplify notation
Q, R, A, B, N, C, P, d = lq.Q, lq.R, lq.A, lq.B, lq.N, lq.C, lq.P, lq.d
# Some useful matrices
s1 = Q + lq.bet * (B'P*B)
s2 = lq.bet * (B'P*A) + N
s3 = lq.bet * (A'P*A)
# Compute F as (Q + B'PB)^{-1} (beta B'PA)
lq.F = s1 \ s2
# Shift P back in time one step
new_P = R - s2'lq.F + s3
# Recalling that tr(AB) = tr(BA)
new_d = lq.bet * (d + tr(P * C * C'))
# Set new state
lq.P, lq.d = new_P, new_d
end
@doc doc"""
Computes value and policy functions in infinite horizon model.
##### Arguments
- `lq::LQ` : instance of `LQ` type
##### Returns
- `P::ScalarOrArray` : n x n matrix in value function representation ``V(x) = x'Px + d``
- `d::Real` : Constant in value function representation
- `F::ScalarOrArray` : Policy rule that specifies optimal control in each period
##### Notes
This function updates the `P`, `d`, and `F` fields on the `lq` instance in
addition to returning them
"""
function stationary_values!(lq::LQ)
# simplify notation
Q, R, A, B, N, C = lq.Q, lq.R, lq.A, lq.B, lq.N, lq.C
# solve Riccati equation, obtain P
A0, B0 = sqrt(lq.bet) * A, sqrt(lq.bet) * B
P = solve_discrete_riccati(A0, B0, R, Q, N)
# Compute F
s1 = Q .+ lq.bet * (B' * P * B)
s2 = lq.bet * (B' * P * A) .+ N
F = s1 \ s2
# Compute d
d = lq.bet * tr(P * C * C') / (1 - lq.bet)
# Bind states
lq.P, lq.F, lq.d = P, F, d
end
"""
Non-mutating routine for solving for `P`, `d`, and `F` in infinite horizon model
See docstring for `stationary_values!` for more explanation
"""
function stationary_values(lq::LQ)
_lq = LQ(copy(lq.Q),
copy(lq.R),
copy(lq.A),
copy(lq.B),
copy(lq.C),
copy(lq.N),
bet=copy(lq.bet),
capT=lq.capT,
rf=copy(lq.rf))
stationary_values!(_lq)
return _lq.P, _lq.F, _lq.d
end
"""
Private method implementing `compute_sequence` when state is a scalar
"""
function _compute_sequence(lq::LQ, x0::T, policies) where T
capT = length(policies)
x_path = Vector{T}(undef, capT+1)
u_path = Vector{T}(undef, capT)
x_path[1] = x0
u_path[1] = -(first(policies)*x0)
w_path = lq.C * randn(capT+1)
for t = 2:capT
f = policies[t]
x_path[t] = lq.A*x_path[t-1] + lq.B*u_path[t-1] + w_path[t]
u_path[t] = -(f*x_path[t])
end
x_path[end] = lq.A*x_path[capT] + lq.B*u_path[capT] + w_path[end]
x_path, u_path, w_path
end
"""
Private method implementing `compute_sequence` when state is a scalar
"""
function _compute_sequence(lq::LQ, x0::Vector{T}, policies) where T
# Ensure correct dimensionality
n, j, k = size(lq.C, 1), size(lq.C, 2), size(lq.B, 2)
capT = length(policies)
A, B, C = lq.A, reshape(lq.B, n, k), reshape(lq.C, n, j)
x_path = Matrix{T}(undef, n, capT+1)
u_path = Matrix{T}(undef, k, capT)
w_path = C*randn(j, capT+1)
x_path[:, 1] = x0
u_path[:, 1] = -(first(policies)*x0)
for t = 2:capT
f = policies[t]
x_path[:, t] = A*x_path[: ,t-1] + B*u_path[:, t-1] + w_path[:, t]
u_path[:, t] = -(f*x_path[:, t])
end
x_path[:, end] = A*x_path[:, capT] + B*u_path[:, capT] + w_path[:, end]
x_path, u_path, w_path
end
@doc doc"""
Compute and return the optimal state and control sequence, assuming innovation ``N(0,1)``
##### Arguments
- `lq::LQ` : instance of `LQ` type
- `x0::ScalarOrArray`: initial state
- `ts_length::Integer(100)` : maximum number of periods for which to return process. If `lq` instance is finite horizon type, the sequenes are returned only for `min(ts_length, lq.capT)`
##### Returns
- `x_path::Matrix{Float64}` : An `n x T+1` matrix, where the t-th column represents ``x_t``
- `u_path::Matrix{Float64}` : A `k x T` matrix, where the t-th column represents ``u_t``
- `w_path::Matrix{Float64}` : A `n x T+1` matrix, where the t-th column represents `lq.C*N(0,1)`
"""
function compute_sequence(lq::LQ, x0::ScalarOrArray, ts_length::Integer=100)
# Compute and record the sequence of policies
if isa(lq.capT, Nothing)
stationary_values!(lq)
policies = fill(lq.F, ts_length)
else
capT = min(ts_length, lq.capT)
policies = Vector{typeof(lq.F)}(undef, capT)
for t = capT:-1:1
update_values!(lq)
policies[t] = lq.F
end
end
_compute_sequence(lq, x0, policies)
end