flatquad_landing_experiment.py

#!/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)