def test_smoother(): utils.set_rng_seed(1) torch.set_default_dtype(torch.float64) sigma = torch.tensor([10.]) rho = torch.tensor([28.]) beta = torch.tensor([8. / 3.]) C = torch.randn(2, 3) dt = 0.04 sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint') B = 1 T = 200 xs = [torch.randn(B, 1, 3)] for t in range(T - 1): xs.append(sys.step(torch.tensor([0.] * B), xs[-1])) x = torch.cat(xs, dim=1).detach() x.requires_grad = True y = x.detach() sys = BasicLagrangianSystem(qdim=3, dt=dt) t = torch.stack([torch.arange(T), torch.arange(T)]).to(torch.get_default_dtype()) x0 = torch.zeros_like(x) x0.requires_grad = True obscrit = GaussianObservationCriterion(torch.ones(3), t, y) dyncrit = DELCriterion(t) # Test GroupSOSCriterion crit = GroupSOSCriterion([obscrit, dyncrit]) xsm, metrics = NLSsmoother(x0, crit, sys, solver_kwargs={ 'verbose': 2, 'tr_rho': 0., 'max_nfev': 2 }) print('Passed.') # Test BlockSparseGroupSOSCriterion # crit = BlockSparseGroupSOSCriterion([obscrit, dyncrit]) # xsm, metrics = NLSsmoother(torch.zeros_like(x), crit, sys) print('Passed.')
def test_learner(): utils.set_rng_seed(1) torch.set_default_dtype(torch.float64) sigma = torch.tensor([10.]) rho = torch.tensor([28.]) beta = torch.tensor([8. / 3.]) C = torch.randn(2, 3) dt = 0.04 sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint') B = 1 T = 200 xs = [torch.randn(B, 1, 3)] for t in range(T - 1): xs.append(sys.step(torch.tensor([0.] * B), xs[-1])) x = torch.cat(xs, dim=1).detach() t = torch.stack([torch.arange(T), torch.arange(T)]).to(torch.get_default_dtype()) dyncrit = GaussianDynamicsCriterion(torch.ones(3), t) params = list(sys.parameters())[:2] vparams = parameters_to_vector(params) true_vparams = vparams.clone() vparams += torch.randn_like(vparams) * 0.01 vector_to_parameters(vparams, params) opt_result = learner(sys, [dyncrit], [x], ['torch_minimize'], [params], [{}], opt_kwargs_list=[{ 'nan_line_search': True, 'scheduler': 'ExponentialLR', 'scheduler_kwargs': { 'gamma': 0.99, } }])[0] params = list(sys.parameters())[:2] vparams = parameters_to_vector(params) error = (vparams - true_vparams).norm().item() assert np.allclose(error, 0., atol=1e-4), 'Error=%.3e' % error print('Passed.')
def train_horizon_model(H, datadir): # extract hyperparameters hp = hyperparameters[H] lr = hp['lr'] n_epochs = hp['n_epochs'] save_file = hp['save_file'] logdir = hp['logdir'] # load data utils.set_rng_seed(1) y_mean, y_std, u_mean, u_std = load_statistics(datadir) train_data, train_trgt = load_helidata(datadir, 'train') test_data, test_trgt = load_helidata(datadir, 'test') valid_data, valid_trgt = load_helidata(datadir, 'valid') # Define net neural_net_kwargs = dict(input_size=10 * H, hidden_sizes=[32] * 8, output_size=6, activation='tanh', gain=1.0, ln=False) net = LagModel(neural_net_kwargs, H) logger.setup(logdir, action='d') # Train system = train_net(net, train_data, train_trgt, test_data, test_trgt, valid_data, valid_trgt, y_std, lr, logdir=logdir, H=H, n_epochs=n_epochs) torch.save(system.state_dict(), save_file)
def main(): utils.set_rng_seed(1) # load the data for dset in ['0p01', '0p05', '0p10', '0p20', '0p30', '0p40', '0p50', '1p0']: for damped in ['', 'damped_']: data = torch.load('./datasets/%sdubpen_qddot.td'%damped) data_ = torch.load('./datasets/%sdubpen_%s.td'%(damped,dset)) dt = 0.05 t_, q_, y = data_[:] t, q, dq, ddq = data[:] B, T, qdim = q.shape with torch.no_grad(): std = 0.05 * 10 smoothed_q, smoothed_dq, smoothed_ddq = kalman_smooth_and_diff(y,dt, em_Q=False) # new! if plot: for b in range(B): for n in range(2): plt.subplot(3,2,1+n) plt.plot(smoothed_q[b,:,n]) plt.plot(q_[b,:,n], '--') plt.plot(y[b,:,n], alpha=0.5) plt.subplot(3,2,3+n) plt.plot(smoothed_dq[b,:,n]) plt.plot(dq[b,:,n], '--') plt.subplot(3,2,5+n) plt.plot(smoothed_ddq[b,:,n]) plt.plot(ddq[b,:,n], '--') plt.show() dataset = TensorDataset(t_, smoothed_q, smoothed_dq, smoothed_ddq) torch.save(dataset, './datasets/%sdubpen_%s_smoothed.td'%(damped, dset))
def train(datadir, logdir, lr, n_epochs, H, save_file='lstm.th'): utils.set_rng_seed(1) y_mean, y_std, u_mean, u_std = load_statistics(datadir) train_data, train_trgt = load_helidata(datadir, 'train') test_data, test_trgt = load_helidata(datadir, 'test') valid_data, valid_trgt = load_helidata(datadir, 'valid') net = LSTMModel(**DEFAULT_LSTM_KWARGS) logger.setup(logdir, action='d') # Train system = train_net(net, train_data, train_trgt, test_data, test_trgt, valid_data, valid_trgt, y_std, lr, logdir=logdir, H=H, n_epochs=n_epochs) torch.save(system.state_dict(), save_file)
def main(logdir, datafile, valdatafile, subset, obsmodel, practice, wfac, xdim, hotstartdir, seed, parallel): utils.set_rng_seed(seed) torch.set_default_dtype(torch.float64) logger.setup(logdir) # load train data traindata = np.load(datafile) u = torch.tensor(traindata['u']) y = torch.tensor(traindata['y']) valdata = np.load(valdatafile) val_u = torch.tensor(valdata['u']) val_y = torch.tensor(valdata['y']) _, y_std, _, _ = load_statistics('./datasets/split_normalized') if practice == 1: inds = np.random.choice(466, 10, replace=False) u = u[inds] y = y[inds] B, T, ydim = y.shape _, _, udim = u.shape t = torch.stack([torch.arange(T)] * B).to(torch.get_default_dtype()) # specify system sid = loadmat(opj(hotstartdir, 'SID_%dD.mat' % xdim)) A = torch.tensor(sid['A']).to(torch.get_default_dtype()) Bsys = torch.tensor(sid['B']).to(torch.get_default_dtype()) C = torch.tensor(sid['C']).to(torch.get_default_dtype()) D = torch.tensor(sid['D']).to(torch.get_default_dtype()) NN = deepcopy(DEFAULT_NN) NN['gain'] = 0.1 sys = DiscreteLinear(xdim, udim, ydim, A=A, B=Bsys, C=C, D=D, obsModel=NN) # specify smoothing criteria smoothing_criteria = [] vfac = 1.0 for b in range(B): obscrit = GaussianObservationCriterion(vfac * torch.ones(ydim), t[b:b + 1], y[b:b + 1], u=u[b:b + 1]) dyncrit = GaussianDynamicsCriterion(wfac * torch.ones(xdim), t[b:b + 1], u=u[b:b + 1]) smoothing_criteria.append(GroupSOSCriterion([obscrit, dyncrit])) smooth_solver_kwargs = {'verbose': 0, 'tr_rho': 0.1} # specify learning criteria learning_criteria = [ GaussianObservationCriterion(torch.ones(ydim), t, y, u=u), GaussianDynamicsCriterion(torch.ones(xdim), t, u=u) ] learning_params = [list(sys._obs.parameters()), list(sys._dyn.parameters())] learning_opts = ['torch_minimize', 'torch_minimize'] learner_opt_kwargs = [{ 'method': 'Adam', 'lr': 5e-4, 'nepochs': 500, 'tr_rho': 0.5 }, { 'method': 'Adam', 'lr': 1e-3, 'nepochs': 500, 'tr_rho': 0.5 }] # save params run_params = dict(seed=seed, subset=subset, xdim=xdim, vfac=vfac, wfac=wfac, learning_opts=learning_opts, learner_opt_kwargs=learner_opt_kwargs, smooth_solver_kwargs=smooth_solver_kwargs, practice=practice, obsmodel=obsmodel) with open(opj(logdir, 'run_params.json'), 'w') as f: json.dump(run_params, f) # instantiate CEEM class Tracker: def __init__(self): return tracker = Tracker() tracker.best_val_rmse = np.inf def ecb(epoch): torch.save(sys.state_dict(), os.path.join(logger.get_dir(), 'ckpts', 'model_{}.th'.format(epoch))) y_pred = gen_ypred_model(sys, val_u, val_y) rms = compute_rms(val_y[:, 25:], y_pred[:, 25:], y_std) val_rmse = float(rms.mean()) logger.logkv('test/val_rmse', val_rmse) if val_rmse < tracker.best_val_rmse: tracker.best_val_rmse = val_rmse torch.save(sys.state_dict(), os.path.join(logger.get_dir(), 'ckpts', 'best_model.th')) return epoch_callbacks = [ecb] def tcb(epoch): # TODO return False termination_callback = tcb ceem = CEEM(smoothing_criteria, learning_criteria, learning_params, learning_opts, epoch_callbacks, termination_callback, parallel=parallel) # run CEEM x0 = 0.01 * torch.randn(B, T, xdim) ceem.train(xs=x0, sys=sys, nepochs=5000, smooth_solver_kwargs=smooth_solver_kwargs, learner_opt_kwargs=learner_opt_kwargs, subset=subset)
def train(seed, logdir, sys_seed, k, b): from mpl_toolkits.mplot3d import Axes3D import matplotlib.pyplot as plt ystd = 0.01 # ystd = 0. torch.set_default_dtype(torch.float64) logger.setup(logdir, action='d') N = 128 n = 3 * k B = b true_system = default_lorenz_system(k, obsdif=2) utils.set_rng_seed(sys_seed) xdim = true_system.xdim ydim = true_system.ydim dt = true_system._dt x0mean = torch.tensor([[-6] * k + [-6] * k + [24.] * k]).unsqueeze(0) # simulate true_dynamics over IC distribution x_test = x0mean.repeat(1024, 1, 1) x_test += 5.0 * torch.randn_like(x_test) x_test = x_test.detach() t_test = torch.zeros(1024, 1) tgt_test = true_system.step_derivs(t_test, x_test).detach() ## simulate the true system xs = [x0mean.repeat(B, 1, 1)] xs[0] += 2.5 * torch.randn_like(xs[0]) with torch.no_grad(): for t in range(N - 1): xs.append(true_system.step(torch.tensor([0.] * B), xs[-1])) xs = torch.cat(xs, dim=1) fig = plt.figure() for b in range(B): ax = fig.add_subplot(int(np.ceil(B / 2.)), 2, b + 1, projection='3d') for k_ in range(k): plot3d(plt.gca(), xs[b, :, k_], xs[b, :, k + k_], xs[b, :, 2 * k + k_], linestyle='--', alpha=0.5) plt.savefig(os.path.join(logger.get_dir(), 'figs/traj_%d.png' % b), dpi=300) # plt.show() plt.close() t = torch.tensor(range(N)).unsqueeze(0).expand(B, -1).to( torch.get_default_dtype()) y = true_system.observe(t, xs).detach() # seed for real now utils.set_rng_seed(seed) y += ystd * torch.randn_like(y) # prep system system = deepcopy(true_system) true_params = parameters_to_vector(true_system.parameters()) utils.set_rng_seed(seed) params = true_params * ( (torch.rand_like(true_params) - 0.5) / 5. + 1.) # within 10% vector_to_parameters(params, system.parameters()) params = list(system.parameters()) # specify smoothing criteria smoothing_criteria = [] for b in range(B): obscrit = GaussianObservationCriterion(1.0 * torch.ones(ydim), t[b:b + 1], y[b:b + 1]) dyncrit = GaussianDynamicsCriterion(1e0 * torch.ones(xdim), t[b:b + 1]) smoothing_criteria.append(GroupSOSCriterion([obscrit, dyncrit])) smooth_solver_kwargs = {'verbose': 0, 'tr_rho': 0.01} # specify learning criteria learning_criteria = [GaussianDynamicsCriterion(1e0 * torch.ones(xdim), t)] learning_params = [params] # learning_opts = ['scipy_minimize'] # learner_opt_kwargs = {'method': 'Nelder-Mead', 'tr_rho': 0.1, # 'options':{'adaptive':True}} # learner_opt_kwargs = {'method': 'BFGS', 'tr_rho': 0.1, # 'options':{'disp':True}} learning_opts = ['torch_minimize'] # learner_opt_kwargs = { # 'method': 'Adam', # 'lr': 5e-4, # 'tr_rho': 0.1, # 'nepochs': 200, # 'max_grad_norm': 10.0 # } learner_opt_kwargs = {'method': 'LBFGS'} # instantiate CEEM def ecb(epoch): params = list(system.parameters()) vparams = parameters_to_vector(params) error = (vparams - true_params).norm().item() logger.logkv('test/log10_paramerror', np.log10(error)) return epoch_callbacks = [ecb] class Last10Errors: def __init__(self): return last_10_errors = Last10Errors last_10_errors._arr = [] def tcb(epoch): with torch.no_grad(): tgt_test_pr = system.step_derivs(t_test, x_test) error = float(torch.nn.functional.mse_loss(tgt_test_pr, tgt_test)) logger.logkv('test/log10_error', np.log10(error)) last_10_errors._arr.append(np.log10(error)) if len(last_10_errors._arr) > 100: last_10_errors._arr = last_10_errors._arr[-100:] l10err = torch.tensor(last_10_errors._arr) convcrit = float((l10err.min() - l10err.max()).abs()) logger.logkv('test/log10_convcrit', np.log10(convcrit)) if convcrit < 1e-3: return True return False termination_callback = tcb ecb(-1) tcb(-1) logger.dumpkvs() ceem = CEEM(smoothing_criteria, learning_criteria, learning_params, learning_opts, epoch_callbacks, termination_callback, parallel=min(4, B)) # run CEEM # x0 = torch.zeros_like(xs) x0 = xs + torch.randn_like(xs) ceem.train(xs=x0, sys=system, nepochs=5000, smooth_solver_kwargs=smooth_solver_kwargs, learner_opt_kwargs=learner_opt_kwargs) return
def main(): utils.set_rng_seed(2) torch.set_default_dtype(torch.float64) dt = 0.05 sys = LagrangianDoublePendulum(dt, 1., 1., 1., 1., 10.) q1 = torch.rand(16, 1, 2) * np.pi - np.pi / 2 q2 = q1.clone() qs = [q1, q2] for t in tqdm(range(200)): qt = qs[-2].detach() qtp1 = qs[-1].detach() nq = sys.variational_step(qt, qtp1, oneatatime=True) qs.append(nq) x = torch.cat(qs, dim=1).detach() B, T, _ = x.shape t = torch.arange(T).unsqueeze(0).repeat(B, 1).float() y_p01 = (x + 0.01 * torch.randn_like(x)).detach() y_p05 = (x + 0.05 * torch.randn_like(x)).detach() y_p10 = (x + 0.1 * torch.randn_like(x)).detach() y_p20 = (x + 0.2 * torch.randn_like(x)).detach() y_p30 = (x + 0.3 * torch.randn_like(x)).detach() y_p40 = (x + 0.4 * torch.randn_like(x)).detach() y_p50 = (x + 0.5 * torch.randn_like(x)).detach() y_1p0 = (x + 1.0 * torch.randn_like(x)).detach() dataset = TensorDataset(t, x, y_p01) torch.save(dataset, './datasets/dubpen_0p01.td') dataset = TensorDataset(t, x, y_p05) torch.save(dataset, './datasets/dubpen_0p05.td') dataset = TensorDataset(t, x, y_p10) torch.save(dataset, './datasets/dubpen_0p10.td') dataset = TensorDataset(t, x, y_p20) torch.save(dataset, './datasets/dubpen_0p20.td') dataset = TensorDataset(t, x, y_p30) torch.save(dataset, './datasets/dubpen_0p30.td') dataset = TensorDataset(t, x, y_p40) torch.save(dataset, './datasets/dubpen_0p40.td') dataset = TensorDataset(t, x, y_p50) torch.save(dataset, './datasets/dubpen_0p50.td') dataset = TensorDataset(t, x, y_1p0) torch.save(dataset, './datasets/dubpen_1p0.td') x_ = (x[:, 1:] + x[:, :-1]) * 0.5 dx = (x[:, 1:] - x[:, :-1]) / dt x = x_ ddx = sys.compute_qddot(x, dx) B, T, _ = x.shape t = torch.arange(T).unsqueeze(0).repeat(B, 1).float() dataset = TensorDataset(t, x, dx, ddx) torch.save(dataset, './datasets/dubpen_qddot.td') if plot: for b in range(16): plt.subplot(8, 2, b + 1) plt.plot(x[b]) plt.show()
def test_particleem(): utils.set_rng_seed(1) torch.set_default_dtype(torch.float64) # setup system sigma = torch.tensor([10.]) rho = torch.tensor([28.]) beta = torch.tensor([8. / 3.]) C = torch.randn(2, 3) dt = 0.04 sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint') # simulate B = 2 T = 20 xs = [torch.randn(B, 1, 3)] for t in range(T - 1): xs.append(sys.step(torch.tensor([0.] * B), xs[-1])) xtr = torch.cat(xs, dim=1).detach() y = sys.observe(0., xtr).detach() # y += torch.rand_like(y) * 0.01 t = torch.stack([torch.arange(T)] * B).to(torch.get_default_dtype()) # params = list(sys.parameters()) vparams = parameters_to_vector(params) true_vparams = vparams.clone() vparams *= 1. + (torch.rand_like(vparams) - 0.5) * 2 * 0.025 vector_to_parameters(vparams, params) Px0 = torch.eye(3) Q = 0.1 * torch.eye(3) R = 0.1 * torch.eye(2) Np = 300 fapf = faPF(Np, sys, Q, R, Px0) def callback(epoch): pass trainer = SAEMTrainer( fapf, y, gamma_sched=lambda x: 0.8, xlen_cutoff=3, max_k=10, ) trainer.train(params, callbacks=[callback]) assert True if True: print('Passed.') else: print('Failed.')
def test_smoother(): utils.set_rng_seed(1) torch.set_default_dtype(torch.float64) sigma = torch.tensor([10.]) rho = torch.tensor([28.]) beta = torch.tensor([8. / 3.]) C = torch.randn(2, 3) dt = 0.04 sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint') B = 1 T = 200 xs = [torch.randn(B, 1, 3)] for t in range(T - 1): xs.append(sys.step(torch.tensor([0.] * B), xs[-1])) x = torch.cat(xs, dim=1).detach() x.requires_grad = True y = sys.observe(0., x).detach() # y += torch.rand_like(y) * 0.01 t = torch.stack([torch.arange(T), torch.arange(T)]).to(torch.get_default_dtype()) x0 = torch.zeros_like(x) obscrit = GaussianObservationCriterion(torch.ones(2), t, y) dyncrit = GaussianDynamicsCriterion(torch.ones(3), t) # Test GroupSOSCriterion crit = GroupSOSCriterion([obscrit, dyncrit]) xsm, metrics = NLSsmoother(x0, crit, sys, solver_kwargs={ 'verbose': 2, 'tr_rho': 0. }) err = float((xsm - x).norm()) assert err < 1e-8, 'Smoothing Error: %.3e' % err print('Passed.') # Test BlockSparseGroupSOSCriterion crit = BlockSparseGroupSOSCriterion([obscrit, dyncrit]) xsm, metrics = NLSsmoother(torch.zeros_like(x), crit, sys) err = float((xsm - x).norm()) assert err < 1e-8, 'Smoothing Error: %.3e' % err print('Passed.')
def train(seed, logdir, ystd=0.1, wstd=0.01, sys_seed=4): print('\n\n\n##### SEED %d #####\n\n' % seed) torch.set_default_dtype(torch.float64) logger.setup(logdir, action='d') # Number of timesteps in the trajectory T = 128 n = 3 # Batch size B = 4 k = 1 utils.set_rng_seed(sys_seed) sys = default_lorenz_attractor() dt = sys._dt utils.set_rng_seed(seed) # simulate the system x0mean = torch.tensor([[-6] * k + [-6] * k + [24.] * k]) x0mean = x0mean.unsqueeze(0).repeat(B, 1, 1) # Rollout with noise Q = (wstd**2) * torch.eye(sys.xdim) R = (ystd**2) * torch.eye(sys.ydim) Px0 = 5.0 * torch.eye(sys.xdim) Qpdf = MultivariateNormal(torch.zeros((B, 1, sys.xdim)), Q.unsqueeze(0).unsqueeze(0)) Rpdf = MultivariateNormal(torch.zeros((B, 1, sys.ydim)), R.unsqueeze(0).unsqueeze(0)) Px0pdf = MultivariateNormal(x0mean, Px0.unsqueeze(0).unsqueeze(0)) xs = [Px0pdf.sample()] ys = [sys.observe(0, xs[0]) + Rpdf.sample()] for t in range(T - 1): tinp = torch.tensor([t] * B).unsqueeze(1).to(dtype=torch.get_default_dtype()) xs.append(sys.step(tinp, xs[-1]) + Qpdf.sample()) ys.append(sys.observe(tinp, xs[-1]) + Rpdf.sample()) x = torch.cat(xs, dim=1) ys = torch.cat(ys, dim=1) m = ys.shape[-1] fig = plt.figure() for b in range(B): ax = fig.add_subplot(int(np.ceil(B / 2.)), 2, b + 1, projection='3d') for k_ in range(k): plot3d(plt.gca(), x[b, :, k_], x[b, :, k + k_], x[b, :, 2 * k + k_], linestyle='--', alpha=0.5) plt.savefig(opj(logdir, 'traintrajs.png'), dpi=300) true_system = deepcopy(sys) Np = 100 params = [sys._sigma, sys._rho, sys._beta] true_vparams = parameters_to_vector(params) pert_vparams = true_vparams * ( (torch.rand_like(true_vparams) - 0.5) / 5 + 1.0) vector_to_parameters(pert_vparams, params) fapf = faPF(Np, sys, Q, R, Px0) timer = {'start_time': timeit.default_timer()} def callback(epoch): logger.logkv('test/rho', float(sys._rho)) logger.logkv('test/sigma', float(sys._sigma)) logger.logkv('test/beta', float(sys._beta)) logger.logkv( 'test/rho_pcterr_log10', float( torch.log10( (true_system._rho - sys._rho).abs() / true_system._rho))) logger.logkv( 'test/sigma_pcterr_log10', float( torch.log10((true_system._sigma - sys._sigma).abs() / true_system._sigma))) logger.logkv( 'test/beta_pcterr_log10', float( torch.log10((true_system._beta - sys._beta).abs() / true_system._beta))) logger.logkv('time/epochtime', timeit.default_timer() - timer['start_time']) timer['start_time'] = timeit.default_timer() return callback(-1) logger.dumpkvs() trainer = SAEMTrainer( fapf, ys, # gamma_sched=lambda x: HarmonicDecayScheduler(x, a=50.), gamma_sched=lambda x: 0.8) trainer.train(params, callbacks=[callback])
def train(seed, logdir, sys_seed, ystd, wstd): torch.set_default_dtype(torch.float64) logger.setup(logdir, action='d') # Number of timesteps in the trajectory T = 128 n = 3 # Batch size B = 1 k = 1 utils.set_rng_seed(sys_seed) true_system = default_lorenz_attractor() dt = true_system._dt utils.set_rng_seed(43) # simulate the system x0mean = torch.tensor([[-6] * k + [-6] * k + [24.] * k]).unsqueeze(0) # seed for real now utils.set_rng_seed(seed) # Rollout with noise xs = [x0mean] xs[0] += 5. * torch.randn_like(xs[0]) with torch.no_grad(): for t in range(T - 1): xs.append( true_system.step(torch.tensor([0.] * B), xs[-1]) + wstd * torch.randn_like(xs[-1])) xs = torch.cat(xs, dim=1) t = torch.tensor(range(T)).unsqueeze(0).to(torch.get_default_dtype()) y = true_system.observe(t, xs).detach() y += ystd * torch.randn_like(y) # Observation noise # prep system system = deepcopy(true_system) true_params = parameters_to_vector(true_system.parameters()) params = true_params * ((torch.rand_like(true_params) - 0.5) / 5. + 1.) # within 10% vector_to_parameters(params, system.parameters()) params = list(system.parameters()) # specify smoothing criteria B = 1 smoothing_criteria = [] for b in range(B): obscrit = GaussianObservationCriterion(torch.ones(2), t[b:b + 1], y[b:b + 1]) dyncrit = GaussianDynamicsCriterion(wstd / ystd * torch.ones(3), t[b:b + 1]) smoothing_criteria.append(GroupSOSCriterion([obscrit, dyncrit])) smooth_solver_kwargs = {'verbose': 0, 'tr_rho': 0.001} # specify learning criteria learning_criteria = [GaussianDynamicsCriterion(torch.ones(3), t)] learning_params = [params] learning_opts = ['scipy_minimize'] learner_opt_kwargs = {'method': 'Nelder-Mead', 'tr_rho': 0.01} # instantiate CEEM def ecb(epoch): logger.logkv('test/rho', float(system._rho)) logger.logkv('test/sigma', float(system._sigma)) logger.logkv('test/beta', float(system._beta)) logger.logkv('test/rho_pcterr_log10', float(torch.log10((true_system._rho - system._rho).abs() / true_system._rho))) logger.logkv( 'test/sigma_pcterr_log10', float(torch.log10((true_system._sigma - system._sigma).abs() / true_system._sigma))) logger.logkv( 'test/beta_pcterr_log10', float(torch.log10((true_system._beta - system._beta).abs() / true_system._beta))) return epoch_callbacks = [ecb] class Last10Errors: def __init__(self): return last_10_errors = Last10Errors last_10_errors._arr = [] def tcb(epoch): params = list(system.parameters()) vparams = parameters_to_vector(params) error = (vparams - true_params).norm().item() last_10_errors._arr.append(float(error)) logger.logkv('test/log10_error', np.log10(error)) if len(last_10_errors._arr) > 10: last_10_errors._arr = last_10_errors._arr[-10:] l10err = torch.tensor(last_10_errors._arr) convcrit = float((l10err.min() - l10err.max()).abs()) logger.logkv('test/log10_convcrit', np.log10(convcrit)) if convcrit < 1e-4: return True return False termination_callback = tcb ceem = CEEM(smoothing_criteria, learning_criteria, learning_params, learning_opts, epoch_callbacks, termination_callback) # run CEEM x0 = torch.zeros_like(xs) ceem.train(xs=x0, sys=system, nepochs=500, smooth_solver_kwargs=smooth_solver_kwargs, learner_opt_kwargs=learner_opt_kwargs) return float(system._sigma), float(system._rho), float(system._beta)
def test_sys(): utils.set_rng_seed(1) torch.set_default_dtype(torch.float64) # test LorenzAttractor sigma = torch.tensor([10.]) rho = torch.tensor([28.]) beta = torch.tensor([8. / 3.]) C = torch.randn(2, 3) dt = 0.04 sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint') B = 5 T = 20 xs = [torch.randn(B, 1, 3)] for t in range(T - 1): xs.append(sys.step(torch.tensor([0.] * B), xs[-1])) x = torch.cat(xs, dim=1).detach() x.requires_grad = True y = sys.observe(0., x) y += torch.rand_like(y) * 0.01 t = torch.stack([torch.arange(T), torch.arange(T)]).to(torch.get_default_dtype()) check_sys(sys, t, x, y) # test SpringMassDamper n = 4 M = D = K = torch.tensor([[1., 2.], [2., 5.]]) dt = 0.1 method = 'midpoint' sys = SpringMassDamper(M, D, K, dt, method=method) B = 5 T = 20 xs = [torch.randn(B, 1, n)] for t in range(T - 1): xs.append(sys.step(torch.tensor([0.] * B), xs[-1])) x = torch.cat(xs, dim=1).detach() x.requires_grad = True y = sys.observe(0., x).detach() y += torch.rand_like(y) * 0.01 t = torch.stack([torch.arange(T), torch.arange(T)]).to(torch.get_default_dtype()) check_sys(sys, t, x, y) # test DiscreteLinear xdim = 2 ydim = 3 udim = 2 sys = DiscreteLinear(xdim, udim, ydim) x = torch.randn(B, T, xdim) u = torch.randn(B, T, udim) y = torch.randn(B, T, ydim) t = torch.stack([torch.arange(T), torch.arange(T)]).to(torch.get_default_dtype()) check_sys(sys, t, x, y, u=u)
def test_ceem(): utils.set_rng_seed(1) torch.set_default_dtype(torch.float64) # setup system sigma = torch.tensor([10.]) rho = torch.tensor([28.]) beta = torch.tensor([8. / 3.]) C = torch.randn(2, 3) dt = 0.04 sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint') # simulate B = 2 T = 20 xs = [torch.randn(B, 1, 3)] for t in range(T - 1): xs.append(sys.step(torch.tensor([0.] * B), xs[-1])) xtr = torch.cat(xs, dim=1).detach() y = sys.observe(0., xtr).detach() # y += torch.rand_like(y) * 0.01 t = torch.stack([torch.arange(T)] * B).to(torch.get_default_dtype()) # params = list(sys.parameters()) vparams = parameters_to_vector(params) true_vparams = vparams.clone() vparams *= 1. + (torch.rand_like(vparams) - 0.5) * 2 * 0.025 vector_to_parameters(vparams, params) # specify smoothing criteria smoothing_criteria = [] for b in range(B): obscrit = GaussianObservationCriterion(torch.ones(2), t[b:b + 1], y[b:b + 1]) dyncrit = GaussianDynamicsCriterion(torch.ones(3), t[b:b + 1]) smoothing_criteria.append(GroupSOSCriterion([obscrit, dyncrit])) smooth_solver_kwargs = {'verbose': 0, 'tr_rho': 0.001} # specify learning criteria learning_criteria = [GaussianDynamicsCriterion(torch.ones(3), t)] learning_params = [params] learning_opts = ['scipy_minimize'] learner_opt_kwargs = {'method': 'Nelder-Mead', 'tr_rho': 0.01} # learning_opts = ['torch_minimize'] # learner_opt_kwargs = {'method': 'Adam', # 'tr_rho': 0.01} # instantiate CEEM def ecb(epoch): return epoch_callbacks = [ecb] def tcb(epoch): params = list(sys.parameters()) vparams = parameters_to_vector(params) error = (vparams - true_vparams).norm().item() logger.logkv('test/log10_error', np.log10(error)) return error < 5e-3 termination_callback = tcb ceem = CEEM(smoothing_criteria, learning_criteria, learning_params, learning_opts, epoch_callbacks, termination_callback, parallel=2) hp_scheduler = {'tr_rho': utils.LambdaScheduler(0.01, lambda x: x)} # run CEEM x0 = torch.zeros_like(xtr) ceem.train(xs=x0, sys=sys, nepochs=150, smooth_solver_kwargs=smooth_solver_kwargs, learner_opt_kwargs=learner_opt_kwargs, subset=1, hp_schedulers=hp_scheduler) assert tcb(0) if tcb(0): print('Passed.') else: print('Failed.')
def train(seed, logdir, ystd=0.1, wstd=0.01, sys_seed=4): print('\n\n\n##### SEED %d #####\n\n'%seed) torch.set_default_dtype(torch.float64) logger.setup(logdir, action='d') # Number of timesteps in the trajectory T = 128 n = 3 # Batch size B = 4 k = 1 utils.set_rng_seed(sys_seed) sys = default_lorenz_attractor() dt = sys._dt utils.set_rng_seed(seed) # simulate the system x0mean = torch.tensor([[-6] * k + [-6] * k + [24.] * k]) x0mean = x0mean.unsqueeze(0).repeat(B,1,1) # Rollout with noise Q = (wstd ** 2) * torch.eye(sys.xdim) R = (ystd ** 2) * torch.eye(sys.ydim) Px0 = 5.0 * torch.eye(sys.xdim) Qpdf = MultivariateNormal(torch.zeros((B,1,sys.xdim)), Q.unsqueeze(0).unsqueeze(0)) Rpdf = MultivariateNormal(torch.zeros((B,1,sys.ydim)), R.unsqueeze(0).unsqueeze(0)) Px0pdf = MultivariateNormal(x0mean, Px0.unsqueeze(0).unsqueeze(0)) xs = [Px0pdf.sample()] ys = [sys.observe(0, xs[0]) + Rpdf.sample()] for t in range(T-1): tinp = torch.tensor([t] * B).unsqueeze(1).to(dtype=torch.get_default_dtype()) xs.append(sys.step(tinp, xs[-1]) + Qpdf.sample()) ys.append(sys.observe(tinp, xs[-1]) + Rpdf.sample()) x = torch.cat(xs, dim=1) y = torch.cat(ys, dim=1) t = torch.tensor(range(T)).unsqueeze(0).to(torch.get_default_dtype()).repeat(B,1) m = y.shape[-1] fig = plt.figure() for b in range(B): ax = fig.add_subplot(int(np.ceil(B / 2.)), 2, b + 1, projection='3d') for k_ in range(k): plot3d(plt.gca(), x[b, :, k_], x[b, :, k + k_], x[b, :, 2 * k + k_], linestyle='--', alpha=0.5) plt.savefig(opj(logdir, 'traintrajs.png'), dpi=300) # prep system true_system = sys system = deepcopy(true_system) true_params = parameters_to_vector(true_system.parameters()) params = true_params * ((torch.rand_like(true_params) - 0.5) / 5. + 1.) # within 10% vector_to_parameters(params, system.parameters()) params = list(system.parameters()) # specify smoothing criteria smoothing_criteria = [] for b in range(B): obscrit = GaussianObservationCriterion(torch.ones(2), t[b:b + 1], y[b:b + 1]) dyncrit = GaussianDynamicsCriterion(wstd / ystd * torch.ones(3), t[b:b + 1]) smoothing_criteria.append(GroupSOSCriterion([obscrit, dyncrit])) smooth_solver_kwargs = {'verbose': 0, 'tr_rho': 0.001} # specify learning criteria learning_criteria = [GaussianDynamicsCriterion(torch.ones(3), t)] learning_params = [params] learning_opts = ['scipy_minimize'] learner_opt_kwargs = {'method': 'Nelder-Mead', 'tr_rho': 0.01} # instantiate CEEM timer = {'start_time':timeit.default_timer()} def ecb(epoch): logger.logkv('test/rho', float(system._rho)) logger.logkv('test/sigma', float(system._sigma)) logger.logkv('test/beta', float(system._beta)) logger.logkv('test/rho_pcterr_log10', float(torch.log10((true_system._rho - system._rho).abs() / true_system._rho))) logger.logkv( 'test/sigma_pcterr_log10', float(torch.log10((true_system._sigma - system._sigma).abs() / true_system._sigma))) logger.logkv( 'test/beta_pcterr_log10', float(torch.log10((true_system._beta - system._beta).abs() / true_system._beta))) logger.logkv('time/epochtime', timeit.default_timer() - timer['start_time']) timer['start_time'] = timeit.default_timer() return epoch_callbacks = [ecb] class Last10Errors: def __init__(self): return last_10_errors = Last10Errors last_10_errors._arr = [] def tcb(epoch): params = list(system.parameters()) vparams = parameters_to_vector(params) error = (vparams - true_params).norm().item() last_10_errors._arr.append(float(error)) logger.logkv('test/log10_error', np.log10(error)) if len(last_10_errors._arr) > 10: last_10_errors._arr = last_10_errors._arr[-10:] l10err = torch.tensor(last_10_errors._arr) convcrit = float((l10err.min() - l10err.max()).abs()) logger.logkv('test/log10_convcrit', np.log10(convcrit)) if convcrit < 1e-4: return True return False termination_callback = tcb ceem = CEEM(smoothing_criteria, learning_criteria, learning_params, learning_opts, epoch_callbacks, termination_callback) # run CEEM x0 = torch.zeros_like(x) ecb(-1) logger.dumpkvs() ceem.train(xs=x0, sys=system, nepochs=100, smooth_solver_kwargs=smooth_solver_kwargs, learner_opt_kwargs=learner_opt_kwargs) return float(system._sigma), float(system._rho), float(system._beta)
def train(seed, logdir, sys_seed, k, b): from mpl_toolkits.mplot3d import Axes3D import matplotlib.pyplot as plt ystd = 0.01 # ystd = 0. torch.set_default_dtype(torch.float64) logger.setup(logdir, action='d') N = 128 n = 3 * k B = b true_system = default_lorenz_system(k, obsdif=2) utils.set_rng_seed(sys_seed) xdim = true_system.xdim ydim = true_system.ydim dt = true_system._dt x0mean = torch.tensor([[-6] * k + [-6] * k + [24.] * k]).unsqueeze(0) # simulate true_dynamics over IC distribution x_test = x0mean.repeat(1024, 1, 1) x_test += 5.0 * torch.randn_like(x_test) x_test = x_test.detach() t_test = torch.zeros(1024, 1) tgt_test = true_system.step_derivs(t_test, x_test).detach() Q = (0.01**2) * torch.eye(true_system.xdim) R = (ystd**2) * torch.eye(true_system.ydim) Px0 = 2.5**2 * torch.eye(true_system.xdim) ## simulate the true system xs = [x0mean.repeat(B, 1, 1)] xs[0] += 2.5 * torch.randn_like(xs[0]) with torch.no_grad(): for t in range(N - 1): xs.append(true_system.step(torch.tensor([0.] * B), xs[-1])) xs = torch.cat(xs, dim=1) fig = plt.figure() for b in range(B): ax = fig.add_subplot(int(np.ceil(B / 2.)), 2, b + 1, projection='3d') for k_ in range(k): plot3d(plt.gca(), xs[b, :, k_], xs[b, :, k + k_], xs[b, :, 2 * k + k_], linestyle='--', alpha=0.5) plt.savefig(os.path.join(logger.get_dir(), 'figs/traj_%d.png' % b), dpi=300) # plt.show() plt.close() t = torch.tensor(range(N)).unsqueeze(0).expand(B, -1).to( torch.get_default_dtype()) y = true_system.observe(t, xs).detach() # seed for real now utils.set_rng_seed(seed) y += ystd * torch.randn_like(y) # prep system system = deepcopy(true_system) true_params = parameters_to_vector(true_system.parameters()) utils.set_rng_seed(seed) params = true_params * ( (torch.rand_like(true_params) - 0.5) / 5. + 1.) # within 10% vector_to_parameters(params, system.parameters()) params = list(system.parameters()) Np = 100 fapf = faPF(Np, system, Q, R, Px0) timer = {'start_time': timeit.default_timer()} def ecb(epoch): logger.logkv('time/epoch', epoch) params = list(system.parameters()) vparams = parameters_to_vector(params) error = (vparams - true_params).norm().item() logger.logkv('test/log10_paramerror', np.log10(error)) logger.logkv('time/epochtime', timeit.default_timer() - timer['start_time']) timer['start_time'] = timeit.default_timer() with torch.no_grad(): tgt_test_pr = system.step_derivs(t_test, x_test) error = float(torch.nn.functional.mse_loss(tgt_test_pr, tgt_test)) logger.logkv('test/log10_error', np.log10(error)) return epoch_callbacks = [ecb] ecb(-1) logger.dumpkvs() trainer = SAEMTrainer( fapf, y, # gamma_sched=lambda x: HarmonicDecayScheduler(x, a=50.), gamma_sched=lambda x: 0.2, max_k=5000, xlen_cutoff=15, ) trainer.train(params, callbacks=epoch_callbacks) return