-
Notifications
You must be signed in to change notification settings - Fork 0
/
flatquad_landing_experiment.py
executable file
·512 lines (371 loc) · 17.1 KB
/
flatquad_landing_experiment.py
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
#!/usr/bin/env python
import argparse
import time
from operator import itemgetter
import diffrax
import ipdb
import jax
import jax.numpy as np
import matplotlib.pyplot as pl
import meshcat
import meshcat.geometry as geom
import meshcat.transformations as tf
import numpy as onp
import tqdm
import wandb
import levelsets
import pontryagin_utils
import visualiser
from misc import *
# from jax import config
# config.update("jax_enable_x64", True)
# random helper functions {{{
# transform between "old" (= local coordinates) and "new" (= embedded in R^n)
# representation.
def old_to_new(x):
return np.concatenate([
x[0:2], # posx, posy
np.array([np.sin(x[2]), np.cos(x[2])]), # angle embedding
x[3:]
])
def new_to_old(x):
# verified experimentally: thetas = np.arctan2(np.sin(thetas), np.cos(thetas))
# because old_to_new is not globally invertible this inverts its restriction on
# the domain -pi/2 < theata < pi/2 or something like that.
return np.concatenate([
x[0:2], # posx, posy
np.array([np.arctan2(x[2], x[3])]),
x[4:]
])
# }}}
def define_problem_params():
# classic 2D quad type thing. 6D` state.
# update, 6D manifold embedded in R^7.
m = 20 # kg
g = 9.81 # m/s^2
r = 0.5 # m
I = m * (r/2)**2 # kg m^2 / radian (???)
umax = m * g * 1.2 / 2 # 20% above hover thrust
# remove time arguments sometime now that we're mostly treating
# infinite horizon, time-invariant problems?
def f(x, u):
# unpack for easier names
Fl, Fr = u
posx, posy, sinPhi, cosPhi, vx, vy, omega = x
# Phi' = omega
# d/dt sin(Phi) = cos(Phi) Phi' = cosPhi omega
# d/dt cos(Phi) = -sin(Phi) Phi' = -sinPhi omega
xdot = np.array([
vx,
vy,
cosPhi * omega,
-sinPhi * omega,
-sinPhi * (Fl + Fr) / m,
cosPhi * (Fl + Fr) / m - g,
(Fr-Fl) * r / I,
])
return xdot
x_eq = np.array([0, 0, 0, 1, 0, 0, 0], dtype=float)
def l(x, u):
Fl, Fr = u
posx, posy, sin_Phi, cos_Phi, vx, vy, omega = x
# penalise deviation from cos(Phi)=1, sin(Phi)=0 just in cartesian ambient space
# derivatives should be the same still (bc sin'(0) = 1)
state_length_scales = np.array([0.3, 0.3, np.deg2rad(30), np.deg2rad(30), .5, .5, np.deg2rad(120)])
Q = np.diag(1/state_length_scales**2)
state_cost = (x - x_eq).T @ Q @ (x - x_eq)
# can we just set an input penalty that is zero at hover?
# penalise x acc, y acc [m/s^2], angular acc [rad/s^2] here
# this here is basically a state-dependent linear map of the inputs, i.e. M(x) u with M(x) a 3x2 matrix.
# the overall input cost will be acc.T M(x).T Q M(x) acc, so for each state it is still a nice quadratic in u.
accelerations = np.array([
-sin_Phi * (Fl + Fr) / m,
cos_Phi * (Fl + Fr) / m - g,
(Fr - Fl) * r / I,
])
accelerations_lengthscale = np.array([1, 1, 1])
input_cost = accelerations.T @ np.diag(1/accelerations_lengthscale**2) @ accelerations
return state_cost + input_cost
problem_params = {
'system_name': 'flatquad',
# dynamics X x U -> TxX, stage cost X x U -> R
'f': f,
'l': l,
# state & input space dimensions
# if manifold, the dimension of the ambient space, not the manifold!
'nx': 7,
'nu': 2,
'state_names': ("x", "y", "sinPhi", "cosPhi", "vx", "vy", "omega"),
'u_eq': np.ones(2) * m * g / 2,
'x_eq': x_eq,
# if ever treating slightly bigger systems it would pay to frame this
# as a general convex polytope described by Ax <= b.
'U_interval': [np.zeros(2), umax*np.ones(2)],
# the value level below which we accept the LQR solution as correct.
'V_f': 0.001,
'V_max': 2000.,
# 'V_max': 50.,
# constraint equation defining the state space manifold as its 0-levelset.
# if R^n, set this to None
# number of constraint equations = codimension of manifold.
# atm only codimension 1 is supported, because this makes finding
# an orthonormal basis for the normal space trivial.
# in this case only the unit circle for angle parameterisation.
# / 2 so its jacobian is normalised.
'm': lambda x: (x[2]**2 + x[3]**2 - 1) / 2,
# projection operation onto the manifold -- great for resetting if
# we stray off the manifold due to numerical errors.
'project_M': lambda x: x.at[2:4].set(x[2:4] / np.linalg.norm(x[2:4])),
'x_extent': np.array([
20, 20, # x and y, [m]
1., 1., # sinPhi and cosPhi [1] (but irrelevant -- see sampling fct)
20, 20, # vx and vy, [m/s]
20*np.pi # omega [rad/s]
]),
}
return problem_params
def base_algo_params():
algo_params = {
# PRNG seed
'seed': 32,
# ODE SOLVER PARAMS
'pontryagin_solver_atol': 1e-4,
'pontryagin_solver_rtol': 1e-4,
'dtmin': 0.01,
'dtmax': 0.5,
# project back to manifold after each solver step. only possible if
# problem_params['project_M'] correctly defined.
'project_manifold': True,
# with throw=True we can set this pretty tight - it will just stop early.
# will have to make sure ourselves that this is not a problem
'pontryagin_solver_maxsteps': 128,
# not very relevant if we can just "resume" the trajectory in a later solve
# also maybe it makes sense to stop based on value, like stop after we reach sth like 10x
# the current value level? then we pervent spending lots of effort in "difficult" (=high l(x, u))
# state space regions.
'pontryagin_solver_T': 5.,
# (this was not used for a long time)
# in theory ||vxx|| can become infinite - meaning we solve an ODE with finite escape time.
# this happenn when many optimal trajectories originate from a small region (or a point in the limit)
# to avoid this we just stop calculating the trajectory once ||vxx|| exceeds this bound.
# hopefully the state space will still be sufficiently covered. In regions where ||vxx|| would
# have been very high we will just have to accept the interpolation instead.
'pontryagin_solver_vxx': False,
'vxx_max_norm': 1e4,
# causes it not to quit when hitting maxsteps. probably still all subsequent
# results will be unusable due to evaluating solutions outside their domain giving NaN
'throw': False,
# NN ARCHITECTURE & TRAINING
'nn_type': 'leaky',
# 'nn_layerdims': (16, 16, 16),
# 'nn_layerdims': (32, 32, 32),
# 'nn_layerdims': (64, 64, 64),
# 'nn_layerdims': (128, 128, 128),
# 'nn_layerdims': (256, 256, 256),
# 'nn_layerdims': (512, 512, 512),
# 'nn_layerdims': (32, 32, 32, 32),
# 'nn_layerdims': (128, 8),
# nicer for the launch script
'nn_n_layers': 3,
'nn_layer_dim': 256,
'nn_batchsize': 32,
'nn_N_epochs': 512,
'nn_train_fraction': .98,
'lr_staircase': False,
'lr_staircase_steps': 8,
'lr_init': 0.01,
'lr_final': 0.001,
'weight_decay': .003,
'nn_ensemble_size': 4,
'nn_warm_start': True,
'nn_warmstart_fraction': 1/2,
'nn_value_sweep': True,
'nn_progressbar': True,
# NN LOSS FUNCTION
# relative importance of the losses for v, vx, vxx.
# mostly we care about representing vx with great accuracy,
# the other two can be thought of as "hints"/priors/inductive biases
# to fit the correct vx function.
# update: vxx not used anymore, leave it at 0 or update lots of code
'nn_sobolev_weight_v': 1.,
'nn_sobolev_weight_vx': 9.,
'nn_sobolev_weight_vxx': 0.,
# width of the quadratic regions in smoothed huber loss.
'vx_loss_d': 0.3,
'v_loss_d': 0.1,
# above those thresholds relative loss is used
'min_important_v': 1.,
'min_important_vx': 1.,
# penalisation of the extra value derivative which is defined in the ambient space
# but normal to the state manifold.
'vx_normal_regularisation': 0.001,
# this is not a proper "prior" in the bayesian sense, but rather
# just an additional weak loss term that makes the value function
# large-ish at the problematic state of being upside down but
# otherwise at equilibrium.
'prior_strength': 0.01,
'v_prior': 10.,
'inv_vx_loss_fadeout': 5.,
# MAIN ALGO
# only take a subsample of data for active learning. dense sample
# close to current level set, less dense sample further down.
# the uncertainty bound we wish to satisfy.
# sigma_max(mu) = simga_max_abs + simga_max_rel * mu
'sigma_max_abs': 0.5,
'sigma_max_rel': 0.05,
# value band for training = [v_k / thin_data_denominator, v_next_target]
'thin_data': True,
'thin_data_denominator': 5,
# initial data generation. 'uniform' or 'lqr' for nicer distribution.
'initial_shooting': 'lqr',
# the value level we include in the initial learning round.
'v_init': 50,
# number of proposals per active learning iteration.
# larger = nicer! but don't kill our poor RAM
'initial_batchsize': 128,
'active_learning_batchsize': 128,
'include_future_data': False, # shit idea
'consider_old_data': True, # less shit idea?
'relative_kernel_lengthscale': 1/4,
# the max. time horizon by which we aim to grow the known level set
# in one iteration.
'T_value_target': 1.,
'vk_estimator': 'k_exceptions',
'pruning_strategy': 'conservative',
'proposal_sampling_distribution': 'uniform',
'proposal_strategy': 'max_kernel_adaptive',
'proposal_kernel_scaling': .5,
# the sublevel set Vk must contain at least this fraction of test points
# which are below the sigma target to qualify as "learned".
# only applies for 'vk_estimator' == 'relaxed'.
'frac_certain_in_Vk': .99,
# OUTPUT & VISUALISATION
'wandb': False,
'savefigs': False,
'wandbfigs': False,
'showfigs': True,
'ipdb_interval': 8,
# set this in euler launch script to filter wandb. does nothing otw.
'sweep_name': 'default',
# EVALUATION
'eval': '', # instead of None bc argparse wants same type.
}
def sample_states_batched(key, N, extent, log_min_scale=0):
# vmapped version of the above, but also scales down half the
# points with logspace'd distribution.
# maybe the "scaling" should also be part of the sampling fct?
# stochastic not determinstic? probably only cosmetic though
keys = jax.random.split(key, N)
# if log_min_scale != 0, this will scale half the points down with a
# logarithmically scaled factor, while the other half will stay the same.
log_min_scale = -np.abs(log_min_scale)
scales = np.clip(np.logspace(log_min_scale, -log_min_scale, N), -np.inf, 1.)[:, None]
pts = jax.vmap(sample_state, in_axes=(0, None, 0))(keys, extent, scales)
return pts
def sample_state(key, extent, scale=1.):
# sample points "uniformly" from "the whole state space".
# problem specific function! thus outside in problem_params.
# key: usual PRNG key
# extent: np.array of shape (nx,). [-extent, extent] are the box
# bounds for uniform sampling. Manifold states cosPhi, sinPhi
# treated separately so their "extent" is irrelevant.
# scale: scales the sample by some scalar.
# TODO think about what happens when x_eq != 0 -- just add it here?
# maybe (especially for higher dims) ellipsids are better? a bit like this:
# - sample from unit normal
# - transform magnitude of samples such that they are uniform within unit ball
# (inverse transform normcdf chi squared something, i think I did this once)
# - squash with matrix A to transform to ellipse {z: || z.T inv(A).T inv(A) z || <= 1 }
# - sample from different scaled versions of this ellipse to avoid the soap bubble effect :)
# but this is just an intuitive hunch, because for a uniform box most
# of the volume is at the corners, where we might not want it. also
# these effects probably don't really kick in at like 6 to 12 dims.
# separate the "flat" R^n part and the manifold part.
rnkey, manifoldkey = jax.random.split(key)
# generate uniform points from a box in R^n
x_pt = jax.random.uniform(
key=rnkey,
shape=extent.shape,
minval=-extent,
maxval= extent,
) * scale
# for the manifold part: uniform sampling from unit circle.
# if we instead generate this by drawing Phi ~ U[-pi, pi]
# we can apply the same scaling logic to the angle and only
# then convert to ambient space representation...
# generate 2D gaussian & normalise
xy = jax.random.normal(key=manifoldkey, shape=(2,))
xy = xy / np.linalg.norm(xy)
# indices of sinPhi and cosPhi states.
assert problem_params['state_names'][2] == 'sinPhi'
assert problem_params['state_names'][3] == 'cosPhi'
x_pt = x_pt.at[2:4].set(xy)
return x_pt
# just pass on the entire functions :)
algo_params['sample_state'] = sample_state
algo_params['sample_states_batched'] = sample_states_batched
return algo_params
if __name__ == '__main__':
problem_params = define_problem_params()
algo_params = base_algo_params()
# argparser based on the algo_params dict.
parser = argparse.ArgumentParser()
arg_types = (bool, int, float, str)
# thanks stackoverflow
# https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse
def _str_to_bool(s):
"""Convert string to bool (in argparse context)."""
if s.lower() not in ['true', 'false']:
raise ValueError('Need bool; got %r' % s)
return {'true': True, 'false': False}[s.lower()]
def add_boolean_argument(parser, name, default=False):
"""Add a boolean argument to an ArgumentParser instance."""
group = parser.add_mutually_exclusive_group()
group.add_argument(
name, nargs='?', default=default, const=True, type=_str_to_bool)
group.add_argument('--no' + name, dest=name, action='store_false')
for k in algo_params:
t = type(algo_params[k])
if t in arg_types:
if t == bool:
add_boolean_argument(parser, f'--{k}', default=algo_params[k])
else:
parser.add_argument(f'--{k}', type=type(algo_params[k]), default=algo_params[k])
commandline_args = parser.parse_args()
# now, put the arguments back into the algo_params dict
for k in algo_params:
if type(algo_params[k]) in arg_types:
new_arg = getattr(commandline_args, k)
old_arg = algo_params[k]
if type(new_arg) != type(old_arg):
raise ValueError(f'argument {k} has type {type(new_arg)} but should have type {type(old_arg)}')
algo_params[k] = new_arg
# levelsets.evaluate('euler_runs/8dgpt7uo', problem_params, algo_params)
# levelsets.evaluate('euler_runs/ff5mij89', problem_params, algo_params)
# levelsets.evaluate('euler_runs/uqf3ybp8', problem_params, algo_params)
if algo_params['eval'] == '':
# hacking this to make remaining larger batchsize run worth it.
if algo_params['active_learning_batchsize'] > 256:
# if more data, ratio>1, we want less epochs.
ratio = algo_params['active_learning_batchsize'] / 256
algo_params['nn_warmstart_fraction'] = algo_params['nn_warmstart_fraction'] / ratio
levelsets.main(problem_params, algo_params)
else:
# algo_params['wandb'] = False
# levelsets.evaluate(algo_params['eval'], problem_params, algo_params)
# levelsets.evaluate_from_wandb(algo_params['eval'], problem_params)
run_dir = algo_params['eval']
run_id = run_dir.split('/')[-1]
print('getting algo_params from wandb...')
api = wandb.Api()
sysname = problem_params['system_name']
# thanks to https://github.com/wandb/wandb/issues/5122 :)
runs = api.runs(path=f'mbjd-projects/levelsets_{sysname}', filters={'name': {'$in': [run_id]}})
assert len(runs) == 1, 'expected to get just 1 run per id...'
run = runs[0]
algo_params = run.config
base = base_algo_params()
algo_params['sample_state'] = base['sample_state']
algo_params['sample_states_batched'] = base['sample_states_batched']
algo_params['wandb'] = False # otw it wants to log the results...
levelsets.evaluate(run_dir, problem_params, algo_params)