-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathproject3.jl
4241 lines (3484 loc) · 144 KB
/
project3.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
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
### A Pluto.jl notebook ###
# v0.20.4
using Markdown
using InteractiveUtils
# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
#! format: off
quote
local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
local el = $(esc(element))
global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
el
end
#! format: on
end
# ╔═╡ 14964632-98d8-4a2f-b2f6-e3f28b558803
# ╠═╡ show_logs = false
using StanfordAA228V
# ╔═╡ 173388ab-207a-42a6-b364-b2c1cb335f6b
# ╠═╡ show_logs = false
begin
import MarkdownLiteral: @mdx
using ProgressLogging
using Downloads
using TOML
using Test
using Base64
using PlutoUI
using Distributions
using Random
using Plots
using ForwardDiff
import ForwardDiff: gradient, hessian, jacobian
using IntervalArithmetic
using Parameters
using BSON
using GridInterpolations
using LinearAlgebra
using LazySets
using JuMP
using SCS
using Flux
default(fontfamily="Computer Modern", framestyle=:box) # LaTeX-style plotting
md"> _Additional package management._"
end
# ╔═╡ 117d0059-ce1a-497e-8667-a0c2ef20c632
md"""
# Project 3: Estimate reachable sets
_Please wait until the entire notebook is finished loading before proceeding (you may get temporary errors)._
"""
# ╔═╡ d7643abe-4619-4859-b2e3-9e932fe53b2f
highlight(md"""_See the three **"⟶ Task"** sections below for where to fill out the algorithms._""")
# ╔═╡ 78181077-5548-459d-970d-1d8a9d63b72c
# ╔═╡ da5b4000-0bce-4fc2-be85-dada21264ca3
textbook_details([
"Chapter 8. _Reachability for Linear Systems_",
"Chapter 9. _Reachability for Nonlinear Systems_",
])
# ╔═╡ 0456a732-2672-4108-a241-db9ae879a913
# ╔═╡ 6e8ab7c9-fb49-4d89-946d-c7d7588c199a
md"""
## Julia/Pluto tips
Useful tips you may be interested in regarding Julia, Pluto, and plotting.
[](https://sisl.github.io/AA228VLectureNotebooks/media/html/julia_pluto_session.html)
[](https://sisl.github.io/AA228VLectureNotebooks/media/html/julia_plotting.html)
"""
# ╔═╡ a21612a1-1092-4892-9132-629833e7c867
# ╔═╡ ec776b30-6a30-4643-a22c-e071a365d50b
md"""
## Hints and FAQs
Expand the sections below for some helpful hints or refer to the FAQs.
[](https://github.com/sisl/AA228V-FAQs)
"""
# ╔═╡ 5aafffc6-8e3b-4591-99f6-6ee3b8082786
html_expand("Check out the <code>LazySets.jl</code> docs.", md"""
The documentation for the `LazySets` package is an excellent resource:
- [Latest documentation](https://juliareach.github.io/LazySets.jl/dev/)
- [`overapproximate`](https://juliareach.github.io/LazySets.jl/dev/lib/approximations/overapproximate/): Specifically the method that as an optional `[ε]::Real` error tolerance.
- [`box_approximation`](https://juliareach.github.io/LazySets.jl/dev/lib/approximations/box_approximation/)
""")
# ╔═╡ bd8cfc1d-cf0e-42e2-8e46-c763f142e8f3
html_expand("Stuck on the nonlinear system (medium)? Expand for hints on what to try.", hint(Markdown.parse("""Check out the _natural inclusion_ techniques in section 9.2 of the textbook.
We do _not_ restrict the number of vertices of your sets. So you can try other approaches detailed in chapter 9 _Reachability for Nonlinear System_ in the textbook.""")))
# ╔═╡ 5d389282-466d-47d8-ae69-0794cf620f27
html_expand("Stuck on the nonlinear neural network system (large)? Expand for hints on what to try.", hint(Markdown.parse("""Check out the _natural inclusion_ techniques in section 9.2 of the textbook.
We do _not_ restrict the number of vertices of your sets.
But if you're not satisfied with the (overly approximated) sets from natural inclusion, then check out the section on using `NeuralVerification.jl` ([link](#neural-verification)).
Note that algorithms like _Taylor inclusion_ and _conservative linearization_ are **unsound** for the nonlinear neural network because the ReLU activation functions are not continuous and differentiable (see details in section 9.7 of the textbook).""")))
# ╔═╡ 6bad6e8b-c021-41d2-afbb-bcd0242138dd
# ╔═╡ 1167fdfb-b097-4d18-b982-b1786794f8cf
html_space()
# ╔═╡ 55521bf3-eb1d-4ea1-9c18-aa8ba4461bff
html_expand("Expand for <code>SmallSystem</code> code.", md"""
```julia
## Agent
struct ProportionalController <: Agent
k
end
(c::ProportionalController)(s, a=missing) = c.k' * s
Πo(agent::ProportionalController) = agent.k'
## Environment
@with_kw struct MassSpringDamper <: Environment
m = 1.0
k = 10.0
c = 2.0
dt = 0.05
end
function (env::MassSpringDamper)(s, a)
return Ts(env) * s + Ta(env) * a
end
Ts(env::MassSpringDamper) = [1 env.dt; -env.k*env.dt/env.m 1-env.c*env.dt/env.m]
Ta(env::MassSpringDamper) = [0 env.dt/env.m]'
Ps(env::MassSpringDamper) = Product([Uniform(-0.2, 0.2), Uniform(-1e-12, 1e-12)])
𝒮₁(env::MassSpringDamper) = Hyperrectangle(low=[-0.2, 0.0], high=[0.2, 0.0])
## Sensor
struct AdditiveNoiseSensor <: Sensor
Do
end
(sensor::AdditiveNoiseSensor)(s) = sensor(s, rand(Do(sensor, s)))
(sensor::AdditiveNoiseSensor)(s, x) = s + x
Do(sensor::AdditiveNoiseSensor, s) = sensor.Do
Os(sensor::AdditiveNoiseSensor) = I
```
""")
# ╔═╡ 17fa8557-9656-4347-9d44-213fd3b635a6
Markdown.parse("""
## Small system
The system is comprised of an `agent`, environment (`env`), and `sensor`.
""")
# ╔═╡ e93d0515-c7d4-4af6-9041-fda738af5caa
begin
local ϵ = 0.5
sys_small = System(
ProportionalController([0.0, -1.0]),
MassSpringDamper(),
AdditiveNoiseSensor(Product([Uniform(-ϵ, ϵ), Uniform(-ϵ, ϵ)]))
)
end;
# ╔═╡ 6f3e24de-094c-49dc-b892-6721b3cc54ed
SmallSystem::Type = typeof(sys_small) # Type used for multiple dispatch
# ╔═╡ 592e4e77-8104-4464-8e10-ee2834c7c0ab
Markdown.parse("""
## Small specification \$\\psi\$
The specification \$\\psi\$ (written `\\psi<TAB>` in code) indicates what the system should do:
\$\$\\psi(\\tau) = \\square\\left( |p| < 0.3 \\right)\$\$
i.e., "the state position \$p\$ in the trajectory \$\\tau\$ should _always_ (\$\\square\$) be less than \$0.3\$, anything else is a failure."
We use the `AvoidSetSpecification` to define a _failure set_ where we want to ensure that the reachable set \$\\mathcal{R}_{1:d}\$ does not intersect with its complement:
\$\$\\mathcal{R}_{1:d} \\cap \\neg\\psi = \\varnothing\$\$
""")
# ╔═╡ ab4c6807-5b4e-4688-b794-159e26a1599b
ψ_small = AvoidSetSpecification(
HalfSpace([1.0, 0.0], -0.3) ∪ HalfSpace([-1.0, 0.0], -0.3));
# ╔═╡ 402c0eaa-727f-4c54-89ec-64c3dfb8002c
fbaseline(sys,ψ,seeds) =
aggregate_performance(estimate_probability_baseline, sys, ψ; seeds);
# ╔═╡ 92f20cc7-8bc0-4aea-8c70-b0f759748fbf
Markdown.parse("""
## ⟶ **Task (Small)**: Estimate the reachable set
Please fill in the following `estimate_reachable_sets` function.
""")
# ╔═╡ a003beb6-6235-455c-943a-e381acd00c0e
start_code()
# ╔═╡ c494bb97-14ef-408c-9de1-ecabe221eea6
end_code()
# ╔═╡ 8082ce45-6e93-4d98-8f90-79935deadec8
# ╔═╡ 38f3d8cf-21cf-4c77-bd45-a618b9b2e1cd
highlight(md"""
We recommend adding `@progress` to for loops in the `reachable` algorithms, e.g.:
```julia
@progress for d in 1:alg.h
# ...
end
```
This way you get incremental feedback that the algorithms are making progress.
""")
# ╔═╡ e2418154-4471-406f-b900-97905f5d2f59
html_quarter_space()
# ╔═╡ 1789c8b5-b314-4aba-ad44-555be9a85984
md"""
# 📊 Small Test
We'll automatically test your `estimate_reachable_sets(::SmallSystem, ψ)` function below.
"""
# ╔═╡ 97ddd327-94b8-4c2e-a79d-f5304725c25b
md"""
## Passing test criteria
Find a sequence of sets over time $t$ to depth $d$ that:
1. Over approximates the optimal reachable set using max $4$ vertices per time step.
2. The sets must not intersect with the failure region.
3. The sets must not be _under approximations_ of the optimal reachable sets.
"""
# ╔═╡ fda151a1-5069-44a8-baa1-d7903bc89797
html_space()
# ╔═╡ 6ec8a963-726d-4738-9a82-7e0b26b90b16
html_expand("Expand for <code>MediumSystem</code> code.", md"""
```julia
## Agent
struct ProportionalController <: Agent
k
end
(c::ProportionalController)(s, a=missing) = c.k' * s
Πo(agent::ProportionalController) = agent.k'
## Environment
@with_kw struct InvertedPendulum <: Environment
m::Float64 = 1.0
l::Float64 = 1.0
g::Float64 = 10.0
dt::Float64 = 0.05
ω_max::Float64 = 8.0
a_max::Float64 = 2.0
end
function (env::InvertedPendulum)(s, a, xs=missing)
θ, ω = s[1], s[2]
dt, g, m, l = env.dt, env.g, env.m, env.l
ω = ω + (3g / (2 * l) * sin(θ) + 3 * a / (m * l^2)) * dt # No `clamp`
θ = θ + ω * dt
return [θ, ω]
end
# Initial state distribution
Ps(env::InvertedPendulum) = Product([Uniform(-π/16, π/16), Uniform(-1.0, 1.0)])
## Sensor
struct AdditiveNoiseSensor <: Sensor
Do
end
(sensor::AdditiveNoiseSensor)(s) = sensor(s, rand(Do(sensor, s)))
(sensor::AdditiveNoiseSensor)(s, x) = s + x
Do(sensor::AdditiveNoiseSensor, s) = sensor.Do
Os(sensor::AdditiveNoiseSensor) = I
```
""")
# ╔═╡ 86f92b1d-87eb-40b0-ad0b-6b888c5fb9cc
md"""
## Medium system
First we define a convenient struct that holds the parameters of the medium system.
$(highlight(md"May be useful when defining functions such as `sets` and `intervals`.
Please do not change these values."))
"""
# ╔═╡ 0ab3fe4f-a13d-4d92-b3e4-5653f05dafe7
@with_kw struct PendulumParameters
disturbance_mag = 0.01
θmin = -π/16
θmax = π/16
ωmin = -1.0
ωmax = 1.0
end
# ╔═╡ 613443b1-ac8b-4570-b4c7-d107d63a36cd
md"""
An inverted pendulum is comprised of a `ProportionalController` with an `AdditiveNoiseSensor`.
"""
# ╔═╡ 99ccbce1-ac0b-4744-9c26-cb951490f482
md"""
We redefine the transition function as follows and remove clamping that was in Projects 1 and 2.
"""
# ╔═╡ 022f242e-f839-4b6a-b6ff-6ad3b09470a4
function (env::InvertedPendulum)(s, a, xs=missing)
θ, ω = s[1], s[2]
dt, g, m, l = env.dt, env.g, env.m, env.l
ω = ω + (3g / (2 * l) * sin(θ) + 3 * a / (m * l^2)) * dt # No `clamp`
θ = θ + ω * dt
return [θ, ω]
end
# ╔═╡ 77637b5e-e3ce-4ecd-90fc-95611af18002
begin
medium_params = PendulumParameters()
local disturbance_mag = medium_params.disturbance_mag
local θmin, θmax = medium_params.θmin, medium_params.θmax
local ωmin, ωmax = medium_params.ωmin, medium_params.ωmax
local agent = ProportionalController([-15.0, -8.0])
local env = InvertedPendulum()
# Different than Projects 1 and 2 (to be bounded)
local sensor = AdditiveNoiseSensor(Product([
Uniform(-disturbance_mag, disturbance_mag),
Uniform(-disturbance_mag, disturbance_mag)]))
# Different than Projects 1 and 2 (to be bounded)
StanfordAA228V.Ps(env::InvertedPendulum) =
Product([Uniform(θmin, θmax), Uniform(ωmin, ωmax)])
sys_medium = System(agent, env, sensor)
end;
# ╔═╡ c4c0328d-8cb3-41d5-9740-0197cbf760c2
MediumSystem::Type = typeof(sys_medium) # Type used for multiple dispatch
# ╔═╡ 8c78529c-1e00-472c-bb76-d984b37235ab
begin
# Different than Projects 1 and 2
StanfordAA228V.get_depth(sys::MediumSystem) = 21
Markdown.MD(
md"""
# 2️⃣ **Medium**: Inverted Pendulum
The medium system is a swinging inverted pendulum.
- It uses a proportional controller to keep it upright.
- The state is comprised of the angle $\theta$ and angular velocity $\omega$ making $s = [\theta, \omega]$
- Actions are left/right adjustments in the range $[-2, 2]$
- Disturbances $x$ are treated as additive noise: $x \sim \mathcal{N}(\mathbf{0}, 0.1^2I)$
""",
depth_highlight(sys_medium)
)
end
# ╔═╡ fd8c851a-3a42-41c5-b0fd-a12085543c9b
Markdown.MD(
Markdown.parse("""
# 1️⃣ **Small**: Mass-spring-damper
The small system is a simple _mass-spring-damper_. The mass-spring-damper system consists of a mass \$m\$ attached to a wall by a spring with spring constant \$k\$ and a damper with damping coefficient \$c\$.
- The state \$s\$ is the position (relative to the resting point) \$p\$ and velocity \$v\$ of the mass, \$s = [p,v]\$
- The action is the force \$\\beta\$ applies to the mass.
- The rollout depth is \$d=$(get_depth(sys_small))\$.
- Disturbances are a noisy measurement of the state with uniform noise with \$\\epsilon = 0.5\$:
\$\$\\begin{align}
x_p \\sim \\mathcal{U}(-\\epsilon, \\epsilon) \\\\
x_v \\sim \\mathcal{U}(-\\epsilon, \\epsilon)
\\end{align}\$\$
_See Appendix A.4 for more system details._
"""),
depth_highlight(sys_small)
)
# ╔═╡ 6a7b4f2b-187a-40a3-842f-126d220d40ed
md""" $t =$ $(@bind t_small Slider(1:get_depth(sys_small), show_value=true))
**Change the time $t$ slider to see the reachable set per time step.**"""
# ╔═╡ b1e9bd40-a401-4630-9a1f-d61b276e72f7
md"""
## Medium specification $\psi$
The inverted pendulum specification $\psi$ indicates what the system should do:
$$\psi(\tau) = \square\big(|\theta| < \pi/4\big)$$
i.e., "the absolute value of the pendulum angle $\theta$ (first element of the state $s$) in the trajectory $\tau$ should _always_ ($\square$) be less than $\pi/4$, anything else is a failure."
"""
# ╔═╡ fe272c1b-421c-49de-a513-80c7bcefdd9b
ψ_medium = LTLSpecification(@formula □(s -> abs(s[1]) < π / 4));
# ╔═╡ bac5c489-553c-436f-b332-8a8e97126a51
html_quarter_space()
# ╔═╡ 1da9695f-b7fc-46eb-9ef9-12160246018d
Markdown.parse("""
## ⟶ **Task (Medium)**: Estimate the reachable set
Please fill in the following `estimate_reachable_sets` function.
""")
# ╔═╡ 0606d827-9c70-4a79-afa7-14fb6b806546
start_code()
# ╔═╡ cb7b9b9f-59da-4851-ab13-c451c26117df
@medium function estimate_reachable_sets(sys::MediumSystem, ψ)
# TODO: WRITE YOUR CODE HERE
end
# ╔═╡ 759534ca-b40b-4824-b7ec-3a5c06cbd23e
end_code()
# ╔═╡ 97dbe1e4-8045-4213-866f-6921c733fbeb
# ╔═╡ b1cf81ad-e5cb-40d7-b365-abda3fc67a88
highlight(md"""
We recommend adding `@progress` to for loops in the `reachable` algorithms, e.g.:
```julia
@progress for d in 1:alg.h
# ...
end
```
This way you get incremental feedback that the algorithms are making progress.
""")
# ╔═╡ a228c1ac-62cb-4c18-89b9-bda4c3b1c5bb
html_quarter_space()
# ╔═╡ da2d692a-8378-435e-bd6b-c0e65caef542
md"""
# 📊 Medium Test
We'll automatically test your `estimate_reachable_sets(::MediumSystem, ψ)` function below.
"""
# ╔═╡ 50199170-f120-48ee-879c-bbdc11618f1e
md""" $t =$ $(@bind t_medium Slider(1:get_depth(sys_medium), show_value=true))
**Change the time $t$ slider to see the reachable set per time step.**"""
# ╔═╡ 428575d2-66a9-4a19-9eb0-cdf9b55277f7
md"""
## Passing test criteria
Find a sequence of sets over time $t$ to depth $d$ that:
1. Over approximates the reachable set (no limit on max vertices).
2. The sets _may_ intersect with the failure region.
- Although, for better leaderboard scores you may want to avoid this.
3. The sets must not be _under approximations_ of the approximately optimal reachable sets.
"""
# ╔═╡ 60ab8107-db65-4fb6-aeea-d4978aed77bd
html_space()
# ╔═╡ f8ea2983-c2d0-40ea-b949-9fc478ea45f8
Markdown.parse("""_The figure above shows several states_ \$s_t\$ _sampled from trajectories and a single trajectory shown as a line._""")
# ╔═╡ 5aae63b2-54ec-421b-84ec-0d4bc9c00c10
html_expand("Expand for <code>LargeSystem</code> code (surrogate).", md"""
```julia
## Agent
struct NoAgent <: Agent end
(c::NoAgent)(s, a=missing) = nothing
Distributions.pdf(c::NoAgent, s, x) = 1.0
## Environment
@with_kw struct ContinuumWorldSurrogate <: Environment
cw::ContinuumWorld = ContinuumWorld()
model::Chain
disturbance_mag = 0.1
end
(env::ContinuumWorldSurrogate)(s, a) = env(s, a, rand(Ds(env, s, a)))
(env::ContinuumWorldSurrogate)(s, a, x) = env.model(s) + x # Call neural network
Ps(env::ContinuumWorldSurrogate) = Ps(env.cw)
Ds(env::ContinuumWorldSurrogate, s, a) =
Product([
Uniform(-env.disturbance_mag, env.disturbance_mag),
Uniform(-env.disturbance_mag, env.disturbance_mag)])
## Sensor
struct IdealSensor <: Sensor end
(sensor::IdealSensor)(s) = s
(sensor::IdealSensor)(s, x) = sensor(s)
Distributions.pdf(sensor::IdealSensor, s, xₛ) = 1.0
```
""")
# ╔═╡ 790f1562-d6ff-4c44-b1ea-1b0ab1dced85
html_expand("Expand for <code>LargeSystem</code> code (original).", md"""
```julia
## Agent
struct InterpAgent <: Agent
grid::RectangleGrid
Q
end
(c::InterpAgent)(s) = argmax([interpolate(c.grid, q, s) for q in c.Q])
(c::InterpAgent)(s, x) = c(s)
Distributions.pdf(c::InterpAgent, o, xₐ) = 1.0
function load_cw_policy(filename::String)
res = BSON.load(filename)
grid = res[:grid]
Q = res[:Q]
return grid, Q
end
## Environment
@with_kw struct ContinuumWorld <: Environment
size = [10, 10] # dimensions
terminal_centers = [[4.5,4.5],[6.5,7.5]] # obstacle and goal centers
terminal_radii = [0.5, 0.5] # radius of obstacle and goal
directions = [[0,1],[0,-1],[-1,0],[1,0]] # up, down, left, right
end
(env::ContinuumWorld)(s, a) = env(s, a, rand(Ds(env, s, a)))
function (env::ContinuumWorld)(s, a, x)
dir = env.directions[a]
return s .+ dir
end
Ps(env::ContinuumWorld) = Product([Uniform(0,1), Uniform(0,1)])
Ds(env::ContinuumWorld, s, a) = Deterministic()
## Sensor
struct IdealSensor <: Sensor end
(sensor::IdealSensor)(s) = s
(sensor::IdealSensor)(s, x) = sensor(s)
Distributions.pdf(sensor::IdealSensor, s, xₛ) = 1.0
```
""")
# ╔═╡ b608b8b2-5166-419b-ac4e-18e87c93ac01
md"""
## Large system
First we define a convenient struct that holds the parameters of the large system.
$(highlight(md"May be useful when defining functions such as `sets` and `intervals`.
Please do not change these values."))
"""
# ╔═╡ 05892ece-4c38-4c39-841a-f8b9333ba3ef
@with_kw struct ContinuumWorldParameters
disturbance_mag = 0.1 # disturbance magnitude x = [𝒰(-mag, mag), 𝒰(-mag, mag)]
xmin = 0 # x initial state x minimum
xmax = 1 # x initial state x maximum
ymin = 0 # y initial state y minimum
ymax = 1 # y initial state y maximum
end
# ╔═╡ 7d054465-9f80-4dfb-9b5f-76c3977de7cd
md"""
The continuum world system uses the neural network surrogate to predict next states.
"""
# ╔═╡ 005966cb-763d-4346-8fdb-e336a9514e8c
md"""
### Neural network surrogate 🤖
We trained a 3 layer **ReLU-based** neural network to predict next states $\hat{s}'$ given the current state $s$.
"""
# ╔═╡ 94a5b84a-b942-4317-b255-f49cc13532c6
hint(md"""Since we trained a **ReLU-based** network, what algorithms are _not_ applicable (i.e., not **sound**)?
Think: _Is ReLU continuous and differentiable?_""", title="Hint: What can we apply?")
# ╔═╡ 6a38fe32-d862-4d7d-9978-290a152ae575
model = BSON.load("cw_model.bson")[:model]
# ╔═╡ 1ec68a39-8de9-4fd3-be8a-26cf7706d1d6
begin
large_params = ContinuumWorldParameters()
local agent = NoAgent()
local env = ContinuumWorldSurrogate(
cw=ContinuumWorld(),
model=model,
disturbance_mag=large_params.disturbance_mag,
)
local sensor = IdealSensor()
sys_large = System(agent, env, sensor)
end;
# ╔═╡ aa0c4ffc-d7f0-484e-a1e2-7f6f92a3a53d
Markdown.MD(
Markdown.parse("""
# 3️⃣ **Large**: Continuum World
The large system is an agent in the continuum world (continuous grid world).
The objective is to get to the _green_ goal region and avoid the _red_ failure region.
### (Original) Continuum World System
To train our surrogate neural network, we define the "original" continuum world system that uses a lookup-based policy **which we cannot easily apply reachability to**.
- States are the \$[x,y]\$ positions in a \$10\\times10\$ continuous world.
- With an initial state distribution uniformly in \$[0,1]\$ for both \$x\$ and \$y\$.
- Actions \$a\$ are the cardinal directions `up`, `down`, `left`, `right`.
- It uses an interpolated lookup-table policy.
- The transition dynamics apply the action \$a\$ to the current state.
- e.g., \$s' = s + a\$ where \$a \\in \\big\\{[0,1],\\, [0,-1],\\, [-1,0],\\, [1,0]\\big\\}\$
- We do not apply disturbances to the original environment.
### (Surrogate) Continuum World System
- Same state space and initial state distribution as the original system.
- The agent does not directly take actions, but instead will predict the next state \$\\hat{s}'\$ using the trained _**neural network surrogate**_.
- Disturbances \$x\$ are applied to the predicted next state \$\\hat{s}'\$ as "slip" noise. The agent slips in a random direction modeled by adding a random vector \$x = [x_x, x_y]\$ sampled from two independent uniform distributions:
\$\$
\\begin{align}
x_x &\\sim \\mathcal{U}(-$(large_params.disturbance_mag), $(large_params.disturbance_mag)) \\\\
x_y &\\sim \\mathcal{U}(-$(large_params.disturbance_mag), $(large_params.disturbance_mag))
\\end{align}\$\$
- Thus, the transition dynamics apply the disturbances to the predicted next state:
\$\$s' = \\hat{s}' + x\$\$
_Modified from the system details in Appendix A.7._"""),
depth_highlight(sys_large)
)
# ╔═╡ 29279d51-1162-4fa8-bdd5-f0ff0e4ef968
LargeSystem::Type = typeof(sys_large) # Type used for multiple dispatch
# ╔═╡ 9f739929-1cd3-4935-b229-ae3aeac7e131
begin
ThisProject = Project3
max_vertices(sys::SmallSystem) = 4
max_vertices(sys::MediumSystem) = Inf
max_vertices(sys::LargeSystem) = Inf
end;
# ╔═╡ 60f72d30-ab80-11ef-3c20-270dbcdf0cc4
Markdown.MD(
Markdown.parse("""
**Task**: Estimating the reachable sets of three systems.
- **Small system**: Mass-spring-damper (_**linear system**_).
- A maximum of \$n = $(format(max_vertices(sys_small); latex=true))\$ vertices per time step is allotted for the small system.
- **Medium system**: Swinging inverted pendulum (_**nonlinear system**_).
- No restriction on number of vertices.
- **Large system**: Continuum world (_**nonlinear neural network**_).
- No restriction on number of vertices.
"""),
highlight(md"""Your job is to write the following function that returns the estimated reachable sets for each time step from $1$ to depth $d$ (use `get_depth(sys)` to get `d`)."""),
md"""
```julia
estimate_reachable_sets(sys, ψ)
```
##### Note on the return type:
The return type of `estimate_reachable_sets` should either be:
- a [`UnionSet`](https://juliareach.github.io/LazySets.jl/dev/lib/lazy_operations/UnionSet) or [`UnionSetArray`](https://juliareach.github.io/LazySets.jl/dev/lib/lazy_operations/UnionSet/#def_UnionSetArray) (from `LazySets.jl`)
- **This is already the type that is output from the `reachable` algorithms in the textbook.**
- or a `Vector{<:LazySet}` (see [docs](https://juliareach.github.io/LazySets.jl/dev/lib/interfaces/LazySet/#LazySets.LazySet) for the `LazySet` subtypes)
- If you want to run your own custom reachability code, we also accept a `Vector` of `LazySet`s for each time step.
""")
# ╔═╡ c4fa9af9-1a79-43d7-9e8d-2854652a4ea2
html_expand("Stuck on the linear system (small)? Expand for hints on what to try.", hint(Markdown.parse("""Check out the _set propagation_ techniques in section 8.2 of the textbook.
Being a linear system, we can solve the reachability problem _**exactly**_. But we restrict the maximum number of vertices per time step to be \$n = $(format(max_vertices(sys_small); latex=true))\$ to make it more interesting.""")))
# ╔═╡ fc2d34da-258c-4460-a0a4-c70b072f91ca
@small function estimate_reachable_sets(sys::SmallSystem, ψ; n=max_vertices(sys))
# TODO: WRITE YOUR CODE HERE
end
# ╔═╡ 307afd9c-6dac-4a6d-89d7-4d8cabfe3fe5
Markdown.MD(
md"""
$(@bind rerun_small LargeCheckBox(text="⟵ Click to re-run the <code>SmallSystem</code> evaluation."))""",
Markdown.parse("""
↑ This will re-run **`estimate_reachable_sets(::SmallSystem, ψ)`** and re-save **`$(get_filename(sys_small, ThisProject))`**
_Uncheck this to load results from the file._
""")
)
# ╔═╡ 38f26afd-ffa5-48d6-90cc-e3ec189c2bf1
Markdown.MD(
md"""
$(@bind rerun_medium LargeCheckBox(text="⟵ Click to re-run the <code>MediumSystem</code> evaluation."))""",
Markdown.parse("""
↑ This will run **`estimate_reachable_sets(::MediumSystem, ψ)`** and save **`$(get_filename(sys_medium, ThisProject))`**
_Uncheck this to load results from the file._
""")
)
# ╔═╡ 4eeaa9ae-eac5-478a-aca5-82de3dda24f7
submission_details(@bind(directory_trigger, OpenDirectory(@__DIR__)), ThisProject,
[SmallSystem, MediumSystem, LargeSystem])
# ╔═╡ 0c520f93-49ce-45eb-899d-a31105d856c8
if directory_trigger
@info "Opening local directory..."
sleep(1)
end
# ╔═╡ 2a81d0a0-a609-4d20-bddb-e5cbe65caf4b
md"""
### Training code
If you're interested in the `Flux.jl` training code, expand this section below.
"""
# ╔═╡ 3ae14115-f7d6-495b-a3fa-d13d0b0cdd54
html_expand("Expand for the <code>Flux</code> training code.", md"""
You can use `CUDA` if your machine supports CUDA-based GPUs.
```julia
using Flux
using Metal # Training on Apple silicon
```
We define the "original" system to collect data from:
```julia
local grid, Q = load_cw_policy(joinpath(@__DIR__, "cw_policy.bson"))
cw_agent_orig = InterpAgent(grid, Q)
cw_env_orig = ContinuumWorld()
cw_sensor_orig = IdealSensor()
sys_large_orig = System(cw_agent_orig, cw_env_orig, cw_sensor_orig)
# Type used for multiple dispatch
LargeSystemOriginal::Type = typeof(sys_large_orig)
```
For data collection, we collect $(s, s')$ pairs which consist of $s = [x,y]$ and $s' = [x', y']$ state values (using the original continuum world system).
```julia
function data_collection(sys::LargeSystemOriginal; n=100_000, seed=4)
Random.seed!(seed)
𝒟 = []
d = get_depth(sys)
for i in 1:n
τ = rollout(sys; d)
𝐬 = [deepcopy(step.s) for step in τ[1:end-1]]
𝐬′ = [deepcopy(step.s) for step in τ[2:end]]
push!(𝒟, collect(zip(𝐬, 𝐬′))...)
end
return 𝒟
end
```
We train a simple 3 (hidden) layer ReLU-based neural network to predict next states $\hat{s}'$ given the current state $s$. We minimize the mean-squared error (MSE) loss function using the Adam optimizer with a learning rate of $\alpha = 5{e-}4$ for $100$ epochs.
```julia
function train(sys::LargeSystemOriginal; epochs=100, lr=5e-4, seed=4)
𝒟 = data_collection(sys; seed)
𝐬 = first.(𝒟)
𝐬′ = last.(𝒟)
inputs::Matrix{Float32} = hcat(first.(𝐬), last.(𝐬))'
targets::Matrix{Float32} = hcat(first.(𝐬′), last.(𝐬′))'
loader = Flux.DataLoader((inputs, targets), batchsize=length(𝐬)÷50)
Random.seed!(seed)
model = Chain(
Dense(2, 6, relu),
Dense(6, 6, relu),
Dense(6, 6, relu),
Dense(6, 2) # next state s′
)
model = gpu(model)
opt = Flux.setup(Flux.Adam(lr), model)
@withprogress for epoch in 1:epochs
for (s, s′) in loader
s, s′ = gpu(s), gpu(s′)
loss, grads = Flux.withgradient(model) do m
ŝ′ = m(s)
Flux.mse(ŝ′, s′)
end
Flux.update!(opt, model, grads[1])
@logprogress "[$epoch]\n$loss" epoch/epochs
end
end
model = cpu(model)
model = fmap(f64, model)
return model
end
```
Then we run `train` and save the model.
```julia
model = train(sys_large_orig)
BSON.@save "cw_model.bson" model
```
""")
# ╔═╡ be426908-3fee-4ecd-b054-2497ce9a2e50
md"""
## Large specification $\psi$
The continuum world specification $\psi$ indicates what the system should do:
$$\psi(\tau) = \lozenge(G) \wedge \square(\neg F)$$
where
$$\begin{gather}
G(s_t): \lVert s_t - \text{goal} \rVert_2 \le r_g \tag{$s_t$ is in the goal region} \\
F(s_t): \lVert s_t - \text{fail} \rVert_2 \le r_f \tag{$s_t$ is in the obstace region}
\end{gather}$$
where $r_g$ and $r_f$ are the goal region and obstacle region radii, respectively.
i.e., "the agent must _eventually_ ($\lozenge$) reach the goal and _always_ ($\square$) avoid the obstace region."
"""
# ╔═╡ 258e14c4-9a2d-4515-9a8f-8cd96f31a6ff
begin
local goal = sys_large.env.cw.terminal_centers[2]
local fail = sys_large.env.cw.terminal_centers[1]
local rg = sys_large.env.cw.terminal_radii[2]
local rf = sys_large.env.cw.terminal_radii[1]
local G = @formula s->norm(s .- goal) ≤ rg
local F = @formula s->norm(s .- fail) ≤ rf
ψ_large = LTLSpecification(@formula ◊(G) ∧ □(¬F))
end;
# ╔═╡ e3d6fdf1-3a9e-446b-8482-49d6f64b652e
html_quarter_space()
# ╔═╡ 23fd490a-74d2-44b4-8a12-ea1460d95f85
Markdown.parse("""
## ⟶ **Task (Large)**: Estimate the reachable set
Please fill in the following `estimate_reachable_sets` function.
""")
# ╔═╡ 18a70925-3c2a-4317-8bbc-c2a096ec56d0
start_code()
# ╔═╡ 3471a623-16af-481a-8f66-5bd1e7890188
@large function estimate_reachable_sets(sys::LargeSystem, ψ)
# TODO: WRITE YOUR CODE HERE
end
# ╔═╡ 4c5210d6-598f-4167-a6ee-93bceda7223b
end_code()
# ╔═╡ ef6bf2ba-748e-4f43-ad53-05d1936c2ba9
# ╔═╡ 0a496e93-5853-46bd-a3c5-6f466df90441
highlight(md"""
We recommend adding `@progress` to for loops in the `reachable` algorithms, e.g.:
```julia
@progress for d in 1:alg.h
# ...
end
```
This way you get incremental feedback that the algorithms are making progress.
""")
# ╔═╡ ba82b4e4-413c-4b78-a777-85d03e3554f4
html_quarter_space()
# ╔═╡ ba780732-7aaa-4e09-9c40-304aa62e564b
html"""
<h3 id='neural-verification'>(Optional) Not satisfied with your solution?</h3>
<p>This <code>import</code> may be a spoiler 🙃</p>
"""
# ╔═╡ 8cc343cc-f5be-4ff4-802f-adcb6a77674a
import NeuralVerification: Ai2, forward_network, network
# ╔═╡ b4a0cef4-25f6-4b5d-a739-6f9583d6fb3d
hint(md"""
Although a simple application of _natural inclusion_ will pass the autograder, you may not be satisfied with this solution for the large problem.
**Expand the section below for some hints on what else to try.**""")
# ╔═╡ 16502371-5139-40de-be3e-0926f55ce405
html_expand("Expand for a hints on a better solution to the large problem.", [
html"<h2hide><code>NeuralVerification.jl</code></h2hide>",
md"""
The `NeuralVerification.jl` package developed by our lab, SISL, has an extensive library of implemented state-of-the-art neural network verification algorithms.
- [Link to the documentation.](https://sisl.github.io/NeuralVerification.jl/latest/)
You can use this package to find a **sound** solution (unlike _Taylor inclusion_ and _conservative linearization_) and a better solution than _natural inclusion_.""",
html"<h2hide>Implement this!</h2hide>",
md"""
Copy the following function `nvreach` to a new cell and implement the comments `1-4` if you want to gain experience with state-of-the-art neural network verification!
```julia
function nvreach(sys::LargeSystem, d)
solver = Ai2() # We recommend using the AI² solver (see [1])
net = network(sys.env.model) # Convert Flux model → NeuralVerification network
𝒮, 𝒳 = sets(sys, 1) # Get initial state set 𝒮 and disturbance set 𝒳
ℛs = LazySet[𝒮] # Initialze reachable sets w/ initial state set 𝒮
@progress for t in 2:d
# TODO:
# 1. Forward pass 𝒮 through the net: forward_network(solver, net, 𝒮)
# 2. Then what? ... \oplus<TAB> may be useful...
# 3. Tip: concretize afterwards.
# 4. Push to ℛs
end
return UnionSetArray([ℛs...])
end
```
Then call this function within your `estimate_reachable_sets(sys::LargeSystem, ψ)` like so:
```julia
d = get_depth(sys)
return nvreach(sys, d)
```
**(Or you can put the internals of `nvreach` directly into `estimate_reachable_sets`)**
If you're curious how to type certain variables used in the `nvreach` function, see the table below.
| Variable | How to produce it |
| :------- | :--------------------------------------------- |
| `𝒮` | `\scrS<TAB>` |
| `𝒮′` | `\scrS<TAB>\prime<TAB>` |
| `𝒮′ ⊕ 𝒳` | `\scrS<TAB>\prime<TAB> \oplus<TAB> \scrX<TAB>` |
| `ℛs` | `\scrR<TAB>s` |
""",
html"<h2hide>AI² solver</h2hide>",
md"""
The `nvreach` function uses the AI² solver [^1], which is a state-of-the-art neural network verification algorithm that works on ReLU-based networks. For more technical details, please refer to the paper below.
- [Link to AI² paper.](https://www.cs.utexas.edu/~swarat/pubs/ai2.pdf)
- [Link to AI² `solve` functionality.](https://github.com/sisl/NeuralVerification.jl/blob/9802458c3055ecdae7ecdebd45f680f11af76c67/src/reachability/ai2.jl#L54)
[^1] T. Gehr, M. Mirman, D. Drashsler-Cohen, P. Tsankov, S. Chaudhuri, and M. Vechev, **"AI²: Safety and Robustness Certification of Neural Networks with Abstract Interpretation,"** in *2018 IEEE Symposium on Security and Privacy (SP)*, 2018.
""",
try LocalResource(joinpath(@__DIR__, "..", "media", "ai2.png")) catch end,
html"<h2hide>Wanna go further?</h2hide>",
hint(md"Maybe a little 🄿 🄰 🅁 🅃 🄸 🅃 🄸 🄾 🄽 🄸 🄽 🄶?"),
])
# ╔═╡ 5c490a85-cff5-46bb-a6fa-330c81d4cd3b
html_quarter_space()
# ╔═╡ ea2d7eb7-d576-415c-ac4c-fea7f90de637
md"""
# 📊 Large Test
We'll automatically test your `estimate_reachable_sets(::LargeSystem, ψ)` function below.
"""
# ╔═╡ 7fe1c3d7-469c-47d9-9d46-e5b8b263edb9
Markdown.MD(
md"""
$(@bind rerun_large LargeCheckBox(text="⟵ Click to re-run the <code>LargeSystem</code> evaluation."))""",
Markdown.parse("""
↑ This will re-run **`estimate_reachable_sets(::LargeSystem, ψ)`** and re-save **`$(get_filename(sys_large, ThisProject))`**
_Uncheck this to load results from the file._
""")
)