[
  {
    "path": "LICENSE",
    "content": "MIT License\n\nCopyright (c) 2018 Robotics and Biology Laboratory, TU Berlin\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
  },
  {
    "path": "README.md",
    "content": "Differentiable Particle Filters\n==================================================\n\nContact\n------------------\n\nRico Jonschkowski (rico.jonschkowski@tu-berlin.de)\n\n\nIntroduction\n------------\n\nThis repository contains our source code for differentiable particle filters (DPFs) described in the paper \"Differentiable Particle Filters: End-to-End Learning with Algorithmic Priors\" (Jonschkowski et al. 2018). This implementation is meant to complement the paper. Our goal is to enable others to reproduce our results and to build on our research by reusing our code. We do not include extensive explanations here -- these can be found in our paper https://arxiv.org/pdf/1805.11122.pdf.\n\nIf you are using this implementation in your research, please consider giving credit by citing our paper:\n\n    @article{jonschkowski18,\n        title = {{Differentiable Particle Filters: End-to-End Learning with Algorithmic Priors}},\n\t    booktitle = {{Proceedings of Robotics: Science and Systems (RSS)}},\n\t    author = {Jonschkowski, Rico and Rastogi, Divyam and Brock, Oliver},\n\t    year = {2018},\n    }\n\nDependencies\n------------\n\nOur code builds on python3 and the following libraries. *It is important that the libraries are installed for python3 not python2.*\n\nnumpy\n\n    sudo apt-get install python3-numpy\n\nmatplotlib \n\n    sudo apt-get install python3-matplotlib\n\nTensorFlow --> https://www.tensorflow.org/install/, e.g. \n\n    pip3 install tensorflow\n\nSonnet --> https://github.com/deepmind/sonnet, e.g. \n\n    pip3 install dm-sonnet\n\n\nSetup\n-----\n\nThe setup script **downloads the data** for training and testing (~2.5GB) and **creates additional folders** (for logging etc.). To perform these steps, simply run the following commands in the main folder of the repository:\n\n    chmod +x setup.sh\n    ./setup.sh\n\nUsage\n-----\n\nAfter all dependencies are installed and setup is done, there is one more thing which needs to be done every time a new shell is opened before the code can be run. In the main repository folder, you need run the following command to append the parent directory to the PYTHONPATH. *Alternatively, you can import the project into the PyCharm IDE and and run `experiments/simple.py` from there. The need for running this command comes from how PyCharm handles relative imports and relative paths.*\n\n\texport PYTHONPATH=\"${PYTHONPATH}:../\"\n\nAfter this is done, you can train and test a differentiable particle filter for global localization in maze 1 by running the following commands in the main folder:\n\n    cd experiments; python3 simple.py; cd ..\n    \nThis command will first train the different models (motion model, observation likelihood estimator, and particle proposer) individually and then train them jointly end-to-end. The command line output will show the current losses on training and validation data (mean +- standard error), where \">>\" indicates a new lowest validation loss. Training will stop if the best validation loss has not decreased for a while (e.g. 200 epochs). You should see something like this (the different plots are generated one after another):\n\n![Screenshot](https://raw.githubusercontent.com/tu-rbo/differentiable-particle-filters/master/screenshot.png)\n\nAs next steps, you can play around with the hyperparameters in `utils/exp_utils.py`, go through the differentiable particle filter code in `methods/dpf.py`, and run other experiments, e.g. applying the filter to the KITTI visual odometry task by running the following command (if your computer has enough memory :D).\n\n    cd experiments; python3 cross_validation_kitti.py; cd ..\n"
  },
  {
    "path": "experiments/__init__.py",
    "content": ""
  },
  {
    "path": "experiments/cross_evaluation.py",
    "content": "import pickle\nimport os\nimport numpy as np\nfrom experiments.distributed_experiments import run_experiment, tracking_exp, planner_agent_exp, learning_curve_exp, noise_test_exp\n\ndef cross(logfile, cross_exp, exp_name='cr'):\n    # load data, choose correct task, method, num_episodes, noise_cond, seq_len\n    with open(logfile, 'rb') as f:\n        log = pickle.load(f)\n    model_path = '../models/' + log['exp_params'][0]['model_path'].split('/models/')[-1] # ['exp_params']['model_path]\n    print(model_path)\n\n    # these are actually already lists so we can pass them on directly\n    task = [log['exp_params'][0]['task']]\n    method = [log['exp_params'][0]['method']]\n    num_episodes = [log['exp_params'][0]['num_episodes']]\n    num_episodes = [log['exp_params'][0]['num_episodes']]\n\n    # define experiment you want to run\n    get_experiment_params, get_train_data_and_eval_iterator = cross_exp('../', exp_name=exp_name, id_extra='',\n                                                                           tasks=task, methods=method, episodes=num_episodes,\n                                                                           num_test_episodes=1000,\n                                                                           run=False)\n\n    run_experiment(get_experiment_params, get_train_data_and_eval_iterator, base_path='../', exp_name=exp_name, id_extra='', load_from_model_path=model_path)\n\n\ndef swapmodels(logfiles, noise_conds, exp_name='swap', flipmodules=False):\n\n    # expect logfiles to be a dict with two keys that match the noise conditions in noise_test,\n    # e.g. {'odom5_imgTG': [log1, log2], 'odom20_imgTG': [log1, log2]}\n    # noise_conds should be a list of the two conditions\n\n    # noise_conds = logfiles.keys()\n    model_paths = dict()\n    for c in noise_conds:\n        model_paths[c] = []\n        for i, logfile in enumerate(logfiles[c]):\n            with open(logfile, 'rb') as f:\n                log = pickle.load(f)\n                model_paths[c].append('../models/' + log['exp_params'][0]['model_path'].split('/models/')[-1])\n                # should be the same for all logfiles, not checked here\n                task = [log['exp_params'][0]['task']]\n                method = [log['exp_params'][0]['method']]\n                num_episodes = [log['exp_params'][0]['num_episodes']]\n\n\n\n    get_experiment_params, get_train_data_and_eval_iterator = noise_test_exp('../', exp_name=exp_name, id_extra='',\n                                                                           tasks=task, methods=method, episodes=num_episodes,\n                                                                           noise_conds=noise_conds,\n                                                                           num_test_episodes=1000,\n                                                                           run=False)\n\n    modules0 = ('mo_noise_generator', 'mo_transition_model')\n    modules1 = ('encoder', 'obs_like_estimator', 'particle_proposer')\n\n    if flipmodules:\n        modules0, modules1 = modules1, modules0\n\n    for variant, (path, module) in {\n        'orig_'+noise_conds[0]: (model_paths[noise_conds[0]][0], None),\n        '%s_%s' % (noise_conds[0], noise_conds[0]): (model_paths[noise_conds[0]], [modules0, modules1]),\n        '%s_%s' % (noise_conds[0], noise_conds[1]): ([model_paths[noise_conds[0]][0], model_paths[noise_conds[1]][0]], [modules0, modules1]),\n        'orig_'+noise_conds[1]: (model_paths[noise_conds[1]][0], None),\n        '%s_%s' % (noise_conds[1], noise_conds[1]): (model_paths[noise_conds[1]], [modules0, modules1]),\n        '%s_%s' % (noise_conds[1], noise_conds[0]): ([model_paths[noise_conds[1]][0], model_paths[noise_conds[0]][0]], [modules0, modules1]),\n        }.items():\n        print('!!! %s %s %s' % (variant, path, module))\n        run_experiment(get_experiment_params, get_train_data_and_eval_iterator, base_path='../', exp_name=exp_name+'/'+variant, id_extra='', load_from_model_path=path, load_modules=module)\n\n\ndef get_all_logs(path, file_ending):\n    return [os.path.join(path, filename) for filename in os.listdir(path)\n              if os.path.isfile(os.path.join(path, filename))\n              # and filename.endswith(file_ending)]\n              and file_ending in filename]\n\ndef cross_lc2pl(method):\n    # for f in get_all_logs('../log/lc', 'nav02_'+method+'_1000'):\n    for f in get_all_logs('../log/lc', 'nav02_'+method+'_'):\n        cross(f, learning_curve_exp, 'lc2lc1')\n        cross(f, planner_agent_exp, 'lc2pl1')\n\ndef cross_pl2lc(method):\n    # for f in get_all_logs('../log/pl', 'nav02_'+method+'_1000'):\n    for f in get_all_logs('../log/pl', 'nav02_'+method):\n        cross(f, learning_curve_exp, 'pl2lc1')\n        cross(f, planner_agent_exp, 'pl2pl1')\n\ndef cross_mx(method):\n    for f in get_all_logs('../log/mx', 'nav02_'+method+'_1000'):\n    # for f in get_all_logs('../log/mx', 'nav02_'+method):\n        cross(f, learning_curve_exp, 'mx2lc')\n        cross(f, planner_agent_exp, 'mx2pl')\n\ndef swap_motion(method):\n    noise_conds = ['odom5_imgTG', 'odom10_imgTG']\n    logs = dict()\n    for c in noise_conds:\n        logs[c] = [f for f in get_all_logs('../log/nt', 'nav02_'+method+'_1000_'+c)]\n        i, j = np.random.choice(len(logs[c]), 2, False)\n        logs[c] = [logs[c][i], logs[c][j]]\n    swapmodels(logs, noise_conds, 'swapmo')\n\ndef swap_measurement(method):\n    noise_conds = ['odom10_imgG', 'odom10_imgTG']\n    logs = dict()\n    for c in noise_conds:\n        logs[c] = [f for f in get_all_logs('../log/nt', 'nav02_'+method+'_1000_'+c)][:2]\n        i, j = np.random.choice(len(logs[c]), 2, False)\n        logs[c] = [logs[c][i], logs[c][j]]\n    swapmodels(logs, noise_conds, 'swapme', flipmodules=True)\n\n# if __name__ == '__main__':\n"
  },
  {
    "path": "experiments/cross_validation_kitti.py",
    "content": "from utils.data_utils_kitti import load_kitti_sequences\nimport tensorflow as tf\nfrom methods.dpf_kitti import DPF\nfrom utils.exp_utils_kitti import get_default_hyperparams\nimport numpy as np\n\ndef run_cross_validation(i):\n\n    print('RUNNING CROSS VALIDATION TRAINING FOR TESTING {}'.format(i))\n\n    model_path = '../models/tmp/cross_validation_ind_e2e/model_trained_ex_{}'.format(i)\n\n    training_subsequences = [j for j in range(11) if j not in [i]]\n\n    # Load all subsequences\n    data = load_kitti_sequences(training_subsequences)\n\n    # Assign weights to all subsequences based on the length of the subsequence\n    weights = np.zeros((data['seq_num'].shape[0],))\n    weights[0] = data['seq_num'][0]\n    weights[1:] = data['seq_num'][1:] - data['seq_num'][:-1]\n    weights = weights/data['seq_num'][-1]\n    data['weights'] = weights\n\n    # reset tensorflow graph\n    tf.reset_default_graph()\n\n    # instantiate method\n    hyperparams = get_default_hyperparams()\n    hyperparams['train']['split_ratio'] = 0.9  # -> 18/2 split\n\n    method = DPF(**hyperparams['global'])\n\n    with tf.Session() as session:\n        # train method and save result in model_path\n        method.fit(session, data, model_path, plot=False, **hyperparams['train'])\n\nif __name__ == '__main__':\n    for i in range(11):\n        run_cross_validation(i)\n"
  },
  {
    "path": "experiments/distributed_experiments.py",
    "content": "import tensorflow as tf\nimport numpy as np\nimport pickle\nimport os\nimport time\nimport itertools\n\nfrom utils.exp_utils import get_default_hyperparams, add_to_log, exp_variables_to_name, print_msg_and_dict, sample_exp_variables\nfrom utils.data_utils import load_data, noisify_data_condition, compute_staticstics, make_batch_iterator, reduce_data, shuffle_data\nfrom utils.method_utils import compute_sq_distance\nfrom methods.dpf import DPF\nfrom methods.rnn import RNN\nfrom methods.odom import OdometryBaseline\n\ndef meta_exp(base_path, id_extra):\n\n    min_counts = []\n    exp_names = ['lc', 'tr', 'nt', 'ab', 'pl', 'mx']\n    funcs = [learning_curve_exp, tracking_exp, noise_test_exp, ablation_test_exp, planner_agent_exp, mix_agent_exp]\n\n    for exp_name, f in zip(exp_names, funcs):\n\n        get_experiment_params, get_train_data_and_eval_iterator = f(base_path, run=False)\n        # check progress for that experiment\n        log_base_path = os.path.join(base_path, 'log', exp_name)\n        min_counts.append(get_experiment_params(log_base_path)[-1])\n        print('Experiment', exp_name, 'has min_count', min_counts[-1])\n\n    min_min_count = np.min(min_counts)\n    sample_list = []\n    for i in range(len(exp_names)):\n        sample_list += [i] * max(0, (min_min_count + 2) - min_counts[i]) * (3 if 'lc' in exp_names[i] else 1)\n    if sample_list == []:\n        sample_list = range(len(exp_names))\n\n    i = sample_list[np.random.choice(len(sample_list))]\n    exp_name = exp_names[i]\n    f = funcs[i]\n\n    print('--> META EXPERIMENT CHOOSES ', exp_name)\n    f(base_path, exp_name, id_extra)\n\n\ndef learning_curve_exp(base_path='', exp_name='lc', id_extra='', tracking=False,\n                       tasks=('nav01', 'nav02', 'nav03'),\n                       methods=('pf_ind', 'pf_e2e', 'pf_ind_e2e', 'lstm', 'ff'),\n                       episodes = (16, 32, 64, 125, 250, 500, 1000), data_dir='100s',\n                       num_test_episodes=20000, run=True):\n\n    def get_experiment_params(base_path):\n\n        variables, min_count = sample_exp_variables(base_path, [tasks, methods, episodes])\n        task, method, num_episodes = variables\n\n        exp_params = {\n            'exp': exp_name,\n            'task': task,\n            'method': method,\n            'num_episodes': num_episodes,\n            'noise_condition': 'odom10_imgTG',\n            'tracking': tracking,\n            'computer': os.uname()[1],\n            'num_test_episodes': num_test_episodes,\n            'eval_batch_size': 16,\n            'eval_seq_len': 50,\n            'data_dir': data_dir,\n            'file_ending': exp_variables_to_name(variables)\n        }\n\n        # match sequence length to task\n        if exp_params['task'] == 'nav01':\n            exp_params['seq_len'] = 20\n        elif exp_params['task'] == 'nav02':\n            exp_params['seq_len'] = 20\n        elif exp_params['task'] == 'nav03':\n            exp_params['seq_len'] = 30\n\n        return exp_params, get_default_hyperparams(), min_count\n\n    def get_train_data_and_eval_iterator(data, exp_params):\n\n        # noisify\n        for k in ['train', 'test']:\n            data[k] = noisify_data_condition(data[k], exp_params['noise_condition'])\n\n        # form batches\n        eval_batch_iterators = {k: make_batch_iterator(data[k], batch_size=exp_params['eval_batch_size'], seq_len=exp_params['eval_seq_len']) for k in ['test']}\n\n        return data['train'], eval_batch_iterators\n\n    if run:\n        # run an experiment with these two functions\n        return run_experiment(get_experiment_params, get_train_data_and_eval_iterator, base_path, exp_name, id_extra)\n    else:\n        return get_experiment_params, get_train_data_and_eval_iterator\n\n\ndef tracking_exp(base_path='', exp_name='tr', id_extra='',\n                       tasks=('nav02',),\n                       methods=('pf_ind', 'pf_e2e', 'pf_ind_e2e', 'lstm', 'odom'),\n                       episodes = (16, 32, 64, 125, 250, 500, 1000), data_dir='100s', num_test_episodes=20000, run=True):\n\n    return learning_curve_exp(base_path, exp_name, id_extra, True, tasks, methods, episodes, data_dir, num_test_episodes, run)\n\ndef planner_agent_exp(base_path='', exp_name='pl', id_extra='',\n                       tasks=('nav02',),\n                       methods=('pf_ind', 'pf_e2e', 'pf_ind_e2e', 'lstm'),\n                       episodes = (16, 32, 64, 125, 250, 500, 1000), data_dir='100s_astar', num_test_episodes=20000, run=True):\n\n    return learning_curve_exp(base_path, exp_name, id_extra, False, tasks, methods, episodes, data_dir, num_test_episodes, run)\n\n\ndef mix_agent_exp(base_path='', exp_name='mx', id_extra='',\n                       tasks=('nav02',),\n                       methods=('pf_ind', 'pf_e2e', 'pf_ind_e2e', 'lstm'),\n                       episodes = (16, 32, 64, 125, 250, 500, 1000), data_dir='100s_mix', num_test_episodes=1000, run=True):\n\n    return learning_curve_exp(base_path, exp_name, id_extra, False, tasks, methods, episodes, data_dir, num_test_episodes, run)\n\n\ndef noise_test_exp(base_path='', exp_name='nt', id_extra='', tracking=False,\n                   tasks=('nav02',),\n                   methods=('pf_ind', 'pf_e2e', 'pf_ind_e2e', 'lstm'),\n                   episodes = (16, 125, 1000),\n                   noise_conds=('odom0_imgTG', 'odom5_imgTG', 'odom10_imgTG', 'odom20_imgTG', 'odomX_imgTG',\n                                  'odom10_imgC', 'odom10_imgG', 'odom10_imgT', 'odom10_imgX'),\n                   data_dir='100s', num_test_episodes=1000,\n                   run=True,\n                   ):\n\n    def get_experiment_params(base_path):\n\n        variables, min_count = sample_exp_variables(base_path, [tasks, methods, episodes, noise_conds])\n        task, method, num_episodes, noise_cond = variables\n\n        exp_params = {\n            'exp': exp_name,\n            'task': task,\n            'method': method,\n            'num_episodes': num_episodes,\n            'noise_condition': noise_cond,\n            'tracking': tracking,\n            'computer': os.uname()[1],\n            'num_test_episodes': num_test_episodes,\n            'eval_batch_size': 16,\n            'eval_seq_len': 50,\n            'data_dir': data_dir,\n            'file_ending': exp_variables_to_name(variables),\n            'seq_len': 20,\n        }\n\n        return exp_params, get_default_hyperparams(), min_count\n\n    def get_train_data_and_eval_iterator(data, exp_params):\n\n        # noisify training data according to sampled noise condition\n        data['train'] = noisify_data_condition(data['train'], exp_params['noise_condition'])\n\n        # create eval batch iterators for every noise condition\n        eval_batch_iterators = dict()\n        for condition in noise_conds:\n            key = 'test' + '_' + condition\n            data[key] = noisify_data_condition(data['test'], condition)\n            eval_batch_iterators[key] = make_batch_iterator(data[key], batch_size=exp_params['eval_batch_size'], seq_len=exp_params['eval_seq_len'])\n\n        return data['train'], eval_batch_iterators\n\n    if run:\n        # run an experiment with these two functions\n        return run_experiment(get_experiment_params, get_train_data_and_eval_iterator, base_path, exp_name, id_extra)\n    else:\n        return get_experiment_params, get_train_data_and_eval_iterator\n\n\ndef ablation_test_exp(base_path='', exp_name='ab', id_extra='', tracking=False,\n                   tasks=('nav02',),\n                   methods=('pf_ind', 'pf_e2e', 'pf_ind_e2e'),\n                   episodes=(16, 125, 1000),\n                   ab_conds=('full', 'learn_odom', 'no_proposer', 'no_inject'),\n                   data_dir='100s',\n                   run=True\n                   ):\n    def get_experiment_params(base_path):\n        variables, min_count = sample_exp_variables(base_path, [tasks, methods, episodes, ab_conds])\n        task, method, num_episodes, ab_cond = variables\n\n        exp_params = {\n            'exp': exp_name,\n            'task': task,\n            'method': method,\n            'num_episodes': num_episodes,\n            'noise_condition': 'odom10_imgTG',\n            'tracking': tracking,\n            'computer': os.uname()[1],\n            'num_test_episodes': 20000,\n            'eval_batch_size': 16,\n            'eval_seq_len': 50,\n            'data_dir': data_dir,\n            'file_ending': exp_variables_to_name(variables),\n            'seq_len': 20,\n        }\n\n        hyper_params = get_default_hyperparams()\n        if ab_cond == 'learn_odom':\n            hyper_params['global']['learn_odom'] = True\n        elif ab_cond == 'no_proposer':\n            hyper_params['global']['use_proposer'] = False\n            hyper_params['global']['propose_ratio'] = 0.0\n        elif ab_cond == 'no_inject':\n            hyper_params['global']['propose_ratio'] = 0.0\n\n        return exp_params, hyper_params, min_count\n\n    def get_train_data_and_eval_iterator(data, exp_params):\n\n        # noisify\n        for k in ['train', 'test']:\n            data[k] = noisify_data_condition(data[k], exp_params['noise_condition'])\n\n        # form batches\n        eval_batch_iterators = {k: make_batch_iterator(data[k], batch_size=exp_params['eval_batch_size'], seq_len=exp_params['eval_seq_len']) for k in ['test']}\n\n        return data['train'], eval_batch_iterators\n\n    if run:\n        # run an experiment with these two functions\n        return run_experiment(get_experiment_params, get_train_data_and_eval_iterator, base_path, exp_name, id_extra)\n    else:\n        return get_experiment_params, get_train_data_and_eval_iterator\n\n\ndef run_experiment(get_experiment_params, get_train_data_and_eval_iterator, base_path, exp_name, id_extra='',\n                   load_from_model_path=None, load_modules=None):\n\n    # construct base paths\n    log_base_path = os.path.join(base_path, 'log', exp_name)\n    if not os.path.exists(log_base_path):\n        os.makedirs(log_base_path)\n    model_base_path = os.path.join(base_path, 'models', exp_name)\n\n    # sample experiment parameters by checking the log for what is most urgent right now\n    exp_params, hyperparams, min_count = get_experiment_params(log_base_path)\n    data_path = os.path.join(base_path, 'data', exp_params['data_dir'])\n\n    id = exp_params['id'] = time.strftime('%Y-%m-%d_%H:%M:%S_') + exp_params['computer'] + str(id_extra) + '_' + exp_params['file_ending']\n    log_path = os.path.join(log_base_path, id)\n    model_path = exp_params['model_path'] = os.path.join(model_base_path, id)\n    if not os.path.exists(model_path):\n        os.makedirs(model_path)\n\n    # load data\n    data = {k: load_data(data_path=data_path, filename=exp_params['task'] + '_' + k) for k in ['train', 'test']}\n    means, stds, state_step_sizes, state_mins, state_maxs = compute_staticstics(data['train'])\n\n    data['train'] = shuffle_data(data['train'])\n    data['train'] = reduce_data(data['train'], exp_params['num_episodes'])\n\n    data['train'], eval_batch_iterators = get_train_data_and_eval_iterator(data, exp_params)\n\n    log = dict()\n\n    # SET THINGS UP\n    tf.reset_default_graph()\n\n    print_msg_and_dict('STARTING EXPERIMENT', exp_params)\n\n    hyperparams['global']['init_with_true_state'] = exp_params['tracking']\n\n    if 'pf' in exp_params['method']:\n        method = DPF(**hyperparams['global'])\n        hyperparams['train']['train_e2e'] = 'e2e' in exp_params['method']\n        hyperparams['train']['train_individually'] = 'ind' in exp_params['method']\n\n    elif 'lstm' in exp_params['method']:\n        method = RNN(**hyperparams['global'])\n    elif 'ff' in exp_params['method']:\n        method = RNN(model='ff', **hyperparams)\n    elif 'odom' in exp_params['method']:\n        method = OdometryBaseline(**hyperparams)\n    else:\n        print('I DONT KNOW THIS METHOD', exp_params['method'])\n\n    with tf.Session() as session:\n\n        t0 = time.time()\n        if load_from_model_path is None:\n            training_log = method.fit(session, data['train'], model_path, **hyperparams['train'])\n        elif type(load_from_model_path) == type([]):\n            for i, (path, modules) in enumerate(zip(load_from_model_path, load_modules)):\n                print('Loading %s from %s' % (modules, path))\n                method.load(session, path, modules=modules, connect_and_initialize=(i==0))\n            training_log = None\n        else:\n            print('Loading model')\n            if load_modules is None:\n                method.load(session, load_from_model_path)\n            else:\n                method.load(session, load_from_model_path, modules=load_modules)\n            training_log = None\n\n        t1 = time.time()\n        add_to_log(log, {'training_duration': t1 - t0})\n\n        print_msg_and_dict('RESULTS after {}s'.format(log['training_duration'][-1]), exp_params)\n\n        for k in sorted(eval_batch_iterators.keys()):\n            results = {'mse': []}\n            result_hist = dict()\n            for i in range(0, exp_params['eval_seq_len'], 10):\n                result_hist[i] = np.zeros(100)\n\n            for eval_batch in eval_batch_iterators[k]:\n\n                predicted_states = method.predict(session, eval_batch, **hyperparams['test'])\n                squared_errors = compute_sq_distance(predicted_states, eval_batch['s'], state_step_sizes)\n\n                for i in result_hist.keys():\n                    result_hist[i] += np.histogram(squared_errors[:, i], bins=100, range=[0.0, 10.0])[0]\n                results['mse'].append(np.mean(squared_errors, axis=0))\n                if len(results['mse']) * exp_params['eval_batch_size'] >= exp_params['num_test_episodes']:\n                   break\n\n            for i in result_hist.keys():\n                result_hist[i] /= len(results['mse']) * exp_params['eval_batch_size']\n            mse = np.stack(results['mse'], axis=0)\n\n            add_to_log(log, {k + '_hist': result_hist,\n                            k + '_mse': np.mean(mse, axis=0),\n                            k + '_mse_se':  np.std(mse, ddof=1, axis=0) / np.sqrt(len(mse))})\n            for i in range(0, len(log[k+'_mse'][-1]), 5):\n                print('{:>10} step {} !! mse: {:.4f}+-{:.4f}'.format(k, i, log[k+'_mse'][-1][i], log[k+'_mse_se'][-1][i]))\n\n    add_to_log(log, {'hyper_params': hyperparams})\n    add_to_log(log, {'exp_params': exp_params})\n    add_to_log(log, {'training': training_log})\n\n    # save result\n    print('Saved log as ', log_path)\n    with open(log_path, 'wb') as f:  # Just use 'w' mode in 3.x\n        pickle.dump(log, f)\n"
  },
  {
    "path": "experiments/evaluation_kitti.py",
    "content": "import tensorflow as tf\n\nfrom methods.dpf_kitti import DPF\nfrom methods.odom import OdometryBaseline\nfrom utils.data_utils_kitti import load_data, noisyfy_data, make_batch_iterator, remove_state, split_data, load_kitti_sequences, make_batch_iterator_for_evaluation, wrap_angle, plot_video\nfrom utils.exp_utils_kitti import get_default_hyperparams\nimport matplotlib.pyplot as plt\nimport numpy as np\n\ndef get_evaluation_stats(model_path='../models/tmp/', test_trajectories=[11], seq_lengths = [100, 200, 400, 800], plot_results=False):\n\n    data = load_kitti_sequences(test_trajectories)\n\n    # reset tensorflow graph\n    tf.reset_default_graph()\n\n    # instantiate method\n    hyperparams = get_default_hyperparams()\n    method = DPF(**hyperparams['global'])\n\n    with tf.Session() as session:\n\n        # load method and apply to new data\n        method.load(session, model_path)\n\n        errors = dict()\n\n        for i, test_traj in enumerate(test_trajectories):\n\n            s_test_traj = data['s'][0:data['seq_num'][i*2]]  # take care of duplicated trajectories (left and right camera)\n            distance = compute_distance_for_trajectory(s_test_traj)\n            errors[test_traj] = dict()\n\n            for seq_len in seq_lengths:\n\n                errors[test_traj][seq_len] = {'trans': [], 'rot': []}\n\n                for start_step in range(0, distance.shape[0], 1):\n\n                    end_step, dist = find_end_step(distance, start_step, seq_len, use_meters=False)  #--> Put use_meters = True for official KITTI benchmark results\n\n                    if end_step == -1:\n                        continue\n\n                    # test_batch_iterator = make_batch_iterator(test_data, seq_len=50)\n                    test_batch_iterator = make_batch_iterator_for_evaluation(data, start_step, trajectory=i, batch_size=1, seq_len=end_step-start_step)\n\n                    batch = next(test_batch_iterator)\n                    batch_input = remove_state(batch, provide_initial_state=True)\n\n                    prediction, particle_list, particle_prob_list = method.predict(session, batch_input, return_particles=True)\n                    error_x = batch['s'][0, -1, 0] - prediction[0, -1, 0]\n                    error_y = batch['s'][0, -1, 1] - prediction[0, -1, 1]\n                    error_trans = np.sqrt(error_x ** 2 + error_y ** 2) / dist\n                    error_rot = abs(wrap_angle(batch['s'][0, -1, 2] - prediction[0, -1, 2]))/dist * 180 / np.pi\n\n                    errors[test_traj][seq_len]['trans'].append(error_trans)\n                    errors[test_traj][seq_len]['rot'].append(error_rot)\n\n                    if plot_results:\n\n                        dim_names = ['pos', 'theta', 'vel_f', 'vel_th']\n                        fig = plt.figure()\n                        ax1 = fig.add_subplot(221)\n                        ax2 = fig.add_subplot(222)\n                        ax3 = fig.add_subplot(223)\n                        ax4 = fig.add_subplot(224)\n\n                        for t in range(particle_list.shape[1]):\n                            dim = 0\n                            ax1.scatter(particle_list[0, t, :, dim], particle_list[0, t, :, dim+1], c=particle_prob_list[0, t, :], cmap='viridis_r', marker='o', s=15, alpha=0.1,\n                                                linewidths=0.05,\n                                                vmin=0.0,\n                                                vmax=0.02)\n\n                            ax1.plot([prediction[0, t, dim]], [prediction[0, t, dim+1]], 'o', markerfacecolor='None', markeredgecolor='b',\n                                             markersize=0.5)\n\n                            ax1.plot([batch['s'][0, t, dim]], [batch['s'][0, t, dim+1]], '+', markerfacecolor='None', markeredgecolor='r',\n                                             markersize=0.5)\n\n                            ax1.set_aspect('equal')\n\n                            dim = 2\n                            ax2.scatter(t * np.ones_like(particle_list[0, t, :, dim]), particle_list[0, t, :, dim], c=particle_prob_list[0, t, :], cmap='viridis_r', marker='o', s=15, alpha=0.1,\n                                                linewidths=0.05,\n                                                vmin=0.0,\n                                                vmax=0.02)\n                                                #np.max(\n                                                    #s_add_probs_list[s, i, :, 0]))  # , vmin=-1/filter.num_particles,)\n                            current_state = prediction[0, t, dim]\n                            ax2.plot([t], [current_state], 'o', markerfacecolor='None', markeredgecolor='k',\n                                             markersize=2.5)\n                            true = batch['s'][0, t, dim]\n                            ax2.plot([t], [true], '+', markerfacecolor='None', markeredgecolor='r',\n                                             markersize=2.5)\n\n                            dim = 3\n                            ax3.scatter(t * np.ones_like(particle_list[0, t, :, dim]), particle_list[0, t, :, dim], c=particle_prob_list[0, t, :], cmap='viridis_r', marker='o', s=15, alpha=0.1,\n                                                linewidths=0.05,\n                                                vmin=0.0,\n                                                vmax=0.02)\n                                                #np.max(\n                                                    #s_add_probs_list[s, i, :, 0]))  # , vmin=-1/filter.num_particles,)\n                            current_state = prediction[0, t, dim]\n                            ax3.plot([t], [current_state], 'o', markerfacecolor='None', markeredgecolor='k',\n                                             markersize=2.5)\n                            true = batch['s'][0, t, dim]\n                            ax3.plot([t], [true], '+', markerfacecolor='None', markeredgecolor='r',\n                                             markersize=2.5)\n\n                            dim = 4\n                            ax4.scatter(t * np.ones_like(particle_list[0, t, :, dim]), particle_list[0, t, :, dim], c=particle_prob_list[0, t, :], cmap='viridis_r', marker='o', s=15, alpha=0.1,\n                                                linewidths=0.05,\n                                                vmin=0.0,\n                                                vmax=0.02)\n\n                            current_state = prediction[0, t, dim]\n                            ax4.plot([t], [current_state], 'o', markerfacecolor='None', markeredgecolor='k',\n                                             markersize=2.5)\n                            true = batch['s'][0, t, dim]\n                            ax4.plot([t], [true], '+', markerfacecolor='None', markeredgecolor='r',\n                                             markersize=2.5)\n\n                        plt.pause(0.05)\n\n                        ax1.set_title(dim_names[0])\n                        ax2.set_title(dim_names[1])\n                        ax3.set_title(dim_names[2])\n                        ax4.set_title(dim_names[3])\n\n    return errors\n\n\n\ndef compute_distance_for_trajectory(s):\n\n    # for ii in range(len(output_oxts_file)):\n    distance = [0]\n    for i in range(1, s.shape[0]):\n        diff_x = s[i, 0, 0] - s[i-1, 0, 0]\n        diff_y = s[i, 0, 1] - s[i-1, 0, 1]\n        dist = distance[-1] + np.sqrt(diff_x ** 2 + diff_y ** 2)\n        distance.append(dist)\n    distance = np.asarray(distance)\n    return distance\n\ndef find_end_step(distance, start_step, length, use_meters=True):\n\n    for i in range(start_step, distance.shape[0]):\n        if (use_meters and distance[i] > (distance[start_step] + length)) or \\\n            (not use_meters and (i - start_step) >= length):\n            return i, distance[i] - distance[start_step]\n    return -1, 0\n\ndef find_all_cross_val_models(model_path):\n    import os\n    models = ([name for name in os.listdir(model_path) if not os.path.isfile(os.path.join(model_path, name))])\n    trajs = [int(name.split('_')[3]) for name in models]\n    return zip(models, trajs)\n\ndef main():\n    plt.ion()\n\n    errors = dict()\n    average_errors = {'trans': {i: [] for i in [100, 200, 400, 800]},\n                      'rot': {i: [] for i in [100, 200, 400, 800]}}\n    model_path = '../models/tmp/cross_validation_ind_e2e/'\n    for model, traj in find_all_cross_val_models(model_path):\n        print('!!! Evaluatng model {} on trajectory {}'.format(model, traj))\n        new_errors = get_evaluation_stats(model_path=model_path+model, test_trajectories=[traj], plot_results=False)\n        errors.update(new_errors)\n        print('')\n        print('Trajectory {}'.format(traj))\n        for seq_len in sorted(errors[traj].keys()):\n            for measure in ['trans', 'rot']:\n                e = errors[traj][seq_len][measure]\n                mean_error = np.mean(e)\n                se_error = np.std(e, ddof=1) / np.sqrt(len(e))\n                average_errors[measure][seq_len].append(mean_error)\n                print('{:>5} error for seq_len {}: {:.4f} +- {:.4f}'.format(measure, seq_len, mean_error, se_error))\n\n        print('Averaged errors:')\n        for measure in ['trans', 'rot']:\n            e_means = []\n            e_ses = []\n            for seq_len in sorted(average_errors[measure].keys()):\n                e = np.array(average_errors[measure][seq_len])\n                e = e[~np.isnan(e)]\n                mean_error = np.mean(e)\n                e_means.append(mean_error)\n                se_error = np.std(e, ddof=1) / np.sqrt(len(e))\n                e_ses.append(se_error)\n                print('{:>5} error for seq_len {}: {:.4f} +- {:.4f}'.format(measure, seq_len, mean_error, se_error))\n            print('{:>5} error averaged over seq_lens: {:.4f} +- {:.4f}'.format(measure, np.mean(e_means), np.std(e_means, ddof=1) / np.sqrt(len(e_means))))\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "experiments/simple.py",
    "content": "import tensorflow as tf\n\nfrom methods.dpf import DPF\nfrom utils.data_utils import load_data, noisyfy_data, make_batch_iterator, remove_state\nfrom utils.exp_utils import get_default_hyperparams\n\n\ndef train_dpf(task='nav01', data_path='../data/100s', model_path='../models/tmp', plot=False):\n\n    # load training data and add noise\n    train_data = load_data(data_path=data_path, filename=task + '_train')\n    noisy_train_data = noisyfy_data(train_data)\n\n    # reset tensorflow graph\n    tf.reset_default_graph()\n\n    # instantiate method\n    hyperparams = get_default_hyperparams()\n    method = DPF(**hyperparams['global'])\n\n    with tf.Session() as session:\n        # train method and save result in model_path\n        method.fit(session, noisy_train_data, model_path, **hyperparams['train'], plot_task=task, plot=plot)\n\n\ndef test_dpf(task='nav01', data_path='../data/100s', model_path='../models/tmp'):\n\n    # load test data\n    test_data = load_data(data_path=data_path, filename=task + '_test')\n    noisy_test_data = noisyfy_data(test_data)\n    test_batch_iterator = make_batch_iterator(noisy_test_data, seq_len=50)\n\n    # reset tensorflow graph\n    tf.reset_default_graph()\n\n    # instantiate method\n    hyperparams = get_default_hyperparams()\n    method = DPF(**hyperparams['global'])\n\n    with tf.Session() as session:\n        # load method and apply to new data\n        method.load(session, model_path)\n        for i in range(10):\n            test_batch = next(test_batch_iterator)\n            test_batch_input = remove_state(test_batch, provide_initial_state=False)\n            result = method.predict(session, test_batch_input, **hyperparams['test'])\n\n\nif __name__ == '__main__':\n    train_dpf(plot=True)\n    test_dpf()\n"
  },
  {
    "path": "methods/__init__.py",
    "content": ""
  },
  {
    "path": "methods/dpf.py",
    "content": "import os\nimport numpy as np\nimport sonnet as snt\nimport tensorflow as tf\nimport matplotlib.pyplot as plt\n\nfrom utils.data_utils import wrap_angle, compute_staticstics, split_data, make_batch_iterator, make_repeating_batch_iterator\nfrom utils.method_utils import atan2, compute_sq_distance\nfrom utils.plotting_utils import plot_maze, show_pause\n\nif tf.__version__ == '1.1.0-rc1' or tf.__version__ == '1.3.0':\n    from tensorflow.python.framework import ops\n    @ops.RegisterGradient(\"FloorMod\")\n    def _mod_grad(op, grad):\n        x, y = op.inputs\n        gz = grad\n        x_grad = gz\n        y_grad = None  # tf.reduce_mean(-(x // y) * gz, axis=[0], keep_dims=True)[0]\n        return x_grad, y_grad\n\n\nclass DPF():\n\n    def __init__(self, init_with_true_state, learn_odom, use_proposer, propose_ratio, proposer_keep_ratio, min_obs_likelihood):\n        \"\"\"\n        :param init_with_true_state:\n        :param learn_odom:\n        :param use_proposer:\n        :param propose_ratio:\n        :param particle_std:\n        :param proposer_keep_ratio:\n        :param min_obs_likelihood:\n        \"\"\"\n\n        # store hyperparameters which are needed later\n        self.init_with_true_state = init_with_true_state\n        self.learn_odom = learn_odom\n        self.use_proposer = use_proposer and not init_with_true_state  # only use proposer if we do not initializet with true state\n        self.propose_ratio = propose_ratio if not self.init_with_true_state else 0.0\n\n        # define some more parameters and placeholders\n        self.state_dim = 3\n        self.placeholders = {'o': tf.placeholder('float32', [None, None, 24, 24, 3], 'observations'),\n                             'a': tf.placeholder('float32', [None, None, 3], 'actions'),\n                             's': tf.placeholder('float32', [None, None, 3], 'states'),\n                             'num_particles': tf.placeholder('float32'),\n                             'keep_prob': tf.placeholder_with_default(tf.constant(1.0), []),\n                             }\n        self.num_particles_float = self.placeholders['num_particles']\n        self.num_particles = tf.to_int32(self.num_particles_float)\n\n        # build learnable modules\n        self.build_modules(min_obs_likelihood, proposer_keep_ratio)\n\n\n    def build_modules(self, min_obs_likelihood, proposer_keep_ratio):\n        \"\"\"\n        :param min_obs_likelihood:\n        :param proposer_keep_ratio:\n        :return: None\n        \"\"\"\n\n        # MEASUREMENT MODEL\n\n        # conv net for encoding the image\n        self.encoder = snt.Sequential([\n            snt.nets.ConvNet2D([16, 32, 64], [[3, 3]], [2], [snt.SAME], activate_final=True, name='encoder/convnet'),\n            snt.BatchFlatten(),\n            lambda x: tf.nn.dropout(x,  self.placeholders['keep_prob']),\n            snt.Linear(128, name='encoder/linear'),\n            tf.nn.relu\n        ])\n\n        # observation likelihood estimator that maps states and image encodings to probabilities\n        self.obs_like_estimator = snt.Sequential([\n            snt.Linear(128, name='obs_like_estimator/linear'),\n            tf.nn.relu,\n            snt.Linear(128, name='obs_like_estimator/linear'),\n            tf.nn.relu,\n            snt.Linear(1, name='obs_like_estimator/linear'),\n            tf.nn.sigmoid,\n            lambda x: x * (1 - min_obs_likelihood) + min_obs_likelihood\n        ], name='obs_like_estimator')\n\n        # motion noise generator used for motion sampling\n        self.mo_noise_generator = snt.nets.MLP([32, 32, self.state_dim], activate_final=False, name='mo_noise_generator')\n\n        # odometry model (if we want to learn it)\n        if self.learn_odom:\n            self.mo_transition_model = snt.nets.MLP([128, 128, 128, self.state_dim], activate_final=False, name='mo_transition_model')\n\n        # particle proposer that maps encodings to particles (if we want to use it)\n        if self.use_proposer:\n            self.particle_proposer = snt.Sequential([\n                snt.Linear(128, name='particle_proposer/linear'),\n                tf.nn.relu,\n                lambda x: tf.nn.dropout(x,  proposer_keep_ratio),\n                snt.Linear(128, name='particle_proposer/linear'),\n                tf.nn.relu,\n                snt.Linear(128, name='particle_proposer/linear'),\n                tf.nn.relu,\n                snt.Linear(128, name='particle_proposer/linear'),\n                tf.nn.relu,\n                snt.Linear(4, name='particle_proposer/linear'),\n                tf.nn.tanh,\n            ])\n\n\n    def measurement_update(self, encoding, particles, means, stds):\n        \"\"\"\n        Compute the likelihood of the encoded observation for each particle.\n\n        :param encoding: encoding of the observation\n        :param particles:\n        :param means:\n        :param stds:\n        :return: observation likelihood\n        \"\"\"\n\n        # prepare input (normalize particles poses and repeat encoding per particle)\n        particle_input = self.transform_particles_as_input(particles, means, stds)\n        encoding_input = tf.tile(encoding[:, tf.newaxis, :], [1,  tf.shape(particles)[1], 1])\n        input = tf.concat([encoding_input, particle_input], axis=-1)\n\n        # estimate the likelihood of the encoded observation for each particle, remove last dimension\n        obs_likelihood = snt.BatchApply(self.obs_like_estimator)(input)[:, :, 0]\n\n        return obs_likelihood\n\n\n    def transform_particles_as_input(self, particles, means, stds):\n        return tf.concat([\n                   (particles[:, :, :2] - means['s'][:, :, :2]) / stds['s'][:, :, :2],  # normalized pos\n                   tf.cos(particles[:, :, 2:3]),  # cos\n                   tf.sin(particles[:, :, 2:3])], # sin\n                  axis=-1)\n\n\n    def propose_particles(self, encoding, num_particles, state_mins, state_maxs):\n        duplicated_encoding = tf.tile(encoding[:, tf.newaxis, :], [1, num_particles, 1])\n        proposed_particles = snt.BatchApply(self.particle_proposer)(duplicated_encoding)\n        proposed_particles = tf.concat([\n            proposed_particles[:,:,:1] * (state_maxs[0] - state_mins[0]) / 2.0 + (state_maxs[0] + state_mins[0]) / 2.0,\n            proposed_particles[:,:,1:2] * (state_maxs[1] - state_mins[1]) / 2.0 + (state_maxs[1] + state_mins[1]) / 2.0,\n            atan2(proposed_particles[:,:,2:3], proposed_particles[:,:,3:4])], axis=2)\n        return proposed_particles\n\n\n    def motion_update(self, actions, particles, means, stds, state_step_sizes, stop_sampling_gradient=False):\n        \"\"\"\n        Move particles according to odometry info in actions. Add learned noise.\n\n        :param actions:\n        :param particles:\n        :param means:\n        :param stds:\n        :param state_step_sizes:\n        :param stop_sampling_gradient:\n        :return: moved particles\n        \"\"\"\n\n        # 1. SAMPLE NOISY ACTIONS\n\n        # add dimension for particles\n        actions = actions[:, tf.newaxis, :]\n\n        # prepare input (normalize actions and repeat per particle)\n        action_input = tf.tile(actions / stds['a'], [1, tf.shape(particles)[1], 1])\n        random_input = tf.random_normal(tf.shape(action_input))\n        input = tf.concat([action_input, random_input], axis=-1)\n\n        # estimate action noise\n        delta = snt.BatchApply(self.mo_noise_generator)(input)\n        if stop_sampling_gradient:\n            delta = tf.stop_gradient(delta)\n\n        # zero-mean the action noise and add to actions\n        delta -= tf.reduce_mean(delta, axis=1, keep_dims=True)\n        noisy_actions = actions + delta\n\n        # 2. APPLY NOISY ACTIONS\n        if self.learn_odom:\n\n            # prepare input (normalize states and actions)\n            state_input = self.transform_particles_as_input(particles, means, stds)\n            action_input = noisy_actions / stds['a']\n            input = tf.concat([state_input, action_input], axis=-1)\n            # estimate state delta, scale it, and apply it\n            state_delta = snt.BatchApply(self.mo_transition_model)(input)\n            new_states = [particles[:, :, i:i+1] + state_delta[:, :, i:i+1] * state_step_sizes[i] for i in range(3)]\n            moved_particles = tf.concat(new_states[:2] + [wrap_angle(new_states[2])], axis=-1)\n\n        else:\n\n            # compute sin and cos of the particles\n            theta = particles[:, :, 2:3]\n            sin_theta = tf.sin(theta)\n            cos_theta = tf.cos(theta)\n            # move the particles using the noisy actions\n            new_x = particles[:, :, 0:1] + (noisy_actions[:, :, 0:1] * cos_theta + noisy_actions[:, :, 1:2] * sin_theta)\n            new_y = particles[:, :, 1:2] + (noisy_actions[:, :, 0:1] * sin_theta - noisy_actions[:, :, 1:2] * cos_theta)\n            new_theta = wrap_angle(particles[:, :, 2:3] + noisy_actions[:, :, 2:3])\n            moved_particles = tf.concat([new_x, new_y, new_theta], axis=-1)\n\n        return moved_particles\n\n\n    def compile_training_stages(self, sess, batch_iterators, particle_list, particle_probs_list, encodings, means, stds, state_step_sizes, state_mins, state_maxs, learning_rate, plot_task):\n\n        # TRAINING!\n        losses = dict()\n        train_stages = dict()\n\n        # TRAIN ODOMETRY\n\n        if self.learn_odom:\n\n            # apply model\n            motion_samples = self.motion_update(self.placeholders['a'][:,1],\n                                                self.placeholders['s'][:, :1],\n                                                means, stds, state_step_sizes,\n                                                stop_sampling_gradient=True)\n\n            # define loss and optimizer\n            sq_distance = compute_sq_distance(motion_samples, self.placeholders['s'][:, 1:2], state_step_sizes)\n            losses['motion_mse'] = tf.reduce_mean(sq_distance, name='loss')\n            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n            var_list = [v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if 'mo_transition_model' in v.name]\n\n            # put everything together\n            train_stages['train_odom'] = {\n                         'train_op': optimizer.minimize(losses['motion_mse'], var_list=var_list),\n                         'batch_iterator_names': {'train': 'train1', 'val': 'val1'},\n                         'monitor_losses': ['motion_mse'],\n                         'validation_loss': 'motion_mse',\n                         'plot': lambda e: self.plot_motion_model(sess, next(batch_iterators['val1']), motion_samples, plot_task) if e % 10 == 0 else None\n                         }\n\n        # TRAIN MOTION MODEL\n\n        # apply model\n        motion_samples = self.motion_update(self.placeholders['a'][:,1],\n                                            tf.tile(self.placeholders['s'][:, :1], [1, self.num_particles, 1]),\n                                            means, stds, state_step_sizes)\n\n        # define loss and optimizer\n        std = 0.01\n        sq_distance = compute_sq_distance(motion_samples, self.placeholders['s'][:, 1:2], state_step_sizes)\n        activations_sample = (1 / self.num_particles_float) / tf.sqrt(2 * np.pi * std ** 2) * tf.exp(\n            -sq_distance / (2.0 * std ** 2))\n        losses['motion_mle'] = tf.reduce_mean(-tf.log(1e-16 + tf.reduce_sum(activations_sample, axis=-1, name='loss')))\n        optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n        var_list = [v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if 'mo_noise_generator' in v.name]\n\n        # put everything together\n        train_stages['train_motion_sampling'] = {\n                     'train_op': optimizer.minimize(losses['motion_mle'], var_list=var_list),\n                     'batch_iterator_names': {'train': 'train1', 'val': 'val1'},\n                     'monitor_losses': ['motion_mle'],\n                     'validation_loss': 'motion_mle',\n                     'plot': lambda e: self.plot_motion_model(sess, next(batch_iterators['val1']), motion_samples, plot_task) if e % 10 == 0 else None\n                     }\n\n        # TRAIN MEASUREMENT MODEL\n\n        # apply model for all pairs of observations and states in that batch\n        test_particles = tf.tile(self.placeholders['s'][tf.newaxis, :, 0], [self.batch_size, 1, 1])\n        measurement_model_out = self.measurement_update(encodings[:, 0], test_particles, means, stds)\n\n        # define loss (correct -> 1, incorrect -> 0) and optimizer\n        correct_samples = tf.diag_part(measurement_model_out)\n        incorrect_samples = measurement_model_out - tf.diag(tf.diag_part(measurement_model_out))\n        losses['measurement_heuristic'] = tf.reduce_sum(-tf.log(correct_samples)) / tf.cast(self.batch_size, tf.float32) \\\n                                          + tf.reduce_sum(-tf.log(1.0 - incorrect_samples)) / tf.cast(self.batch_size * (self.batch_size - 1), tf.float32)\n        optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n        var_list = [v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if 'encoder' in v.name or 'obs_like_estimator' in v.name]\n\n        # put everything together\n        train_stages['train_measurement_model'] = {\n                     'train_op': optimizer.minimize(losses['measurement_heuristic'], var_list=var_list),\n                     'batch_iterator_names': {'train': 'train1', 'val': 'val1'},\n                     'monitor_losses': ['measurement_heuristic'],\n                     'validation_loss': 'measurement_heuristic',\n                     'plot': lambda e: self.plot_measurement_model(sess, batch_iterators['val1'], measurement_model_out) if e % 10 == 0 else None\n                     }\n\n        # TRAIN PARTICLE PROPOSER\n\n        if self.use_proposer:\n\n            # apply model (but only compute gradients until the encoding,\n            # otherwise we would unlearn it and the observation likelihood wouldn't work anymore)\n            proposed_particles = self.propose_particles(tf.stop_gradient(encodings[:, 0]), self.num_particles, state_mins, state_maxs)\n\n            # define loss and optimizer\n            std = 0.2\n            sq_distance = compute_sq_distance(proposed_particles, self.placeholders['s'][:, :1], state_step_sizes)\n            activations = (1 / self.num_particles_float) / tf.sqrt(2 * np.pi * std ** 2) * tf.exp(\n                -sq_distance / (2.0 * std ** 2))\n            losses['proposed_mle'] = tf.reduce_mean(-tf.log(1e-16 + tf.reduce_sum(activations, axis=-1)))\n            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n            var_list = [v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if 'particle_proposer' in v.name]\n\n            # put everything together\n            train_stages['train_particle_proposer'] = {\n                         'train_op': optimizer.minimize(losses['proposed_mle'], var_list=var_list),\n                         'batch_iterator_names': {'train': 'train1', 'val': 'val1'},\n                         'monitor_losses': ['proposed_mle'],\n                         'validation_loss': 'proposed_mle',\n                         'plot': lambda e: self.plot_particle_proposer(sess, next(batch_iterators['val1']), proposed_particles, plot_task) if e % 10 == 0 else None\n                         }\n\n        # END-TO-END TRAINING\n\n        # model was already applyed further up -> particle_list, particle_probs_list\n\n        # define losses and optimizer\n        # first loss (which is being optimized)\n        sq_distance = compute_sq_distance(particle_list, self.placeholders['s'][:, :, tf.newaxis, :], state_step_sizes)\n        activations = particle_probs_list[:, :] / tf.sqrt(2 * np.pi * std ** 2) * tf.exp(\n            -sq_distance / (2.0 * self.particle_std ** 2))\n        losses['mle'] = tf.reduce_mean(-tf.log(1e-16 + tf.reduce_sum(activations, axis=2, name='loss')))\n        # second loss (which we will monitor during execution)\n        pred = self.particles_to_state(particle_list, particle_probs_list)\n\n        sq_distance = compute_sq_distance(pred[:, -1, :], self.placeholders['s'][:, -1, :], state_step_sizes)\n        losses['mse_last'] = tf.reduce_mean(sq_distance)\n        # optimizer\n        optimizer = tf.train.AdamOptimizer(learning_rate)\n\n        # put everything together\n        train_stages['train_e2e'] = {\n                     'train_op': optimizer.minimize(losses['mle']),\n                     'batch_iterator_names': {'train': 'train', 'val': 'val'},\n                     'monitor_losses': ['mse_last', 'mle'],\n                     'validation_loss': 'mse_last',\n                     'plot': lambda e: self.plot_particle_filter(sess, next(batch_iterators['val_ex']), particle_list,\n                                                                 particle_probs_list, self.num_particles, state_step_sizes, plot_task) if e % 1 == 0 else None\n                     }\n\n        return losses, train_stages\n\n\n    def load(self, sess, model_path, model_file='best_validation', statistics_file='statistics.npz', connect_and_initialize=True, modules=('encoder', 'mo_noise_generator', 'mo_transition_model', 'obs_like_estimator', 'particle_proposer')):\n\n        if type(modules) not in [type(list()), type(tuple())]:\n            raise Exception('modules must be a list or tuple, not a ' + str(type(modules)))\n\n        # build the tensorflow graph\n        if connect_and_initialize:\n            # load training data statistics (which are needed to build the tf graph)\n            statistics = dict(np.load(os.path.join(model_path, statistics_file)))\n            for key in statistics.keys():\n                if statistics[key].shape == ():\n                    statistics[key] = statistics[key].item()  # convert 0d array of dictionary back to a normal dictionary\n\n            # connect all modules into the particle filter\n            self.connect_modules(**statistics)\n            init = tf.global_variables_initializer()\n            sess.run(init)\n        else:\n            statistics = None\n\n        # load variables\n        all_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)\n        vars_to_load = []\n        loaded_modules = set()\n        for v in all_vars:\n            for m in modules:\n                if m in v.name:\n                    vars_to_load.append(v)\n                    loaded_modules.add(m)\n\n        print('Loading these modules:', loaded_modules)\n\n        print('%s %s' % (model_path, model_file))\n        print('%r %r' % (model_path, model_file))\n\n        # restore variable values\n        saver = tf.train.Saver(vars_to_load)  # <- var list goes in here\n        saver.restore(sess, os.path.join(model_path, model_file))\n\n        print('Loaded the following variables:')\n        for v in vars_to_load:\n            print(v.name)\n\n        return statistics\n\n\n    def fit(self, sess, data, model_path, train_individually, train_e2e, split_ratio, seq_len, batch_size, epoch_length, num_epochs, patience, learning_rate, dropout_keep_ratio, num_particles, particle_std, plot_task=None, plot=False):\n\n        self.particle_std = particle_std\n\n        # preprocess data\n        data = split_data(data, ratio=split_ratio)\n        epoch_lengths = {'train': epoch_length, 'val': epoch_length*2}\n        batch_iterators = {'train': make_batch_iterator(data['train'], seq_len=seq_len, batch_size=batch_size),\n                           'val': make_repeating_batch_iterator(data['val'], epoch_lengths['val'], batch_size=batch_size, seq_len=seq_len),\n                           'train_ex': make_batch_iterator(data['train'], batch_size=batch_size, seq_len=seq_len),\n                           'val_ex': make_batch_iterator(data['val'], batch_size=batch_size, seq_len=seq_len),\n                           'train1': make_batch_iterator(data['train'], batch_size=batch_size, seq_len=2),\n                           'val1': make_repeating_batch_iterator(data['val'], epoch_lengths['val'], batch_size=batch_size, seq_len=2),\n                           }\n\n        # compute some statistics of the training data\n        means, stds, state_step_sizes, state_mins, state_maxs = compute_staticstics(data['train'])\n\n        # build the tensorflow graph by connecting all modules in the particles filter\n        particles, particle_probs, encodings, particle_list, particle_probs_list = self.connect_modules(means, stds, state_mins, state_maxs, state_step_sizes)\n\n        # define losses and train stages for different ways of training (e.g. training individual models and e2e training)\n        losses, train_stages = self.compile_training_stages(sess, batch_iterators, particle_list, particle_probs_list,\n                                                            encodings, means, stds, state_step_sizes, state_mins,\n                                                            state_maxs, learning_rate, plot_task)\n\n        # initialize variables\n        init = tf.global_variables_initializer()\n        sess.run(init)\n\n        # save statistics and prepare saving variables\n        if not os.path.exists(model_path):\n            os.makedirs(model_path)\n        np.savez(os.path.join(model_path, 'statistics'), means=means, stds=stds, state_step_sizes=state_step_sizes,\n                 state_mins=state_mins, state_maxs=state_maxs)\n        saver = tf.train.Saver()\n        save_path = os.path.join(model_path, 'best_validation')\n\n        # define the training curriculum\n        curriculum = []\n        if train_individually:\n            if self.learn_odom:\n                curriculum += ['train_odom']\n            curriculum += ['train_motion_sampling']\n            curriculum += ['train_measurement_model']\n            if self.use_proposer:\n                curriculum += ['train_particle_proposer']\n        if train_e2e:\n            curriculum += ['train_e2e']\n\n        # split data for early stopping\n        data_keys = ['train']\n        if split_ratio < 1.0:\n            data_keys.append('val')\n\n        # define log dict\n        log = {c: {dk: {lk: {'mean': [], 'se': []} for lk in train_stages[c]['monitor_losses']} for dk in data_keys} for c in curriculum}\n\n        # go through curriculum\n        for c in curriculum:\n\n            stage = train_stages[c]\n            best_val_loss = np.inf\n            best_epoch = 0\n            epoch = 0\n\n            while epoch < num_epochs and epoch - best_epoch < patience:\n                # training\n                for dk in data_keys:\n                    # don't train in the first epoch, just evaluate the initial parameters\n                    if dk == 'train' and epoch == 0:\n                        continue\n                    # set up loss lists which will be filled during the epoch\n                    loss_lists = {lk: [] for lk in stage['monitor_losses']}\n                    for e in range(epoch_lengths[dk]):\n                        # t0 = time.time()\n                        # pick a batch from the right iterator\n                        batch = next(batch_iterators[stage['batch_iterator_names'][dk]])\n\n                        # define the inputs and train/run the model\n                        input_dict = {**{self.placeholders[key]: batch[key] for key in 'osa'},\n                                      **{self.placeholders['num_particles']: num_particles},\n                                      }\n                        if dk == 'train':\n                            input_dict[self.placeholders['keep_prob']] = dropout_keep_ratio\n                        monitor_losses = {l: losses[l] for l in stage['monitor_losses']}\n                        if dk == 'train':\n                            s_losses, _ = sess.run([monitor_losses, stage['train_op']], input_dict)\n                        else:\n                            s_losses = sess.run(monitor_losses, input_dict)\n\n                        for lk in stage['monitor_losses']:\n                            loss_lists[lk].append(s_losses[lk])\n\n                    # after each epoch, compute and log statistics\n                    for lk in stage['monitor_losses']:\n                        log[c][dk][lk]['mean'].append(np.mean(loss_lists[lk]))\n                        log[c][dk][lk]['se'].append(np.std(loss_lists[lk], ddof=1) / np.sqrt(len(loss_lists[lk])))\n\n                # check whether the current model is better than all previous models\n                if 'val' in data_keys:\n                    current_val_loss = log[c]['val'][stage['validation_loss']]['mean'][-1]\n                    if current_val_loss < best_val_loss:\n                        best_val_loss = current_val_loss\n                        best_epoch = epoch\n                        # save current model\n                        saver.save(sess, save_path)\n                        txt = 'epoch {:>3} >> '.format(epoch)\n                    else:\n                        txt = 'epoch {:>3} == '.format(epoch)\n                else:\n                    best_epoch = epoch\n                    saver.save(sess, save_path)\n                    txt = 'epoch {:>3} >> '.format(epoch)\n\n                # after going through all data sets, do a print out of the current result\n                for lk in stage['monitor_losses']:\n                    txt += '{}: '.format(lk)\n                    for dk in data_keys:\n                        if len(log[c][dk][lk]['mean']) > 0:\n                            txt += '{:.2f}+-{:.2f}/'.format(log[c][dk][lk]['mean'][-1], log[c][dk][lk]['se'][-1])\n                    txt = txt[:-1] + ' -- '\n                print(txt)\n\n                # t1 = time.time()\n                # time_deltas.append(t1 - t0)\n\n                if plot:\n                    stage['plot'](epoch)\n\n                epoch += 1\n\n            # after running out of patience, restore the model with lowest validation loss\n            saver.restore(sess, save_path)\n\n        return log\n\n\n    def predict(self, sess, batch, num_particles, return_particles=False, **kwargs):\n        # define input dict, use the first state only if we do tracking\n        input_dict = {self.placeholders['o']: batch['o'],\n                      self.placeholders['a']: batch['a'],\n                      self.placeholders['num_particles']: num_particles}\n        if self.init_with_true_state:\n            input_dict[self.placeholders['s']] = batch['s'][:, :1]\n\n        if return_particles:\n            return sess.run([self.pred_states, self.particle_list, self.particle_probs_list], input_dict)\n        else:\n            return sess.run(self.pred_states, input_dict)\n\n\n    def connect_modules(self, means, stds, state_mins, state_maxs, state_step_sizes):\n\n        # get shapes\n        self.batch_size = tf.shape(self.placeholders['o'])[0]\n        self.seq_len = tf.shape(self.placeholders['o'])[1]\n        # we use the static shape here because we need it to build the graph\n        self.action_dim = self.placeholders['a'].get_shape()[-1].value\n\n        encodings = snt.BatchApply(self.encoder)((self.placeholders['o'] - means['o']) / stds['o'])\n        self.encodings = encodings\n\n        # initialize particles\n        if self.init_with_true_state:\n            # tracking with known initial state\n            initial_particles = tf.tile(self.placeholders['s'][:, 0, tf.newaxis, :], [1, self.num_particles, 1])\n        else:\n            # global localization\n            if self.use_proposer:\n                # propose particles from observations\n                initial_particles = self.propose_particles(encodings[:, 0], self.num_particles, state_mins, state_maxs)\n            else:\n                # sample particles randomly\n                initial_particles = tf.concat(\n                    [tf.random_uniform([self.batch_size, self.num_particles, 1], state_mins[d], state_maxs[d]) for d in\n                     range(self.state_dim)], axis=-1, name='particles')\n\n        initial_particle_probs = tf.ones([self.batch_size, self.num_particles],\n                                         name='particle_probs') / self.num_particles_float\n\n        # assumes that samples has the correct size\n        def permute_batch(x, samples):\n            # get shapes\n            batch_size = tf.shape(x)[0]\n            num_particles = tf.shape(x)[1]\n            sample_size = tf.shape(samples)[1]\n            # compute 1D indices into the 2D array\n            idx = samples + num_particles * tf.tile(\n                tf.reshape(tf.range(batch_size), [batch_size, 1]),\n                [1, sample_size])\n            # index using the 1D indices and reshape again\n            result = tf.gather(tf.reshape(x, [batch_size * num_particles, -1]), idx)\n            result = tf.reshape(result, tf.shape(x[:,:sample_size]))\n            return result\n\n\n        def loop(particles, particle_probs, particle_list, particle_probs_list, additional_probs_list, i):\n\n            num_proposed_float = tf.round((self.propose_ratio ** tf.cast(i, tf.float32)) * self.num_particles_float)\n            num_proposed = tf.cast(num_proposed_float, tf.int32)\n            num_resampled_float = self.num_particles_float - num_proposed_float\n            num_resampled = tf.cast(num_resampled_float, tf.int32)\n\n            if self.propose_ratio < 1.0:\n\n                # resampling\n                basic_markers = tf.linspace(0.0, (num_resampled_float - 1.0) / num_resampled_float, num_resampled)\n                random_offset = tf.random_uniform([self.batch_size], 0.0, 1.0 / num_resampled_float)\n                markers = random_offset[:, None] + basic_markers[None, :]  # shape: batch_size x num_resampled\n                cum_probs = tf.cumsum(particle_probs, axis=1)\n                marker_matching = markers[:, :, None] < cum_probs[:, None, :]  # shape: batch_size x num_resampled x num_particles\n                samples = tf.cast(tf.argmax(tf.cast(marker_matching, 'int32'), dimension=2), 'int32')\n                standard_particles = permute_batch(particles, samples)\n                standard_particle_probs = tf.ones([self.batch_size, num_resampled])\n                standard_particles = tf.stop_gradient(standard_particles)\n                standard_particle_probs = tf.stop_gradient(standard_particle_probs)\n\n                # motion update\n                standard_particles = self.motion_update(self.placeholders['a'][:, i], standard_particles, means, stds, state_step_sizes)\n\n                # measurement update\n                standard_particle_probs *= self.measurement_update(encodings[:, i], standard_particles, means, stds)\n\n            if self.propose_ratio > 0.0:\n\n                # proposed particles\n                proposed_particles = self.propose_particles(encodings[:, i], num_proposed, state_mins, state_maxs)\n                proposed_particle_probs = tf.ones([self.batch_size, num_proposed])\n\n\n            # NORMALIZE AND COMBINE PARTICLES\n            if self.propose_ratio == 1.0:\n                particles = proposed_particles\n                particle_probs = proposed_particle_probs\n\n            elif self.propose_ratio == 0.0:\n                particles = standard_particles\n                particle_probs = standard_particle_probs\n\n            else:\n                standard_particle_probs *= (num_resampled_float / self.num_particles_float) / tf.reduce_sum(standard_particle_probs, axis=1, keep_dims=True)\n                proposed_particle_probs *= (num_proposed_float / self.num_particles_float) / tf.reduce_sum(proposed_particle_probs, axis=1, keep_dims=True)\n                particles = tf.concat([standard_particles, proposed_particles], axis=1)\n                particle_probs = tf.concat([standard_particle_probs, proposed_particle_probs], axis=1)\n\n            # NORMALIZE PROBABILITIES\n            particle_probs /= tf.reduce_sum(particle_probs, axis=1, keep_dims=True)\n\n            particle_list = tf.concat([particle_list, particles[:, tf.newaxis]], axis=1)\n            particle_probs_list = tf.concat([particle_probs_list, particle_probs[:, tf.newaxis]], axis=1)\n\n            return particles, particle_probs, particle_list, particle_probs_list, additional_probs_list, i + 1\n\n        # reshapes and sets the first shape sizes to None (which is necessary to keep the shape consistent in while loop)\n        particle_list = tf.reshape(initial_particles,\n                                   shape=[self.batch_size, -1, self.num_particles, self.state_dim])\n        particle_probs_list = tf.reshape(initial_particle_probs, shape=[self.batch_size, -1, self.num_particles])\n        additional_probs_list = tf.reshape(tf.ones([self.batch_size, self.num_particles, 4]), shape=[self.batch_size, -1, self.num_particles, 4])\n\n        # run the filtering process\n        particles, particle_probs, particle_list, particle_probs_list, additional_probs_list, i = tf.while_loop(\n            lambda *x: x[-1] < self.seq_len, loop,\n            [initial_particles, initial_particle_probs, particle_list, particle_probs_list, additional_probs_list,\n             tf.constant(1, dtype='int32')], name='loop')\n\n        # compute mean of particles\n        self.pred_states = self.particles_to_state(particle_list, particle_probs_list)\n        self.particle_list = particle_list\n        self.particle_probs_list = particle_probs_list\n\n        return particles, particle_probs, encodings, particle_list, particle_probs_list\n\n    def particles_to_state(self, particle_list, particle_probs_list):\n        mean_position = tf.reduce_sum(particle_probs_list[:, :, :, tf.newaxis] * particle_list[:, :, :, :2], axis=2)\n        mean_orientation = atan2(\n            tf.reduce_sum(particle_probs_list[:, :, :, tf.newaxis] * tf.cos(particle_list[:, :, :, 2:]), axis=2),\n            tf.reduce_sum(particle_probs_list[:, :, :, tf.newaxis] * tf.sin(particle_list[:, :, :, 2:]), axis=2))\n        return tf.concat([mean_position, mean_orientation], axis=2)\n\n\n    def plot_motion_model(self, sess, batch, motion_samples, task):\n\n        # define the inputs and train/run the model\n        input_dict = {**{self.placeholders[key]: batch[key] for key in 'osa'},\n                      **{self.placeholders['num_particles']: 100},\n                      }\n\n        s_motion_samples = sess.run(motion_samples, input_dict)\n\n        plt.figure('Motion Model')\n        plt.gca().clear()\n        plot_maze(task)\n        for i in range(min(len(s_motion_samples), 10)):\n            plt.quiver(s_motion_samples[i, :, 0], s_motion_samples[i, :, 1], np.cos(s_motion_samples[i, :, 2]), np.sin(s_motion_samples[i, :, 2]), color='blue', width=0.001, scale=100)\n            plt.quiver(batch['s'][i, 0, 0], batch['s'][i, 0, 1], np.cos(batch['s'][i, 0, 2]), np.sin(batch['s'][i, 0, 2]), color='black', scale=50, width=0.003)\n            plt.quiver(batch['s'][i, 1, 0], batch['s'][i, 1, 1], np.cos(batch['s'][i, 1, 2]), np.sin(batch['s'][i, 1, 2]), color='red', scale=50, width=0.003)\n\n        plt.gca().set_aspect('equal')\n        plt.pause(0.01)\n\n\n    def plot_measurement_model(self, sess, batch_iterator, measurement_model_out):\n\n        batch = next(batch_iterator)\n\n        # define the inputs and train/run the model\n        input_dict = {**{self.placeholders[key]: batch[key] for key in 'osa'},\n                      **{self.placeholders['num_particles']: 100},\n                      }\n\n        s_measurement_model_out = sess.run(measurement_model_out, input_dict)\n\n        plt.figure('Measurement Model Output')\n        plt.gca().clear()\n        plt.imshow(s_measurement_model_out, interpolation=\"nearest\", cmap=\"coolwarm\")\n        plt.pause(0.01)\n\n\n\n    def plot_particle_proposer(self, sess, batch, proposed_particles, task):\n\n        # define the inputs and train/run the model\n        input_dict = {**{self.placeholders[key]: batch[key] for key in 'osa'},\n                      **{self.placeholders['num_particles']: 100},\n                      }\n\n        s_samples = sess.run(proposed_particles, input_dict)\n\n        plt.figure('Particle Proposer')\n        plt.gca().clear()\n        plot_maze(task)\n\n        for i in range(min(len(s_samples), 10)):\n            color = np.random.uniform(0.0, 1.0, 3)\n            plt.quiver(s_samples[i, :, 0], s_samples[i, :, 1], np.cos(s_samples[i, :, 2]), np.sin(s_samples[i, :, 2]), color=color, width=0.001, scale=100)\n            plt.quiver(batch['s'][i, 0, 0], batch['s'][i, 0, 1], np.cos(batch['s'][i, 0, 2]), np.sin(batch['s'][i, 0, 2]), color=color, scale=50, width=0.003)\n\n        plt.pause(0.01)\n\n\n    def plot_particle_filter(self, sess, batch, particle_list,\n                        particle_probs_list, num_particles, state_step_sizes, task):\n\n        num_particles = 1000\n        head_scale = 1.5\n        quiv_kwargs = {'scale_units': 'xy', 'scale': 1. / 40., 'width': 0.003, 'headlength': 5 * head_scale,\n                       'headwidth': 3 * head_scale, 'headaxislength': 4.5 * head_scale}\n        marker_kwargs = {'markersize': 4.5, 'markerfacecolor': 'None', 'markeredgewidth': 0.5}\n\n        color_list = plt.cm.tab10(np.linspace(0, 1, 10))\n        colors = {'lstm': color_list[0], 'pf_e2e': color_list[1], 'pf_ind_e2e': color_list[2], 'pf_ind': color_list[3],\n                  'ff': color_list[4], 'odom': color_list[4]}\n\n        pred, s_particle_list, s_particle_probs_list = self.predict(sess, batch, num_particles,\n                                                                      return_particles=True)\n\n        num_steps = 20  # s_particle_list.shape[1]\n\n        for s in range(1):\n\n            plt.figure(\"example {}\".format(s), figsize=[12, 5.15])\n            plt.gca().clear()\n\n            for i in range(num_steps):\n                ax = plt.subplot(4, 5, i + 1, frameon=False)\n                plt.gca().clear()\n\n                plot_maze(task, margin=5, linewidth=0.5)\n\n                if i < num_steps - 1:\n                    ax.quiver(s_particle_list[s, i, :, 0], s_particle_list[s, i, :, 1],\n                              np.cos(s_particle_list[s, i, :, 2]), np.sin(s_particle_list[s, i, :, 2]),\n                              s_particle_probs_list[s, i, :], cmap='viridis_r', clim=[.0, 2.0 / num_particles],\n                              alpha=1.0,\n                              **quiv_kwargs\n                              )\n\n                    current_state = batch['s'][s, i, :]\n                    plt.quiver(current_state[0], current_state[1], np.cos(current_state[2]),\n                               np.sin(current_state[2]), color=\"red\", **quiv_kwargs)\n\n                    plt.plot(current_state[0], current_state[1], 'or', **marker_kwargs)\n                else:\n\n                    ax.plot(batch['s'][s, :num_steps, 0], batch['s'][s, :num_steps, 1], '-', linewidth=0.6, color='red')\n                    ax.plot(pred[s, :num_steps, 0], pred[s, :num_steps, 1], '-', linewidth=0.6,\n                            color=colors['pf_ind_e2e'])\n\n                    ax.plot(batch['s'][s, :1, 0], batch['s'][s, :1, 1], '.', linewidth=0.6, color='red', markersize=3)\n                    ax.plot(pred[s, :1, 0], pred[s, :1, 1], '.', linewidth=0.6, markersize=3,\n                            color=colors['pf_ind_e2e'])\n\n                plt.subplots_adjust(left=0.0, bottom=0.0, right=1.0, top=1.0, wspace=0.001, hspace=0.1)\n                plt.gca().set_aspect('equal')\n                plt.xticks([])\n                plt.yticks([])\n\n        show_pause(pause=0.01)\n"
  },
  {
    "path": "methods/dpf_kitti.py",
    "content": "import os\nimport numpy as np\nimport sonnet as snt\nimport tensorflow as tf\nimport matplotlib.pyplot as plt\n\nfrom utils.data_utils_kitti import wrap_angle, compute_statistics, split_data, make_batch_iterator, make_repeating_batch_iterator, rotation_matrix, load_data_for_stats\nfrom utils.method_utils import atan2, compute_sq_distance\nfrom utils.plotting_utils import plot_maze, show_pause\nfrom datetime import datetime\n\nif tf.__version__ == '1.1.0-rc1' or tf.__version__ == '1.2.0':\n    from tensorflow.python.framework import ops\n    @ops.RegisterGradient(\"FloorMod\")\n    def _mod_grad(op, grad):\n        x, y = op.inputs\n        gz = grad\n        x_grad = gz\n        y_grad = None  # tf.reduce_mean(-(x // y) * gz, axis=[0], keep_dims=True)[0]\n        return x_grad, y_grad\n\n\nclass DPF():\n\n    def __init__(self, init_with_true_state, learn_odom, use_proposer, propose_ratio, proposer_keep_ratio, min_obs_likelihood, learn_gaussian_mle):\n        \"\"\"\n        :param init_with_true_state:\n        :param learn_odom:\n        :param use_proposer:\n        :param propose_ratio:\n        :param particle_std:\n        :param proposer_keep_ratio:\n        :param min_obs_likelihood:\n        \"\"\"\n\n        # store hyperparameters which are needed later\n        self.init_with_true_state = init_with_true_state\n        self.learn_odom = learn_odom\n        self.use_proposer = use_proposer and not init_with_true_state  # only use proposer if we do not initializet with true state\n        self.propose_ratio = propose_ratio if not self.init_with_true_state else 0.0\n\n        # define some more parameters and placeholders\n        self.state_dim = 5\n        self.action_dim = 3\n        self.observation_dim = 6\n        self.placeholders = {'o': tf.placeholder('float32', [None, None, 50, 150, self.observation_dim], 'observations'),\n                             'a': tf.placeholder('float32', [None, None, 3], 'actions'),\n                             's': tf.placeholder('float32', [None, None, 5], 'states'),\n                             'num_particles': tf.placeholder('float32'),\n                             'keep_prob': tf.placeholder_with_default(tf.constant(1.0), []),\n                             'is_training': tf.placeholder_with_default(tf.constant(False), [])\n                             }\n        self.num_particles_float = self.placeholders['num_particles']\n        self.num_particles = tf.to_int32(self.num_particles_float)\n\n        # build learnable modules\n        self.build_modules(min_obs_likelihood, proposer_keep_ratio, learn_gaussian_mle)\n\n\n    def build_modules(self, min_obs_likelihood, proposer_keep_ratio, learn_gaussian_mle):\n        \"\"\"\n        :param min_obs_likelihood:\n        :param proposer_keep_ratio:\n        :return: None\n        \"\"\"\n\n        # MEASUREMENT MODEL\n\n        # conv net for encoding the image\n        self.encoder = snt.Sequential([\n            snt.nets.ConvNet2D([16, 16, 16, 16], [[7, 7], [5, 5], [5, 5], [5, 5]], [[1,1], [1, 2], [1, 2], [2, 2]], [snt.SAME], activate_final=True, name='encoder/convnet'),\n            snt.BatchFlatten(),\n            lambda x: tf.nn.dropout(x,  self.placeholders['keep_prob']),\n            snt.Linear(128, name='encoder/linear'),\n            tf.nn.relu\n        ])\n\n        # observation likelihood estimator that maps states and image encodings to probabilities\n        self.obs_like_estimator = snt.Sequential([\n            snt.Linear(128, name='obs_like_estimator/linear'),\n            tf.nn.relu,\n            snt.Linear(128, name='obs_like_estimator/linear'),\n            tf.nn.relu,\n            snt.Linear(1, name='obs_like_estimator/linear'),\n            tf.nn.sigmoid,\n            lambda x: x * (1 - min_obs_likelihood) + min_obs_likelihood\n        ], name='obs_like_estimator')\n\n        # motion noise generator used for motion sampling\n        if learn_gaussian_mle:\n            self.mo_noise_generator = snt.nets.MLP([32, 32, 4], activate_final=False, name='mo_noise_generator')\n        else:\n            self.mo_noise_generator = snt.nets.MLP([32, 32, 2], activate_final=False, name='mo_noise_generator')\n\n        # odometry model (if we want to learn it)\n        if self.learn_odom:\n            self.mo_transition_model = snt.nets.MLP([128, 128, 128, self.state_dim], activate_final=False, name='mo_transition_model')\n\n        # particle proposer that maps encodings to particles (if we want to use it)\n        if self.use_proposer:\n            self.particle_proposer = snt.Sequential([\n                snt.Linear(128, name='particle_proposer/linear'),\n                tf.nn.relu,\n                lambda x: tf.nn.dropout(x,  proposer_keep_ratio),\n                snt.Linear(128, name='particle_proposer/linear'),\n                tf.nn.relu,\n                snt.Linear(128, name='particle_proposer/linear'),\n                tf.nn.relu,\n                snt.Linear(128, name='particle_proposer/linear'),\n                tf.nn.relu,\n                snt.Linear(4, name='particle_proposer/linear'),\n                tf.nn.tanh,\n            ])\n\n        self.noise_scaler1 = snt.Module(lambda x: x * tf.exp(10 * tf.get_variable('motion_sampler/noise_scaler1', initializer=np.array(0.0, dtype='float32'))))\n        self.noise_scaler2 = snt.Module(lambda x: x * tf.exp(10 * tf.get_variable('motion_sampler/noise_scaler2', initializer=np.array(0.0, dtype='float32'))))\n\n\n    def custom_build(self, inputs):\n        \"\"\"A custom build method to wrap into a sonnet Module.\"\"\"\n        outputs = snt.Conv2D(output_channels=16, kernel_shape=[7, 7], stride=[1, 1])(inputs)\n        outputs = tf.nn.relu(outputs)\n        outputs = snt.Conv2D(output_channels=16, kernel_shape=[5, 5], stride=[1, 2])(outputs)\n        outputs = tf.nn.relu(outputs)\n        outputs = snt.Conv2D(output_channels=16, kernel_shape=[5, 5], stride=[1, 2])(outputs)\n        outputs = tf.nn.relu(outputs)\n        outputs = snt.Conv2D(output_channels=16, kernel_shape=[5, 5], stride=[2, 2])(outputs)\n        outputs = tf.nn.relu(outputs)\n        outputs = tf.nn.dropout(outputs,  self.placeholders['keep_prob'])\n        outputs = snt.BatchFlatten()(outputs)\n        outputs = snt.Linear(128)(outputs)\n        outputs = tf.nn.relu(outputs)\n\n        return outputs\n\n    def measurement_update(self, encoding, particles, means, stds):\n        \"\"\"\n        Compute the likelihood of the encoded observation for each particle.\n\n        :param encoding: encoding of the observation\n        :param particles:\n        :param means:\n        :param stds:\n        :return: observation likelihood\n        \"\"\"\n\n        # prepare input (normalize particles poses and repeat encoding per particle)\n        particle_input = self.transform_particles_as_input(particles, means, stds)\n        encoding_input = tf.tile(encoding[:, tf.newaxis, :], [1,  tf.shape(particles)[1], 1])\n        input = tf.concat([encoding_input, particle_input], axis=-1)\n\n        # estimate the likelihood of the encoded observation for each particle, remove last dimension\n        obs_likelihood = snt.BatchApply(self.obs_like_estimator)(input)[:, :, 0]\n\n        return obs_likelihood\n\n\n    def transform_particles_as_input(self, particles, means, stds):\n        return ((particles - means['s']) / stds['s'])[..., 3:5]\n\n\n    def propose_particles(self, encoding, num_particles, state_mins, state_maxs):\n        duplicated_encoding = tf.tile(encoding[:, tf.newaxis, :], [1, num_particles, 1])\n        proposed_particles = snt.BatchApply(self.particle_proposer)(duplicated_encoding)\n        proposed_particles = tf.concat([\n            proposed_particles[:,:,:1] * (state_maxs[0] - state_mins[0]) / 2.0 + (state_maxs[0] + state_mins[0]) / 2.0,\n            proposed_particles[:,:,1:2] * (state_maxs[1] - state_mins[1]) / 2.0 + (state_maxs[1] + state_mins[1]) / 2.0,\n            atan2(proposed_particles[:,:,2:3], proposed_particles[:,:,3:4])], axis=2)\n        return proposed_particles\n\n\n    def motion_update(self, actions, particles, means, stds, state_step_sizes, learn_gaussian_mle, stop_sampling_gradient=False):\n        \"\"\"\n        Move particles according to odometry info in actions. Add learned noise.\n\n        :param actions:\n        :param particles:\n        :param means:\n        :param stds:\n        :param state_step_sizes:\n        :param stop_sampling_gradient:\n        :return: moved particles\n        \"\"\"\n\n        # 1. SAMPLE NOISY ACTIONS\n\n        # add dimension for particles\n        time_step = 0.103\n\n        if learn_gaussian_mle:\n            actions = tf.concat([particles[:, :, 3:4] - means['s'][:, :, 3:4], particles[:, :, 4:5] - means['s'][:, :, 4:5]], axis=-1)\n\n            # prepare input (normalize actions and repeat per particle)\n            action_input = actions / stds['s'][:, :, 3:5]\n            input = action_input\n\n            # estimate action noise\n            delta = snt.BatchApply(self.mo_noise_generator)(input)\n            delta = tf.concat([delta[:, :, 0:2] * state_step_sizes[3], delta[:, :, 2:4] * state_step_sizes[4]], axis=-1)\n            if stop_sampling_gradient:\n                delta = tf.stop_gradient(delta)\n\n            action_vel_f = tf.random_normal(tf.shape(particles[:, :, 3:4]), mean = delta[:, :, 0:1], stddev = delta[:, :, 1:2])\n            action_vel_rot = tf.random_normal(tf.shape(particles[:, :, 4:5]), mean = delta[:, :, 2:3], stddev = delta[:, :, 3:4])\n\n            heading = particles[:, :, 2:3]\n            sin_heading = tf.sin(heading)\n            cos_heading = tf.cos(heading)\n\n            new_x = particles[:, :, 0:1] + cos_heading * particles[:, :, 3:4] * time_step\n            new_y = particles[:, :, 1:2] + sin_heading * particles[:, :, 3:4] * time_step\n            new_theta = particles[:, :, 2:3] + particles[:, :, 4:5] * time_step\n            wrap_angle(new_theta)\n            new_v = particles[:, :, 3:4] + action_vel_f\n            new_theta_dot = particles[:, :, 4:5] + action_vel_rot\n\n            moved_particles = tf.concat([new_x, new_y, new_theta, new_v, new_theta_dot], axis=-1)\n\n            return moved_particles, delta\n\n        else:\n\n            heading = particles[:, :, 2:3]\n            sin_heading = tf.sin(heading)\n            cos_heading = tf.cos(heading)\n\n            random_input = tf.random_normal(tf.shape(particles[:, :, 3:5]))\n            noise = snt.BatchApply(self.mo_noise_generator)(random_input)\n            noise = noise - tf.reduce_mean(noise, axis=1, keep_dims=True)\n\n            new_z = particles[:, :, 0:1] + cos_heading * particles[:, :, 3:4] * time_step\n            new_x = particles[:, :, 1:2] + sin_heading * particles[:, :, 3:4] * time_step\n            new_theta = wrap_angle(particles[:, :, 2:3] + particles[:, :, 4:5] * time_step)\n\n            new_v = particles[:, :, 3:4] + noise[:, :, :1] * state_step_sizes[3]\n            new_theta_dot = particles[:, :, 4:5] + noise[:, :, 1:] * state_step_sizes[4]\n\n            moved_particles = tf.concat([new_z, new_x, new_theta, new_v, new_theta_dot], axis=-1)\n\n            return moved_particles\n\n\n    def compile_training_stages(self, sess, batch_iterators, particle_list, particle_probs_list, encodings, means, stds, state_step_sizes, state_mins, state_maxs, learn_gaussian_mle, learning_rate, plot_task):\n\n        # TRAINING!\n        losses = dict()\n        train_stages = dict()\n        std = 0.25\n\n        # TRAIN ODOMETRY\n\n        if self.learn_odom:\n\n            # apply model\n            motion_samples = self.motion_update(self.placeholders['a'][:,0],\n                                                self.placeholders['s'][:, :1],\n                                                means, stds, state_step_sizes,\n                                                stop_sampling_gradient=True)\n\n            # define loss and optimizer\n            sq_distance = compute_sq_distance(motion_samples, self.placeholders['s'][:, 1:2], state_step_sizes)\n            losses['motion_mse'] = tf.reduce_mean(sq_distance, name='loss')\n            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n\n            # put everything together\n            train_stages['train_odom'] = {\n                         'train_op': optimizer.minimize(losses['motion_mse']),\n                         'batch_iterator_names': {'train': 'train1', 'val': 'val1'},\n                         'monitor_losses': ['motion_mse'],\n                         'validation_loss': 'motion_mse',\n                         'plot': lambda e: self.plot_motion_model(sess, next(batch_iterators['val2']), motion_samples, plot_task, state_step_sizes) if e % 1 == 0 else None\n                         }\n\n        # TRAIN MOTION MODEL\n\n        if learn_gaussian_mle:\n            motion_samples, motion_params = self.motion_update(self.placeholders['a'][:,1],\n                                                tf.tile(self.placeholders['s'][:, :1], [1, 1, 1]),\n                                                means, stds, state_step_sizes, learn_gaussian_mle)\n\n            # define loss and optimizer\n            diff_in_states = self.placeholders['s'][:, 1:2] - self.placeholders['s'][:, :1]\n            activations_vel_f = (1 / 32) / tf.sqrt(2 * np.pi * motion_params[:, :, 1] ** 2) * tf.exp(\n                -(diff_in_states[:, :, 3] - motion_params[:, :, 0]) ** 2 / (2.0 * motion_params[:, :, 1] ** 2))\n            activations_vel_rot = (1 / 32) / tf.sqrt(2 * np.pi * motion_params[:, :, 3] ** 2) * tf.exp(\n                -(diff_in_states[:, :, 4] - motion_params[:, :, 2]) ** 2 / (2.0 * motion_params[:, :, 3] ** 2))\n            losses['motion_mle'] = tf.reduce_mean(-tf.log(1e-16 + (tf.reduce_sum(activations_vel_f, axis=-1, name='loss1') * tf.reduce_sum(activations_vel_rot, axis=-1, name='loss2'))))\n            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n\n            # put everything together\n            train_stages['train_motion_sampling'] = {\n                         'train_op': optimizer.minimize(losses['motion_mle']),\n                         'batch_iterator_names': {'train': 'train2', 'val': 'val2'},\n                         'monitor_losses': ['motion_mle'],\n                         'validation_loss': 'motion_mle',\n                         'plot': lambda e: self.plot_motion_model(sess, next(batch_iterators['val2']), motion_samples, plot_task, state_step_sizes) if e % 1 == 0 else None\n                         }\n\n        else:\n            motion_samples = self.motion_update(self.placeholders['a'][:,1],\n                                    tf.tile(self.placeholders['s'][:, :1], [1, self.num_particles, 1]),\n                                    means, stds, state_step_sizes, learn_gaussian_mle)\n\n            # define loss and optimizer\n            sq_distance = compute_sq_distance(motion_samples, self.placeholders['s'][:, 1:2], state_step_sizes)\n            activations_sample = (1 / self.num_particles_float) / tf.sqrt(2 * np.pi * std ** 2) * tf.exp(\n                -sq_distance / (2.0 * std ** 2))\n            losses['motion_mle'] = tf.reduce_mean(-tf.log(1e-16 + tf.reduce_sum(activations_sample, axis=-1, name='loss')))\n            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n\n            # put everything together\n            train_stages['train_motion_sampling'] = {\n                         'train_op': optimizer.minimize(losses['motion_mle']),\n                         'batch_iterator_names': {'train': 'train2', 'val': 'val2'},\n                         'monitor_losses': ['motion_mle'],\n                         'validation_loss': 'motion_mle',\n                         'plot': lambda e: self.plot_motion_model(sess, next(batch_iterators['val2']), motion_samples, plot_task, state_step_sizes) if e % 1 == 0 else None\n                         }\n\n        # TRAIN MEASUREMENT MODEL\n\n        # apply model for all pairs of observations and states in that batch\n        test_particles = tf.tile(self.placeholders['s'][tf.newaxis, :, 0], [self.batch_size, 1, 1])\n        measurement_model_out = self.measurement_update(encodings[:, 0], test_particles, means, stds)\n\n        # define loss (correct -> 1, incorrect -> 0) and optimizer\n        correct_samples = tf.diag_part(measurement_model_out)\n        incorrect_samples = measurement_model_out - tf.diag(tf.diag_part(measurement_model_out))\n        losses['measurement_heuristic'] = tf.reduce_sum(-tf.log(correct_samples)) / tf.cast(self.batch_size, tf.float32) \\\n                                          + tf.reduce_sum(-tf.log(1.0 - incorrect_samples)) / tf.cast(self.batch_size * (self.batch_size - 1), tf.float32)\n        optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n\n        # put everything together\n        train_stages['train_measurement_model'] = {\n                     'train_op': optimizer.minimize(losses['measurement_heuristic']),\n                     'batch_iterator_names': {'train': 'train1', 'val': 'val1'},\n                     'monitor_losses': ['measurement_heuristic'],\n                     'validation_loss': 'measurement_heuristic',\n                     'plot': lambda e: self.plot_measurement_model(sess, batch_iterators['val1'], measurement_model_out) if e % 1 == 0 else None\n                     }\n\n        # TRAIN PARTICLE PROPOSER\n\n        if self.use_proposer:\n\n            # apply model (but only compute gradients until the encoding,\n            # otherwise we would unlearn it and the observation likelihood wouldn't work anymore)\n            proposed_particles = self.propose_particles(tf.stop_gradient(encodings[:, 0]), self.num_particles, state_mins, state_maxs)\n\n            # define loss and optimizer\n            std = 0.2\n            sq_distance = compute_sq_distance(proposed_particles, self.placeholders['s'][:, :1], state_step_sizes)\n            activations = (1 / self.num_particles_float) / tf.sqrt(2 * np.pi * std ** 2) * tf.exp(\n                -sq_distance / (2.0 * std ** 2))\n            losses['proposed_mle'] = tf.reduce_mean(-tf.log(1e-16 + tf.reduce_sum(activations, axis=-1)))\n            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n\n            # put everything together\n            train_stages['train_particle_proposer'] = {\n                         'train_op': optimizer.minimize(losses['proposed_mle']),\n                         'batch_iterator_names': {'train': 'train1', 'val': 'val1'},\n                         'monitor_losses': ['proposed_mle'],\n                         'validation_loss': 'proposed_mle',\n                         'plot': lambda e: self.plot_particle_proposer(sess, next(batch_iterators['val1']), proposed_particles, plot_task) if e % 10 == 0 else None\n                         }\n\n\n        # END-TO-END TRAINING\n\n        # model was already applied further up -> particle_list, particle_probs_list\n\n        # define losses and optimizer\n        # first loss (which is being optimized)\n        sq_distance = compute_sq_distance(particle_list[:, :, :, 3:5], self.placeholders['s'][:, :, tf.newaxis, 3:5], state_step_sizes[3:5])\n        activations = particle_probs_list[:, :] / tf.sqrt(2 * np.pi * self.particle_std ** 2) * tf.exp(\n            -sq_distance / (2.0 * self.particle_std ** 2))\n        losses['mle'] = tf.reduce_mean(-tf.log(1e-16 + tf.reduce_sum(activations, axis=2, name='loss')))\n\n        # second loss (which we will monitor during execution)\n        pred = self.particles_to_state(particle_list, particle_probs_list)\n\n        sq_error = compute_sq_distance(pred[:, -1, 0:2], self.placeholders['s'][:, -1, 0:2], [1., 1.])\n        sq_dist = compute_sq_distance(self.placeholders['s'][:, 0, 0:2], self.placeholders['s'][:, -1, 0:2], [1., 1.])\n        losses['m/m'] = tf.reduce_mean(sq_error**0.5/sq_dist**0.5)\n\n        sq_error = compute_sq_distance(pred[:, -1, 2:3], self.placeholders['s'][:, -1, 2:3], [np.pi/180.0])\n        losses['deg/m'] = tf.reduce_mean(sq_error ** 0.5 / sq_dist ** 0.5)\n\n        # optimizer\n        optimizer = tf.train.AdamOptimizer(learning_rate)\n\n        # put everything together\n        train_stages['train_e2e'] = {\n                     'train_op': optimizer.minimize(losses['mle']),\n                     'batch_iterator_names': {'train': 'train', 'val': 'val'},\n                     'monitor_losses': ['m/m', 'deg/m', 'mle'],\n                     'validation_loss': 'deg/m',\n                     'plot': lambda e: self.plot_particle_filter(sess, next(batch_iterators['val_ex']), particle_list,\n                                                                 particle_probs_list, state_step_sizes, plot_task) if e % 1 == 0 else None\n                     }\n\n        return losses, train_stages\n\n\n    def load(self, sess, model_path, model_file='best_validation', statistics_file='statistics.npz', connect_and_initialize=True, modules=('encoder', 'mo_noise_generator', 'mo_transition_model', 'obs_like_estimator', 'particle_proposer')):\n\n        if type(modules) not in [type(list()), type(tuple())]:\n            raise Exception('modules must be a list or tuple, not a ' + str(type(modules)))\n\n        # build the tensorflow graph\n        if connect_and_initialize:\n            # load training data statistics (which are needed to build the tf graph)\n            statistics = dict(np.load(os.path.join(model_path, statistics_file)))\n            for key in statistics.keys():\n                if statistics[key].shape == ():\n                    statistics[key] = statistics[key].item()  # convert 0d array of dictionary back to a normal dictionary\n\n            # connect all modules into the particle filter\n            self.connect_modules(**statistics)\n            init = tf.global_variables_initializer()\n            sess.run(init)\n\n        # load variables\n        all_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)\n        vars_to_load = []\n        loaded_modules = set()\n        for v in all_vars:\n            for m in modules:\n                if m in v.name:\n                    vars_to_load.append(v)\n                    loaded_modules.add(m)\n\n        print('Loading all modules')\n\n        saver = tf.train.Saver()\n        saver.restore(sess, os.path.join(model_path, model_file))\n\n    # def fit(self, sess, data, model_path, train_individually, train_e2e, split_ratio, seq_len, batch_size, epoch_length, num_epochs, patience, learning_rate, dropout_keep_ratio, num_particles, particle_std, plot_task=None, plot=False):\n    def fit(self, sess, data, model_path, train_individually, train_e2e, split_ratio, seq_len, batch_size, epoch_length, num_epochs, patience, learning_rate, dropout_keep_ratio, num_particles, particle_std, learn_gaussian_mle, plot_task=None, plot=False):\n        if plot:\n            plt.ion()\n\n        self.particle_std = particle_std\n\n        mean_loss_for_plot = np.zeros((1,))\n\n        means, stds, state_step_sizes, state_mins, state_maxs = compute_statistics(data)\n\n\n        data = split_data(data, ratio=split_ratio)\n\n        epoch_lengths = {'train': epoch_length, 'val': epoch_length*2}\n        batch_iterators = {'train': make_batch_iterator(data['train'], seq_len=seq_len, batch_size=batch_size),\n                           'val': make_repeating_batch_iterator(data['val'], epoch_lengths['val'], batch_size=batch_size, seq_len=seq_len),\n                           'train_ex': make_batch_iterator(data['train'], batch_size=batch_size, seq_len=seq_len),\n                           'val_ex': make_batch_iterator(data['val'], batch_size=batch_size, seq_len=seq_len),\n                           'train1': make_batch_iterator(data['train'], batch_size=batch_size, seq_len=1),\n                           'train2': make_batch_iterator(data['train'], batch_size=batch_size, seq_len=2),\n                            'val1': make_repeating_batch_iterator(data['val'], epoch_lengths['val'], batch_size=batch_size, seq_len=1),\n                            'val2': make_repeating_batch_iterator(data['val'], epoch_lengths['val'], batch_size=batch_size, seq_len=2),\n                        }\n\n        # build the tensorflow graph by connecting all modules in the particles filter\n        particles, particle_probs, encodings, particle_list, particle_probs_list = self.connect_modules(means, stds, state_mins, state_maxs, state_step_sizes, learn_gaussian_mle)\n\n        # define losses and train stages for different ways of training (e.g. training individual models and e2e training)\n        losses, train_stages = self.compile_training_stages(sess, batch_iterators, particle_list, particle_probs_list,\n                                                            encodings, means, stds, state_step_sizes, state_mins,\n                                                            state_maxs, learn_gaussian_mle, learning_rate, plot_task)\n\n        # initialize variables\n        init = tf.global_variables_initializer()\n        sess.run(init)\n\n        # save statistics and prepare saving variables\n        if not os.path.exists(model_path):\n            os.makedirs(model_path)\n        np.savez(os.path.join(model_path, 'statistics'), means=means, stds=stds, state_step_sizes=state_step_sizes,\n                 state_mins=state_mins, state_maxs=state_maxs)\n        saver = tf.train.Saver()\n        save_path = os.path.join(model_path, 'best_validation')\n\n        # define the training curriculum\n        curriculum = []\n        if train_individually:\n            if self.learn_odom:\n                curriculum += ['train_odom']\n            curriculum += ['train_measurement_model']\n            curriculum += ['train_motion_sampling']\n            if self.use_proposer:\n                curriculum += ['train_particle_proposer']\n        if train_e2e:\n            curriculum += ['train_e2e']\n\n        # split data for early stopping\n        data_keys = ['train']\n        if split_ratio < 1.0:\n            data_keys.append('val')\n\n        # define log dict\n        log = {c: {dk: {lk: {'mean': [], 'se': []} for lk in train_stages[c]['monitor_losses']} for dk in data_keys} for c in curriculum}\n\n        # go through curriculum\n        for c in curriculum:\n\n            stage = train_stages[c]\n            best_val_loss = np.inf\n            best_epoch = 0\n            epoch = 0\n\n            if c == 'train_e2e':\n                saver.save(sess, os.path.join(model_path, 'before_e2e/best_validation'))\n                np.savez(os.path.join(model_path, 'before_e2e/statistics'), means=means, stds=stds, state_step_sizes=state_step_sizes,\n                 state_mins=state_mins, state_maxs=state_maxs)\n            while epoch < num_epochs and epoch - best_epoch < patience:\n                # training\n                for dk in data_keys:\n                    # don't train in the first epoch, just evaluate the initial parameters\n                    if dk == 'train' and epoch == 0:\n                        continue\n                    # set up loss lists which will be filled during the epoch\n                    loss_lists = {lk: [] for lk in stage['monitor_losses']}\n                    for e in range(epoch_lengths[dk]):\n                        # t0 = time.time()\n                        # pick a batch from the right iterator\n                        batch = next(batch_iterators[stage['batch_iterator_names'][dk]])\n                        # define the inputs and train/run the model\n                        input_dict = {**{self.placeholders[key]: batch[key] for key in 'osa'},\n                                      **{self.placeholders['num_particles']: num_particles},\n                                      }\n                        if dk == 'train':\n                            input_dict[self.placeholders['keep_prob']] = dropout_keep_ratio\n                            input_dict[self.placeholders['is_training']] = True\n                        monitor_losses = {l: losses[l] for l in stage['monitor_losses']}\n                        if dk == 'train':\n                            s_losses, _ = sess.run([monitor_losses, stage['train_op']], input_dict)\n                        else:\n                            s_losses = sess.run(monitor_losses, input_dict)\n\n                        for lk in stage['monitor_losses']:\n                            loss_lists[lk].append(s_losses[lk])\n\n                    # after each epoch, compute and log statistics\n                    for lk in stage['monitor_losses']:\n                        log[c][dk][lk]['mean'].append(np.mean(loss_lists[lk]))\n                        log[c][dk][lk]['se'].append(np.std(loss_lists[lk], ddof=1) / np.sqrt(len(loss_lists[lk])))\n\n\n                # check whether the current model is better than all previous models\n                if 'val' in data_keys:\n                    current_val_loss = log[c]['val'][stage['validation_loss']]['mean'][-1]\n                    mean_loss_for_plot = np.append(mean_loss_for_plot,current_val_loss)\n                    if current_val_loss < best_val_loss:\n                        best_val_loss = current_val_loss\n                        best_epoch = epoch\n                        # save current model\n                        saver.save(sess, save_path)\n                        txt = 'epoch {:>3} >> '.format(epoch)\n                    else:\n                        txt = 'epoch {:>3} == '.format(epoch)\n                else:\n                    best_epoch = epoch\n                    saver.save(sess, save_path)\n                    txt = 'epoch {:>3} >> '.format(epoch)\n\n                # after going through all data sets, do a print out of the current result\n                for lk in stage['monitor_losses']:\n                    txt += '{}: '.format(lk)\n                    for dk in data_keys:\n                        if len(log[c][dk][lk]['mean']) > 0:\n                            txt += '{:.2f}+-{:.2f}/'.format(log[c][dk][lk]['mean'][-1], log[c][dk][lk]['se'][-1])\n\n                    txt = txt[:-1] + ' -- '\n                print(txt)\n\n                if plot:\n                    stage['plot'](epoch)\n\n                epoch += 1\n\n            # after running out of patience, restore the model with lowest validation loss\n            saver.restore(sess, save_path)\n\n        return log\n\n\n    def predict(self, sess, batch, return_particles=False, **kwargs):\n        # define input dict, use the first state only if we do tracking\n        input_dict = {self.placeholders['o']: batch['o'],\n                      self.placeholders['a']: batch['a'],\n                      self.placeholders['num_particles']: 100}\n        if self.init_with_true_state:\n            input_dict[self.placeholders['s']] = batch['s'][:, :1]\n\n        if return_particles:\n            return sess.run([self.pred_states, self.particle_list, self.particle_probs_list], input_dict)\n        else:\n            return sess.run(self.pred_states, input_dict)\n\n\n    def connect_modules(self, means, stds, state_mins, state_maxs, state_step_sizes, learn_gaussian_mle=False):\n\n        # get shapes\n        self.batch_size = tf.shape(self.placeholders['o'])[0]\n        self.seq_len = tf.shape(self.placeholders['o'])[1]\n        # we use the static shape here because we need it to build the graph\n        self.action_dim = self.placeholders['a'].get_shape()[-1].value\n\n        encodings = snt.BatchApply(self.encoder)((self.placeholders['o'] - means['o']) / stds['o'])\n\n        # initialize particles\n        if self.init_with_true_state:\n            # tracking with known initial state\n            initial_particles = tf.tile(self.placeholders['s'][:, 0, tf.newaxis, :], [1, self.num_particles, 1])\n        else:\n            # global localization\n            if self.use_proposer:\n                # propose particles from observations\n                initial_particles = self.propose_particles(encodings[:, 0], self.num_particles, state_mins, state_maxs)\n            else:\n                # sample particles randomly\n                initial_particles = tf.concat(\n                    [tf.random_uniform([self.batch_size, self.num_particles, 1], state_mins[d], state_maxs[d]) for d in\n                     range(self.state_dim)], axis=-1, name='particles')\n\n        initial_particle_probs = tf.ones([self.batch_size, self.num_particles],\n                                         name='particle_probs') / self.num_particles_float\n\n        # assumes that samples has the correct size\n        def permute_batch(x, samples):\n            # get shapes\n            batch_size = tf.shape(x)[0]\n            num_particles = tf.shape(x)[1]\n            sample_size = tf.shape(samples)[1]\n            # compute 1D indices into the 2D array\n            idx = samples + num_particles * tf.tile(\n                tf.reshape(tf.range(batch_size), [batch_size, 1]),\n                [1, sample_size])\n            # index using the 1D indices and reshape again\n            result = tf.gather(tf.reshape(x, [batch_size * num_particles, -1]), idx)\n            result = tf.reshape(result, tf.shape(x[:,:sample_size]))\n            return result\n\n\n        def loop(particles, particle_probs, particle_list, particle_probs_list, additional_probs_list, i):\n\n            num_proposed_float = tf.round((self.propose_ratio ** tf.cast(i, tf.float32)) * self.num_particles_float)\n            num_proposed = tf.cast(num_proposed_float, tf.int32)\n            num_resampled_float = self.num_particles_float - num_proposed_float\n            num_resampled = tf.cast(num_resampled_float, tf.int32)\n\n            if self.propose_ratio < 1.0:\n\n                # resampling\n                basic_markers = tf.linspace(0.0, (num_resampled_float - 1.0) / num_resampled_float, num_resampled)\n                random_offset = tf.random_uniform([self.batch_size], 0.0, 1.0 / num_resampled_float)\n                markers = random_offset[:, None] + basic_markers[None, :]  # shape: batch_size x num_resampled\n                cum_probs = tf.cumsum(particle_probs, axis=1)\n                marker_matching = markers[:, :, None] < cum_probs[:, None, :]  # shape: batch_size x num_resampled x num_particles\n                samples = tf.cast(tf.argmax(tf.cast(marker_matching, 'int32'), dimension=2), 'int32')\n                standard_particles = permute_batch(particles, samples)\n                standard_particle_probs = tf.ones([self.batch_size, num_resampled])\n                standard_particles = tf.stop_gradient(standard_particles)\n                standard_particle_probs = tf.stop_gradient(standard_particle_probs)\n\n                # motion update\n                if learn_gaussian_mle:\n                    standard_particles, _ = self.motion_update(self.placeholders['a'][:, i], standard_particles, means, stds, state_step_sizes, learn_gaussian_mle)\n                else:\n                    standard_particles = self.motion_update(self.placeholders['a'][:, i], standard_particles, means, stds, state_step_sizes, learn_gaussian_mle)\n\n\n                # measurement update\n                standard_particle_probs *= self.measurement_update(encodings[:, i], standard_particles, means, stds)\n\n            if self.propose_ratio > 0.0:\n\n                # proposed particles\n                proposed_particles = self.propose_particles(encodings[:, i], num_proposed, state_mins, state_maxs)\n                proposed_particle_probs = tf.ones([self.batch_size, num_proposed])\n\n\n            # NORMALIZE AND COMBINE PARTICLES\n            if self.propose_ratio == 1.0:\n                particles = proposed_particles\n                particle_probs = proposed_particle_probs\n\n            elif self.propose_ratio == 0.0:\n                particles = standard_particles\n                particle_probs = standard_particle_probs\n\n            else:\n                standard_particle_probs *= (num_resampled_float / self.num_particles_float) / tf.reduce_sum(standard_particle_probs, axis=1, keep_dims=True)\n                proposed_particle_probs *= (num_proposed_float / self.num_particles_float) / tf.reduce_sum(proposed_particle_probs, axis=1, keep_dims=True)\n                particles = tf.concat([standard_particles, proposed_particles], axis=1)\n                particle_probs = tf.concat([standard_particle_probs, proposed_particle_probs], axis=1)\n\n            # NORMALIZE PROBABILITIES\n            particle_probs /= tf.reduce_sum(particle_probs, axis=1, keep_dims=True)\n\n            particle_list = tf.concat([particle_list, particles[:, tf.newaxis]], axis=1)\n            particle_probs_list = tf.concat([particle_probs_list, particle_probs[:, tf.newaxis]], axis=1)\n\n            return particles, particle_probs, particle_list, particle_probs_list, additional_probs_list, i + 1\n\n        # reshapes and sets the first shape sizes to None (which is necessary to keep the shape consistent in while loop)\n        particle_list = tf.reshape(initial_particles,\n                                   shape=[self.batch_size, -1, self.num_particles, self.state_dim])\n        particle_probs_list = tf.reshape(initial_particle_probs, shape=[self.batch_size, -1, self.num_particles])\n        additional_probs_list = tf.reshape(tf.ones([self.batch_size, self.num_particles, 4]), shape=[self.batch_size, -1, self.num_particles, 4])\n\n        # run the filtering process\n        particles, particle_probs, particle_list, particle_probs_list, additional_probs_list, i = tf.while_loop(\n            lambda *x: x[-1] < self.seq_len, loop,\n            [initial_particles, initial_particle_probs, particle_list, particle_probs_list, additional_probs_list,\n             tf.constant(1, dtype='int32')], name='loop')\n\n        # compute mean of particles\n        self.pred_states = self.particles_to_state(particle_list, particle_probs_list)\n        self.particle_list = particle_list\n        self.particle_probs_list = particle_probs_list\n\n        return particles, particle_probs, encodings, particle_list, particle_probs_list\n\n    def particles_to_state(self, particle_list, particle_probs_list):\n        mean_position = tf.reduce_sum(particle_probs_list[:, :, :, tf.newaxis] * particle_list[:, :, :, :2], axis=2)\n        mean_orientation = atan2(\n            tf.reduce_sum(particle_probs_list[:, :, :, tf.newaxis] * tf.cos(particle_list[:, :, :, 2:3]), axis=2),\n            tf.reduce_sum(particle_probs_list[:, :, :, tf.newaxis] * tf.sin(particle_list[:, :, :, 2:3]), axis=2))\n        mean_velocity = tf.reduce_sum(particle_probs_list[:, :, :, tf.newaxis] * particle_list[:, :, :, 3:5], axis=2)\n        return tf.concat([mean_position, mean_orientation, mean_velocity], axis=2)\n\n\n    def plot_motion_model(self, sess, batch, motion_samples, task, state_step_sizes):\n\n        # define the inputs and train/run the model\n        input_dict = {**{self.placeholders[key]: batch[key] for key in 'osa'},\n                      **{self.placeholders['num_particles']: 100},\n                      }\n\n        s_motion_samples = sess.run(motion_samples, input_dict)\n\n        plt.figure('Motion Model')\n        plt.gca().clear()\n        for i in range(min(len(s_motion_samples), 10)):\n            plt.scatter(s_motion_samples[i, :, 3] / state_step_sizes[3], s_motion_samples[i, :, 4] / state_step_sizes[4], color='blue', s=1)\n            plt.scatter(batch['s'][i, 0, 3] / state_step_sizes[3], batch['s'][i, 0, 4] / state_step_sizes[4], color='black', s=1)\n            plt.scatter(batch['s'][i, 1, 3] / state_step_sizes[3], batch['s'][i, 1, 4] / state_step_sizes[4], color='red', s=3)\n            plt.plot(batch['s'][i, :2, 3] / state_step_sizes[3], batch['s'][i, :2, 4] / state_step_sizes[4], color='black')\n\n        plt.xlim([0, 200])\n        plt.ylim([-50, 50])\n        plt.xlabel('translational vel')\n        plt.ylabel('angular vel')\n        plt.gca().set_aspect('equal')\n        plt.pause(0.01)\n\n\n    def plot_measurement_model(self, sess, batch_iterator, measurement_model_out):\n\n        batch = next(batch_iterator)\n\n        # define the inputs and train/run the model\n        input_dict = {**{self.placeholders[key]: batch[key] for key in 'osa'},\n                      **{self.placeholders['num_particles']: 100},\n                      }\n\n        s_measurement_model_out = sess.run([measurement_model_out], input_dict)\n\n        plt.figure('Measurement Model Output')\n        plt.gca().clear()\n        plt.imshow(s_measurement_model_out[0], interpolation=\"nearest\", cmap=\"viridis_r\", vmin=0.0, vmax=1.0)\n\n        plt.figure('Measurement Model Input')\n        plt.clf()\n        plt.scatter(batch['s'][:1, 0, 3], batch['s'][:1, 0, 4], marker='x', c=s_measurement_model_out[0][0,:1], vmin=0, vmax=1.0, cmap='viridis_r')\n        plt.scatter(batch['s'][1:, 0, 3], batch['s'][1:, 0, 4], marker='o', c=s_measurement_model_out[0][0,1:], vmin=0, vmax=1.0, cmap='viridis_r')\n        plt.xlabel('x_dot')\n        plt.ylabel('theta_dot')\n        plt.pause(0.01)\n\n\n    def plot_particle_proposer(self, sess, batch, proposed_particles, task):\n\n        # define the inputs and train/run the model\n        input_dict = {**{self.placeholders[key]: batch[key] for key in 'osa'},\n                      **{self.placeholders['num_particles']: 100},\n                      }\n\n        s_samples = sess.run(proposed_particles, input_dict)\n\n        plt.figure('Particle Proposer')\n        plt.gca().clear()\n        plot_maze(task)\n\n        for i in range(min(len(s_samples), 10)):\n            color = np.random.uniform(0.0, 1.0, 3)\n            plt.quiver(s_samples[i, :, 0], s_samples[i, :, 1], np.cos(s_samples[i, :, 2]), np.sin(s_samples[i, :, 2]), color=color, width=0.001, scale=100)\n            plt.quiver(batch['s'][i, 0, 0], batch['s'][i, 0, 1], np.cos(batch['s'][i, 0, 2]), np.sin(batch['s'][i, 0, 2]), color=color, scale=50, width=0.003)\n\n        plt.pause(0.01)\n\n\n    def plot_particle_filter(self, sess, batch, particle_list,\n                        particle_probs_list, state_step_sizes, task):\n\n        s_states, s_particle_list, s_particle_probs_list, \\\n            = sess.run([self.placeholders['s'], particle_list,\n                        particle_probs_list], #self.noise_scaler1(1.0), self.noise_scaler2(2.0)],\n                       {**{self.placeholders[key]: batch[key] for key in 'osa'},\n                        **{self.placeholders['num_particles']: 20},\n                        })\n        # print('learned motion noise factors {:.2f}/{:.2f}'.format(n1, n2))\n\n        num_steps = s_particle_list.shape[1]\n\n        for s in range(3):\n\n            plt.figure('particle_evolution, example {}'.format(s))\n            plt.clf()\n\n            for d in range(5):\n\n                plt.subplot(3, 2, [1, 3, 5, 2, 4][d])\n\n                for i in range(num_steps):\n\n                    plt.scatter(i * np.ones_like(s_particle_list[s, i, :, d]),\n                                s_particle_list[s, i, :, d] / (1 if s == 0 else state_step_sizes[d]),\n                                c=s_particle_probs_list[s, i, :], cmap='viridis_r', marker='o', s=6, alpha=0.5,\n                                linewidths=0.05,\n                                vmin=0.0,\n                                vmax=0.1)\n                    current_state = batch['s'][s, i, d] / (1 if s == 0 else state_step_sizes[d])\n                    plt.plot([i], [current_state], 'o', markerfacecolor='None', markeredgecolor='k',\n                             markersize=2.5)\n\n                plt.xlabel('Time')\n                plt.ylabel('State {}'.format(d))\n\n        show_pause(pause=0.01)\n"
  },
  {
    "path": "methods/odom.py",
    "content": "import numpy as np\n\nfrom utils.data_utils_kitti import wrap_angle\n\nclass OdometryBaseline():\n\n    def __init__(self, *args, **kwargs):\n        pass\n\n    def fit(self, *args, **kwargs):\n        pass\n\n    def predict(self, sess, batch, **kwargs):\n        seq_len = batch['s'].shape[1]\n\n        prediction = np.zeros_like(batch['s'])\n        state = batch['s'][:, 0, :]\n        # print('shape:', batch['s'].shape)\n        prediction[:, 0, :] = state\n        for i in range(1, seq_len):\n\n            action = batch['a'][:, i, :]\n            theta = state[:, 2:3]\n            sin_theta = np.sin(theta)\n            cos_theta = np.cos(theta)\n            new_x = state[:, 0:1] + (action[:, 0:1] * cos_theta + action[:, 1:2] * sin_theta)\n            new_y = state[:, 1:2] + (action[:, 0:1] * sin_theta - action[:, 1:2] * cos_theta)\n            new_theta = wrap_angle(state[:, 2:3] + action[:, 2:3])\n            # copy old and set new particles\n            state = np.concatenate([new_x, new_y, new_theta], axis=-1)\n            prediction[:, i, :] = state\n        return prediction\n\n    def predict_kitti(self, sess, batch, **kwargs):\n        seq_len = batch['s'].shape[1]\n\n        prediction = np.zeros_like(batch['s'])\n        state = batch['s'][:, 0, :]\n        # print('shape:', batch['s'].shape)\n        prediction[:, 0, :] = state\n        for i in range(1, seq_len):\n\n            time = 0.103\n\n            action = batch['a'][:, i, :]\n            heading = state[:, 2:3]\n            wrap_angle(heading)\n            sin_heading = np.sin(heading)\n            cos_heading = np.cos(heading)\n\n            # ang_acc = (noisy_actions[:, :, 1:2] * noisy_actions[:, :, 2:3])/(noisy_actions[:, :, 0:1] ** 2)\n\n            acc_north = action[:, 0:1] * sin_heading + action[:, 1:2] * cos_heading\n            acc_east = - action[:, 1:2] * sin_heading + action[:, 0:1] * cos_heading\n\n            new_north = state[:, 0:1] + state[:, 3:4] * time\n            new_east = state[:, 1:2] + state[:, 4:5] * time\n            new_theta = state[:, 2:3] + state[:, 5:6] * time\n            wrap_angle(new_theta)\n            new_vn = state[:, 3:4] + acc_north * time\n            new_ve = state[:, 4:5] + acc_east * time\n            new_theta_dot = state[:, 5:6] + action[:, 2:3] * time\n\n            state = np.concatenate([new_north, new_east, new_theta, new_vn, new_ve, new_theta_dot], axis=-1)\n            prediction[:, i, :] = state\n        return prediction\n"
  },
  {
    "path": "methods/rnn.py",
    "content": "import tensorflow as tf\nimport sonnet as snt\n\nfrom utils.data_utils import *\nfrom utils.method_utils import compute_sq_distance\n\nclass RNN():\n    def __init__(self, init_with_true_state=False, model='2lstm', **unused_kwargs):\n\n        self.placeholders = {'o': tf.placeholder('float32', [None, None, 24, 24, 3], 'observations'),\n                     'a': tf.placeholder('float32', [None, None, 3], 'actions'),\n                     's': tf.placeholder('float32', [None, None, 3], 'states'),\n                     'keep_prob': tf.placeholder('float32')}\n        self.pred_states = None\n        self.init_with_true_state = init_with_true_state\n        self.model = model\n\n        # build models\n        # <-- observation\n        self.encoder = snt.Sequential([\n            snt.nets.ConvNet2D([16, 32, 64], [[3, 3]], [2], [snt.SAME], activate_final=True, name='encoder/convnet'),\n            snt.BatchFlatten(),\n            lambda x: tf.nn.dropout(x, self.placeholders['keep_prob']),\n            snt.Linear(128, name='encoder/Linear'),\n            tf.nn.relu,\n        ])\n\n        # <-- action\n        if self.model == '2lstm':\n            self.rnn1 = snt.LSTM(512)\n            self.rnn2 = snt.LSTM(512)\n        if self.model == '2gru':\n            self.rnn1 = snt.GRU(512)\n            self.rnn2 = snt.GRU(512)\n        elif self.model == 'ff':\n            self.ff_lstm_replacement = snt.Sequential([\n                snt.Linear(512),\n                tf.nn.relu,\n                snt.Linear(512),\n                tf.nn.relu])\n\n        self.belief_decoder = snt.Sequential([\n            snt.Linear(256),\n            tf.nn.relu,\n            snt.Linear(256),\n            tf.nn.relu,\n            snt.Linear(3)\n        ])\n\n\n    def fit(self, sess, data, model_path, split_ratio, seq_len, batch_size, epoch_length, num_epochs, patience, learning_rate, dropout_keep_ratio, **unused_kwargs):\n\n        # preprocess data\n        data = split_data(data, ratio=split_ratio)\n        epoch_lengths = {'train': epoch_length, 'val': epoch_length*2}\n        batch_iterators = {'train': make_batch_iterator(data['train'], batch_size=batch_size, seq_len=seq_len),\n                           'val': make_repeating_batch_iterator(data['val'], epoch_lengths['val'], batch_size=batch_size, seq_len=seq_len),\n                           'train_ex': make_batch_iterator(data['train'], batch_size=batch_size, seq_len=seq_len),\n                           'val_ex': make_batch_iterator(data['val'], batch_size=batch_size, seq_len=seq_len)}\n        means, stds, state_step_sizes, state_mins, state_maxs = compute_staticstics(data['train'])\n\n        self.connect_modules(means, stds, state_mins, state_maxs, state_step_sizes)\n\n        # training\n\n        sq_dist = compute_sq_distance(self.pred_states, self.placeholders['s'], state_step_sizes)\n        losses = {'mse': tf.reduce_mean(sq_dist),\n                  'mse_last': tf.reduce_mean(sq_dist[:, -1])}\n\n        optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n        gradients = optimizer.compute_gradients(losses['mse'])\n        # clipped_gradients = [(tf.clip_by_value(grad, -1., 1.), var) for grad, var in gradients]\n        train_op = optimizer.apply_gradients(gradients)\n\n        init = tf.global_variables_initializer()\n        sess.run(init)\n\n        # save statistics and prepare saving variables\n        if not os.path.exists(model_path):\n            os.makedirs(model_path)\n        np.savez(os.path.join(model_path, 'statistics'), means=means, stds=stds, state_step_sizes=state_step_sizes,\n                 state_mins=state_mins, state_maxs=state_maxs)\n        saver = tf.train.Saver()\n        save_path = model_path + '/best_validation'\n\n        loss_keys = ['mse_last', 'mse']\n        if split_ratio < 1.0:\n            data_keys = ['train', 'val']\n        else:\n            data_keys = ['train']\n\n        log = {dk: {lk: {'mean': [], 'se': []} for lk in loss_keys} for dk in data_keys}\n\n        best_val_loss = np.inf\n        best_epoch = 0\n        i = 0\n        while i < num_epochs and i - best_epoch < patience:\n            # training\n            loss_lists = dict()\n            for dk in data_keys:\n                loss_lists = {lk: [] for lk in loss_keys}\n                for e in range(epoch_lengths[dk]):\n                    batch = next(batch_iterators[dk])\n                    if dk == 'train':\n                        s_losses, _ = sess.run([losses, train_op], {**{self.placeholders[key]: batch[key] for key in 'osa'},\n                                                                **{self.placeholders['keep_prob']: dropout_keep_ratio}})\n                    else:\n                        s_losses = sess.run(losses, {**{self.placeholders[key]: batch[key] for key in 'osa'},\n                                                            **{self.placeholders['keep_prob']: 1.0}})\n                    for lk in loss_keys:\n                        loss_lists[lk].append(s_losses[lk])\n                # after each epoch, compute and log statistics\n                for lk in loss_keys:\n                    log[dk][lk]['mean'].append(np.mean(loss_lists[lk]))\n                    log[dk][lk]['se'].append(np.std(loss_lists[lk], ddof=1) / np.sqrt(epoch_lengths[dk]))\n\n            # check whether the current model is better than all previous models\n            if 'val' in data_keys:\n                if log['val']['mse_last']['mean'][-1] < best_val_loss:\n                    best_val_loss = log['val']['mse_last']['mean'][-1]\n                    best_epoch = i\n                    # save current model\n                    saver.save(sess, save_path)\n                    txt = 'epoch {:>3} >> '.format(i)\n                else:\n                    txt = 'epoch {:>3} == '.format(i)\n            else:\n                best_epoch = i\n                saver.save(sess, save_path)\n                txt = 'epoch {:>3} >> '.format(i)\n\n            # after going through all data sets, do a print out of the current result\n            for lk in loss_keys:\n                txt += '{}: '.format(lk)\n                for dk in data_keys:\n                    txt += '{:.2f}+-{:.2f}/'.format(log[dk][lk]['mean'][-1], log[dk][lk]['se'][-1])\n                txt = txt[:-1] + ' -- '\n            print(txt)\n\n            i += 1\n\n            # for key in ['train', 'val']:\n            #     batch = next(batch_iterators[key + '_ex'])\n            #     s_states, s_pred_states = sess.run([self.placeholders['s'], self.pred_states], {**{self.placeholders[key]: batch[key] for key in 'osa'},\n            #                            **{self.placeholders['keep_prob']: 1.0}})\n            #\n            #     # s_pred_states = np.argmax(np.reshape(s_pred_states, list(s_pred_states.shape[:2]) + [10,5,8]), axis=2) * 100\n            #\n            #     plt.figure('Example: ' + key)\n            #     plt.gca().clear()\n            #     plot_maze('nav01')\n            #     s_states = np.reshape(s_states, [-1, 3])\n            #     s_pred_states = np.reshape(s_pred_states, [-1, 3])\n            #     plt.plot(s_states[:, 0], s_states[:, 1], 'xb')\n            #     plt.plot(s_pred_states[:, 0], s_pred_states[:, 1], 'xg' if key == 'val' else 'xr')\n            #     errors = np.concatenate([s_states[:, np.newaxis, :], s_pred_states[:, np.newaxis, :]], axis=1)\n            #     plt.plot(errors[:, :, 0].T, errors[:, :, 1].T, '-k')\n            #\n            #     # plt.plot(np.argmax(np.amax(np.amax(np.reshape(s_belief, list(s_belief.shape[:2]) + [10, 5, 8]), axis=4), axis=3), axis=2) * 100 + 50,\n            #     #          np.argmax(np.amax(np.amax(np.reshape(s_belief, list(s_belief.shape[:2]) + [10, 5, 8]), axis=4), axis=2), axis=2) * 100 + 50, 'xg' if key == 'val' else 'xr')\n            #     # plt.plot(s_pred_states[:, :, 0], s_pred_states[:, :, 1], 'xg' if key == 'val' else 'xr')\n            #\n            #     show_pause(pause=0.01)\n            # else:\n            #     print('epoch {} -- mse: {:.4f}'.format(e, log['train']['mse'][-1]))\n            #     # plt.figure('Learning curve: {}'.format(key))\n            #     # plt.gca().clear()\n            #     # plt.plot(log['train'][key], '--k')\n            #     # plt.plot(log['val'][key], '-k')\n            #     # plt.ylim([0, max(log['val'][key])])\n\n        saver.restore(sess, save_path)\n\n        return log\n\n\n    def connect_modules(self, means, stds, state_mins, state_maxs, state_step_sizes):\n\n        # tracking_info_full = tf.tile(((self.placeholders['s'] - means['s']) / stds['s'])[:, :1, :], [1, tf.shape(self.placeholders['s'])[1], 1])\n        tracking_info = tf.concat([((self.placeholders['s'] - means['s']) / stds['s'])[:, :1, :], tf.zeros_like(self.placeholders['s'][:,1:,:])], axis=1)\n        flag = tf.concat([tf.ones_like(self.placeholders['s'][:,:1,:1]), tf.zeros_like(self.placeholders['s'][:,1:,:1])], axis=1)\n\n        preproc_o = snt.BatchApply(self.encoder)((self.placeholders['o'] - means['o']) / stds['o'])\n        # include tracking info\n        if self.init_with_true_state:\n            # preproc_o = tf.concat([preproc_o, tracking_info, flag], axis=2)\n            preproc_o = tf.concat([preproc_o, tracking_info, flag], axis=2)\n            # preproc_o = tf.concat([preproc_o, tracking_info_full], axis=2)\n\n        preproc_a = snt.BatchApply(snt.BatchFlatten())(self.placeholders['a'] / stds['a'])\n        preproc_ao = tf.concat([preproc_o, preproc_a], axis=-1)\n\n        if self.model == '2lstm' or self.model == '2gru':\n            lstm1_out, lstm1_final_state = tf.nn.dynamic_rnn(self.rnn1, preproc_ao, dtype=tf.float32)\n            lstm2_out, lstm2_final_state = tf.nn.dynamic_rnn(self.rnn2, lstm1_out, dtype=tf.float32)\n            belief_list = lstm2_out\n\n        elif self.model == 'ff':\n            belief_list = snt.BatchApply(self.ff_lstm_replacement)(preproc_ao)\n\n        self.pred_states = snt.BatchApply(self.belief_decoder)(belief_list)\n        self.pred_states = self.pred_states * stds['s'] + means['s']\n\n\n    def predict(self, sess, batch, **unused_kwargs):\n        return sess.run(self.pred_states, {**{self.placeholders[key]: batch[key] for key in 'osa'},\n                                           **{self.placeholders['keep_prob']: 1.0}})\n\n    def load(self, sess, model_path, model_file='best_validation', statistics_file='statistics.npz', connect_and_initialize=True):\n\n        # build the tensorflow graph\n        if connect_and_initialize:\n            # load training data statistics (which are needed to build the tf graph)\n            statistics = dict(np.load(os.path.join(model_path, statistics_file)))\n            for key in statistics.keys():\n                if statistics[key].shape == ():\n                    statistics[key] = statistics[key].item()  # convert 0d array of dictionary back to a normal dictionary\n\n            # connect all modules into the particle filter\n            self.connect_modules(**statistics)\n            init = tf.global_variables_initializer()\n            sess.run(init)\n\n        # load variables\n        all_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)\n        for v in all_vars:\n            print(\"%s %r %s\" % (v, v, v.shape))\n\n        # restore variable values\n        saver = tf.train.Saver()  # <- var list goes in here\n        saver.restore(sess, os.path.join(model_path, model_file))\n\n        # print('Loaded the following variables:')\n        # for v in all_vars:\n        #     print(v.name)\n"
  },
  {
    "path": "plotting/__init__.py",
    "content": ""
  },
  {
    "path": "plotting/ab_plot.py",
    "content": "import pickle\nimport numpy as np\nimport matplotlib\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D\nimport itertools\nimport os\n\nresults = None\n\n# matplotlib.rcParams.update({'font.size': 12})\n\ncolor_list = plt.cm.tab10(np.linspace(0, 1, 10))\ncolors = {'lstm': color_list[0], 'pf_e2e': color_list[1], 'pf_ind_e2e': color_list[2], 'pf_ind': color_list[3]}\nlabels = {'lstm': 'LSTM', 'pf_e2e': 'e2e', 'pf_ind_e2e': 'ind+e2e', 'pf_ind': 'ind'}\n# conditions = ['normal', 'no_motion_likelihood', 'learn_odom', 'no_proposer']\n# conditions = ['normal', 'learn_odom', 'no_inject']\n# clabels = {'normal': 'Default', 'no_motion_likelihood': 'W/o motion likelihood', 'learn_odom': 'Learned odometry', 'no_proposer': 'W/o particle proposer', 'no_inject': \"No inject\"}\nconditions = ['full', 'learn_odom', 'no_inject', 'no_proposer']\nclabels = {'full': 'Full', 'learn_odom': 'Learned\\nodometry', 'no_proposer': 'No particle\\nproposer', 'no_inject': \"No particle\\ninjection\"}\ntasks = ['nav02']\nmethods = ['pf_ind', 'pf_e2e', 'pf_ind_e2e']\n\n# load results\nresults = dict()\n\ncount = 0\nfor task in tasks:\n    # log_path = '/home/rbo/Desktop/log/'+task+'_ab1'\n    log_path = '../log/ab'\n    for filename in [f for f in os.listdir(log_path) if os.path.isfile(os.path.join(log_path, f))]:\n        full_filename = os.path.join(log_path, filename)\n        print('loading {}:'.format(count) + full_filename + ' ...')\n        try:\n            # if 'DeepThought' not in filename:\n            # if 'DeepThought' in filename:\n            with open(full_filename, 'rb') as f:\n                result = pickle.load(f)\n                # result_name = result['task'][0] + '/' + result['method'][0] + '/' + str(result['num_episodes'][0]) + '/' + result['condition'][0]\n                result_name = result['exp_params'][0]['file_ending'] #result['exp_params'][0]['task'] + '/' + result['exp_params'][0]['method'] + '/' + str(result['exp_params'][0]['num_episodes']) + '/' + result['exp_params'][0]['ab_cond']\n                for ab_cond in conditions:\n                    if result_name.endswith(ab_cond):\n                        result['exp_params'][0]['ab_cond'] = ab_cond\n                print(result_name)\n                if result_name not in results.keys():\n                    results[result_name] = result\n                else:\n                    for key in result.keys():\n                        if key in results[result_name].keys():\n                            results[result_name][key] += result[key]\n                        else:\n                            results[result_name][key] = result[key]\n                        # print(result_name, key)\n                count += 1\n        except Exception as e:\n            print(e)\n            print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')\n\nfor result_name in results.keys():\n    print(result_name, len(results[result_name]['exp_params'][0]['task']))\n\nprint('Loaded {} results'.format(count))\n\n# print(results['test_errors'].shape, np.mean(results['test_errors']**2, axis=1))\n\n#print('SHAPE', results['test_mse'].shape)\n\n# plt.figure(1)\n# plt.gca().set_color_cycle(None)\n# for method in set(results['method']):\n\ntask = 'nav02'\n# step = 30\nstep = 2\n\nepisodes = [16, 125, 1000]\n# episodes = [1000]\nfig_names = []\n\nmax_1 = 0\nmax_2 = {n: 0 for n in episodes}\n\nmeans = dict()\nses = dict()\n\nfig_name = 'abcolorbar'\nplt.figure(fig_name, [0.8,2.5])\nfig_names.append(fig_name)\n\nvmax=1.0\na = np.array([[0.0, 1.0]])\nimg = plt.imshow(a, cmap=\"viridis_r\", vmin=-0.33*vmax, vmax=vmax)\nplt.gca().set_visible(False)\ncax = plt.axes([0.0, 0.2, 0.1, 0.65])\nplt.colorbar(orientation=\"vertical\", cax=cax, label='Error rate', boundaries=np.linspace(0,1.0,100), ticks=np.linspace(0.0, 1.0, 11))\n\nfor num_episodes in episodes:\n\n    means[num_episodes] = dict()\n    ses[num_episodes] = dict()\n\n    for method in methods:\n\n        means[num_episodes][method] = np.zeros([len(conditions), 5])\n        # means[num_episodes][method] = np.zeros([len(conditions), 50])\n        ses[num_episodes][method] = np.zeros([len(conditions), 5])\n        # ses[num_episodes][method] = np.zeros([len(conditions), 50])\n\n        for c, condition in enumerate(conditions):\n\n            result_name = task + '_' + method + '_' + str(num_episodes) + '_' + condition\n            if result_name in results.keys():\n                result = results[result_name]\n\n                # means[num_episodes][method][c] = np.mean(result['test_mse'], axis=0)\n                # std = np.std(result['test_mse'], axis=0, ddof=1)\n                # ses[num_episodes][method][c] = std / np.sqrt(len(result['test_mse']))\n\n                hist = np.array([[h[i] for i in range(0, 50, 10)] for h in result['test_hist']])  # result x time x sqe [.0, 0.1, .., 10.0]\n                err = 1. - np.sum(hist[:,:,:10], axis=-1) # sqe < 1.0\n                means[num_episodes][method][c] = np.mean(err, axis=0)\n                print(result_name)\n                print(err[:, step])\n                print(np.mean(err, axis=0)[step])\n                print(np.std(err, axis=0, ddof=1)[step], np.sqrt(len(err)))\n                ses[num_episodes][method][c] = np.std(err, axis=0, ddof=1) / np.sqrt(len(err))\n\n            else:\n                # print(result_name, 0)\n                means[num_episodes][method][c] *= np.nan\n                ses[num_episodes][method][c] *= np.nan\n\n\n    means[num_episodes]['min'] = np.stack([means[num_episodes][method] for method in methods], axis=0).min(axis=1)\n\n    fig_name = 'ab1_{}'.format(num_episodes)\n    plt.figure(fig_name, [3,2.5])\n    fig_names.append(fig_name)\n\n    # m = means[num_episodes][method][:,step-1]\n    m = np.stack([means[num_episodes][method][:, step] for method in methods], axis=0)\n    s = np.stack([ses[num_episodes][method][:, step] for method in methods], axis=0)\n    is_min = m == means[num_episodes]['min'][:, None, step]\n\n    # plt.imshow((means[:,:,30-1])**0.5, interpolation='nearest', vmin=0, vmax=15)\n    # plt.imshow(np.log(m.T), interpolation='nearest', vmin=-2.5, vmax=2.5, cmap='viridis')\n    plt.imshow(m.T, interpolation='nearest', vmin=-0.33, vmax=1.0, cmap='viridis_r')\n    # data = np.reshape(np.arange(len(conditions)*len(conditions)), [len(conditions), len(conditions)])\n    # plt.imshow(data, interpolation='nearest', vmin=0, vmax=10)\n    plt.yticks(np.arange(len(conditions)), [clabels[c] for c in conditions])\n    plt.xticks(np.arange(len(methods)), [labels[m] for m in methods])\n\n    # plt.xlabel('Test {}{} noise'.format(noise_in, '.' if noise_in == 'odom' else ''))\n    # plt.ylabel('Training {}{} noise'.format(noise_in, '.' if noise_in == 'odom' else ''))\n\n    # plt.colorbar()\n    #text portion\n    # min_val, max_val, diff = 0., len(conditions), 1.\n    # N_points = (max_val - min_val) / diff\n    ind_array_y = np.arange(0., len(methods), 1.)\n    ind_array_x = np.arange(0., len(conditions), 1.)\n    x, y = np.meshgrid(ind_array_x, ind_array_y)\n\n    for x_val, y_val in zip(x.flatten(), y.flatten()):\n        value = m[int(y_val),int(x_val)]\n        s_value = s[int(y_val),int(x_val)]\n        text = '{:.4s}\\n+-{:.4s}'.format('{:.3f}'.format(value)[1:],'{:.2f}'.format(s_value)[1:])\n        plt.text(y_val, x_val, text, va='center', ha='center', color='white', fontweight='bold' if is_min[int(y_val), int(x_val)] else 'normal')\n        # plt.text(y_val, x_val, text, va='center', ha='center', color='white', fontweight='normal')\n\n        # fig_name = 'nt_diag_{}'.format(num_episodes)\n        # plt.figure(fig_name, [3,2.5])\n        # fig_names.append(fig_name)\n        #\n        # x = np.arange(len(conditions))\n        # m = means[num_episodes][method][:,:,step-1]\n        # s = ses[num_episodes][method][:,:,step-1]\n        # plt.plot(x[:-1], np.diag(m)[:-1], '-', color=colors[method], label=labels[method])\n        # ind = -3 if noise_in == 'odom' else -2\n        # plt.plot(x[ind], np.diag(m)[ind], 'x', color=colors[method])\n        # plt.fill_between(x[:-1], np.diag(m-s)[:-1], np.diag(m+s)[:-1], color=colors[method], alpha=0.5, linewidth=0.0)\n        # plt.xticks(np.arange(len(conditions)-1), [clabels[c] for c in conditions[:-1]])\n        # if noise_in == 'odom':\n        #     plt.xlabel('Gaussian odometry noise in %')\n        # else:\n        #     plt.xlabel('Image noise')\n        #     plt.legend()\n        #\n        # plt.ylabel('Test MSE ({} episodes)'.format(num_episodes))\n        # plt.ylim([0, 2.5])\n        #\n        # fig_name = 'nt_shuffle_{}'.format(num_episodes)\n        # plt.figure(fig_name, [3,2.5])\n        # fig_names.append(fig_name)\n        #\n        # if noise_in == 'odom':\n        #     plt.bar(0.0  - 0.5 + (methods.index(method)+1)/len(methods)*0.8,\n        #             np.diag(m)[-2],\n        #             0.8/len(methods),\n        #             yerr=np.diag(s)[-2],\n        #             color=colors[method], label=labels[method])\n        # plt.bar((1.0 if noise_in == 'odom' else 2.0)  - 0.5 + (methods.index(method)+1)/len(methods)*0.8,\n        #         np.diag(m)[-1], 0.8/len(methods),\n        #         yerr=np.diag(s)[-1],\n        #         color=colors[method])\n        # relative = np.diag(m)[-1] / np.diag(m)[-2]\n        # textpos = np.diag(m)[-1] + np.diag(s)[-1] + 2\n        # if num_episodes == 1000:\n        #     if textpos > 10:\n        #         textpos = 3\n        # elif textpos > 80:\n        #         textpos = 10\n        # plt.text((1.0 if noise_in == 'odom' else 2.0)  - 0.5 + (methods.index(method)+1)/len(methods)*0.8,\n        #          textpos, '×{:.0f}'.format(relative), va='bottom', ha='center',color='black', rotation=90)\n        # plt.xticks([0, 1, 2], ['Both', 'Image', 'Odom.'])\n        # plt.ylim([0,60])\n        # plt.xlabel('Input')\n        # plt.ylabel('Test MSE ({} episodes)'.format(num_episodes))\n        # # plt.legend()\n\n\nfor fn in fig_names:\n    plt.figure(fn)\n    plt.tight_layout()\n    plt.savefig('../plots/ab/{}.pdf'.format(fn), bbox_inches=\"tight\", transparent=True, dpi=600, frameon=True, facecolor='w', pad_inches=0.01)\n    # plt.savefig('../plots/ab/{}.png'.format(fn), bbox_inches=\"tight\", transparent=True, dpi=600, frameon=True, facecolor='w', pad_inches=0.01)\n\n            # plt.bar(np.arange(len(conditions)) + (methods.index(method)+1)/len(methods)*0.8 - 0.5, np.diag(means[:,:,30-1])/np.diag(means[:,:,30-1])[0], 0.8/len(methods), color=colors[method])\n\n\nplt.show()\n"
  },
  {
    "path": "plotting/cross_plot.py",
    "content": "import pickle\nimport numpy as np\nimport matplotlib\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D\nimport itertools\nimport os\n\nresults = None\n\n# matplotlib.rcParams.update({'font.size': 12})\n\ncolor_list = plt.cm.tab10(np.linspace(0, 1, 10))\ncolors = {'lstm': color_list[0], 'pf_e2e': color_list[1], 'pf_ind_e2e': color_list[2], 'pf_ind': color_list[3]}\nlabels = {'lstm': 'LSTM', 'pf_e2e': 'DPF (e2e)', 'pf_ind_e2e': 'DPF (ind+e2e)', 'pf_ind': 'DPF (ind)', 'ff': 'FF', 'odom': 'Odom. baseline'}\n# conditions = ['normal', 'no_motion_likelihood', 'learn_odom', 'no_proposer']\n# conditions = ['normal', 'learn_odom', 'no_inject']\n# clabels = {'normal': 'Default', 'no_motion_likelihood': 'W/o motion likelihood', 'learn_odom': 'Learned odometry', 'no_proposer': 'W/o particle proposer', 'no_inject': \"No inject\"}\nconditions = ['lc2lc', 'pl2lc', 'mx2lc', 'lc2pl', 'pl2pl', 'mx2pl']\nclabels = {'lc2lc':'lc2lc', 'lc2pl':'lc2pl', 'pl2lc':'pl2lc', 'pl2pl':'pl2pl', 'mx2lc': 'mx2lc', 'mx2pl': 'mx2pl'}\ntask = 'nav02'\nmethods = ['pf_ind', 'pf_e2e', 'pf_ind_e2e', 'lstm']\n# methods = ['pf_ind_e2e', 'lstm']\n\n# load results\nresults = dict()\n\ncount = 0\nfor cond in conditions:\n    # log_path = '/home/rbo/Desktop/log/'+task+'_ab1'\n    log_path = '../log/'+cond\n    for filename in [f for f in os.listdir(log_path) if os.path.isfile(os.path.join(log_path, f))]:\n        full_filename = os.path.join(log_path, filename)\n        print('loading {}:'.format(count) + full_filename + ' ...')\n        try:\n            # if 'DeepThought' not in filename:\n            # if 'DeepThought' in filename:\n            with open(full_filename, 'rb') as f:\n                result = pickle.load(f)\n                # result_name = result['task'][0] + '/' + result['method'][0] + '/' + str(result['num_episodes'][0]) + '/' + result['condition'][0]\n                result_name = cond + '_' + result['exp_params'][0]['file_ending'] #result['exp_params'][0]['task'] + '/' + result['exp_params'][0]['method'] + '/' + str(result['exp_params'][0]['num_episodes']) + '/' + result['exp_params'][0]['ab_cond']\n                print(result_name)\n                if result_name not in results.keys():\n                    results[result_name] = result\n                else:\n                    for key in result.keys():\n                        if key in results[result_name].keys():\n                            results[result_name][key] += result[key]\n                        else:\n                            results[result_name][key] = result[key]\n                        # print(result_name, key)\n                count += 1\n        except Exception as e:\n            print(e)\n            print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')\n\nprint()\nfor result_name in results.keys():\n    print(result_name, len(results[result_name]['test_mse']))\n\nprint('Loaded {} results'.format(count))\n\n# print(results['test_errors'].shape, np.mean(results['test_errors']**2, axis=1))\n\n#print('SHAPE', results['test_mse'].shape)\n\n# plt.figure(1)\n# plt.gca().set_color_cycle(None)\n# for method in set(results['method']):\n\ntask = 'nav02'\n# step = 30\nstep = 3\n\nepisodes = [1000]\n# episodes = [1000]\nfig_names = []\n\nmax_1 = 0\nmax_2 = {n: 0 for n in episodes}\n\nmeans = dict()\nses = dict()\n\nfor num_episodes in episodes:\n\n    means[num_episodes] = dict()\n    ses[num_episodes] = dict()\n\n    for method in methods:\n\n        means[num_episodes][method] = np.zeros([len(conditions), 5])\n        # means[num_episodes][method] = np.zeros([len(conditions), 50])\n        ses[num_episodes][method] = np.zeros([len(conditions), 5])\n        # ses[num_episodes][method] = np.zeros([len(conditions), 50])\n\n        for c, condition in enumerate(conditions):\n\n            result_name = condition + '_' + task + '_' + method + '_' + str(num_episodes)\n            if result_name in results.keys():\n                result = results[result_name]\n\n                # means[num_episodes][method][c] = np.mean(result['test_mse'], axis=0)\n                # std = np.std(result['test_mse'], axis=0, ddof=1)\n                # ses[num_episodes][method][c] = std / np.sqrt(len(result['test_mse']))\n\n                hist = np.array([[h[i] for i in range(0, 50, 10)] for h in result['test_hist']])  # result x time x sqe [.0, 0.1, .., 10.0]\n                err = 1. - np.sum(hist[:,:,:10], axis=-1) # sqe < 1.0\n                # err = np.sum(hist[:,:,:10], axis=-1) # sqe < 1.0\n                print(result_name, err)\n                means[num_episodes][method][c] = np.mean(err, axis=0)\n                ses[num_episodes][method][c] = np.std(err, axis=0, ddof=1) / np.sqrt(len(err))\n                print(means[num_episodes][method][c])\n\n            else:\n                # print(result_name, 0)\n                means[num_episodes][method][c] *= np.nan\n                ses[num_episodes][method][c] *= np.nan\n\n\n\n    means[num_episodes]['min'] = np.stack([means[num_episodes][method] for method in methods], axis=0).min(axis=1)\n\n    fig_name = 'ab1_{}'.format(num_episodes)\n    fig = plt.figure(fig_name, [6, 3.5])\n    fig_names.append(fig_name)\n    ax = fig.add_subplot(111)\n    # Turn off axis lines and ticks of the big subplot\n    ax.spines['top'].set_color('none')\n    ax.spines['bottom'].set_color('none')\n    ax.spines['left'].set_color('none')\n    ax.spines['right'].set_color('none')\n    ax.tick_params(labelcolor='w', top='off', bottom='off', left='off', right='off')\n\n    for c, condition in enumerate(conditions):\n        sax = fig.add_subplot(2, 3, c+1)\n        for m, method in enumerate(methods):\n            sax.bar(0.0  - 0.5 + (m+1)/len(methods)*0.8,\n                            means[num_episodes][method][c, step],\n                            0.8/len(methods),\n                            yerr=ses[num_episodes][method][c, step],\n                            color=colors[method], label=labels[method])\n\n            text = '{:.3s}'.format('{:.2f}'.format(means[num_episodes][method][c, step])[1:])\n            plt.text(0.0  - 0.5 + (m+1)/len(methods)*0.8, means[num_episodes][method][c, step] + ses[num_episodes][method][c, step] + 0.05, text, va='center', ha='center', color=colors[method], fontweight='normal')\n\n\n        # sax.set_ylim([0.0, 1.05])\n        sax.set_ylim([0.0, 1.0])\n        sax.set_xticks([])\n        sax.set_yticks([])\n        # if c % 2 == 0:\n        # if c >= 2:\n        if 'lc2' in condition:\n            xlabel = 'A'\n            sax.set_ylabel(('A' if '2lc' in condition else 'B'), fontweight = 'bold')\n        elif 'pl2' in condition:\n            xlabel = 'B'\n        elif 'mx2' in condition:\n            xlabel = 'A+B'\n        if '2pl' in condition:\n            sax.set_xlabel(xlabel, fontweight = 'bold')\n        if c == 0:\n            plt.legend()\n\n    ax.set_xlabel('Trained with policy')\n    ax.set_ylabel('Error rate in test with policy\\n')\n\n    plt.tight_layout(h_pad=0.0, w_pad=0.0, pad=0.0)\n    plt.savefig('../plots/cr/policy.pdf', bbox_inches=\"tight\", transparent=True, dpi=600, frameon=True, facecolor='w', pad_inches=0.01)\n    plt.show()\n\n"
  },
  {
    "path": "plotting/lc_plot.py",
    "content": "import pickle\nimport numpy as np\nimport matplotlib\nimport matplotlib.pyplot as plt\n#from mpl_toolkits.mplot3d import Axes3D\nimport itertools\nfrom collections import namedtuple\nimport os\n\nresults = None\n\n# matplotlib.rcParams.update({'font.size': 12})\n\ncolor_list = plt.cm.tab10(np.linspace(0, 1, 10))\ncolors = {'lstm': color_list[0], 'pf_e2e': color_list[1], 'pf_ind_e2e': color_list[2], 'pf_ind': color_list[3], 'ff': color_list[4], 'odom': color_list[4]}\nlabels = {'lstm': 'LSTM', 'pf_e2e': 'DPF (e2e)', 'pf_ind_e2e': 'DPF (ind+e2e)', 'pf_ind': 'DPF (ind)', 'ff': 'FF', 'odom': 'Odom. baseline'}\n\ndef load_results(base_path='../log/', exp='lc'):\n    results = dict()\n\n    count = 0\n    log_path = os.path.join(base_path, exp)\n    listdir = os.listdir(log_path)\n    for i, filename in enumerate(listdir):\n        full_filename = os.path.join(log_path, filename)\n        # if 'DeepThought' not in full_filename:\n        print('loading ' + full_filename + ' ...')\n        # try:\n        with open(full_filename, 'rb') as f:\n            result = pickle.load(f)\n            print(result['exp_params'][0].keys())\n            result_name = result['exp_params'][0]['task'] + '/' + result['exp_params'][0]['method'] + '/' + str(result['exp_params'][0]['num_episodes'])\n            if result_name not in results.keys():\n                results[result_name] = result\n            else:\n                for key in result.keys():\n                    results[result_name][key] += result[key]\n            count += 1\n        # except Exception as e:\n        #     print(e)\n\n    for task in tasks:\n        for method in methods:\n            for num_episodes in episodes:\n                result_name = task + '/' + method + '/' + str(num_episodes)\n                if result_name in results.keys():\n                    print(result_name, len(results[result_name]['exp_params']))\n                else:\n                    print(result_name, 0)\n\n    print('Loaded {} results'.format(count))\n    return results\n\n# print(results['test_errors'].shape, np.mean(results['test_errors']**2, axis=1))\n\n#print('SHAPE', results['test_mse'].shape)\n\n# plt.figure(1)\n# plt.gca().set_color_cycle(None)\n# for method in set(results['method']):\n\n# step = {'nav01': 20, 'nav02': 20, 'nav03': 20}\n\n# COMPUTE STATISTICS\ndef compute_statistics(results):\n    sqe_means = dict()\n    sqe_ses = dict()\n    acc_means = dict()\n    acc_ses = dict()\n    for task in tasks:\n\n        sqe_means[task] = dict()\n        sqe_ses[task] = dict()\n        acc_means[task] = dict()\n        acc_ses[task] = dict()\n\n        for method in methods:\n            sqe_means[task][method] = []\n            sqe_ses[task][method] = []\n            acc_means[task][method] = []\n            acc_ses[task][method] = []\n            # hist[task][method] = dict()\n            # hist_ses[task][method] = dict()\n            for num_episodes in episodes:\n                result_name = task + '/' + method + '/' + str(num_episodes)\n                if result_name in results.keys():\n                    result = results[result_name]\n                    # hist[task][method][num_episodes] = np.mean([h[step[task]] for h in result['test_hist']], axis=0)\n                    # hist_ses[task][method][num_episodes] = np.std([h[step[task]] for h in result['test_hist']], axis=0, ddof=1) / np.sqrt(len(result['test_hist']))\n                    sqe_means[task][method].append(np.mean(result['test_mse'], axis=0))\n                    sqe_ses[task][method].append(np.std(result['test_mse'], axis=0, ddof=1) / np.sqrt(len(result['test_mse'])))\n\n                    # hist = np.array([h[step[task]] for h in result['test_hist']])  # result x time x sqe [.0, 0.1]\n                    hist = np.array([[h[i] for i in range(0, 50, 10)] for h in result['test_hist']])  # result x time x sqe [.0, 0.1, .., 10.0]\n                    acc = 1. - np.sum(hist[:,:,:10], axis=-1) # sqe < 1.0\n                    acc_means[task][method].append(np.mean(acc, axis=0))\n                    acc_ses[task][method].append(np.std(acc, axis=0, ddof=1) / np.sqrt(len(acc)))\n                else:\n                    sqe_means[task][method].append([np.nan]*max_steps)\n                    sqe_ses[task][method].append([np.nan]*max_steps)\n                    print(num_episodes)\n                    acc_means[task][method].append([np.nan] * (max_steps // 10))\n                    acc_ses[task][method].append([np.nan] * (max_steps // 10))\n\n            sqe_means[task][method] = np.array(sqe_means[task][method])\n            sqe_ses[task][method] = np.nan_to_num(sqe_ses[task][method])\n            acc_means[task][method] = np.array(acc_means[task][method])\n            acc_ses[task][method] = np.nan_to_num(acc_ses[task][method])\n\n    return sqe_means, sqe_ses, acc_means, acc_ses\n\n\ndef plot_learning_curve(means, ses, step, f=lambda x:x, ylabel_func=lambda x: '', ylim_func=None, show_legend=None, divide_by=None, save_extra=''):\n\n    for task in tasks:\n\n        plt.figure('lc'+ save_extra + ' for ' + task, [4,2.5])\n        # plt.plot([125, 125], [0, 1000], ':', color='gray', linewidth=1)\n        # plt.plot([1000, 1000], [0, 1000], ':', color='gray', linewidth=1)\n\n        for method in methods:\n\n            # valid = np.isnan(means[:,step-1]) == False\n            # eps = np.array(episodes)\n\n            if divide_by is None:\n                plt.fill_between(episodes, (f(means[task][method])-np.array(ses[task][method]))[:,step[task]], (f(means[task][method])+np.array(ses[task][method]))[:,step[task]], color=colors[method], alpha=0.3, linewidth=0.0)\n            else:\n                plt.fill_between(episodes,\n                                 (f(means[task][method])-np.array(ses[task][method]))[:,step[task]] / f(means[task][divide_by][:, step[task]]),\n                                 (f(means[task][method])+np.array(ses[task][method]))[:,step[task]] / f(means[task][divide_by][:, step[task]]),\n                                 color=colors[method], alpha=0.3, linewidth=0.0)\n            # plt.plot(125, means[task][method][episodes.index(125), step[task]], 'o', color=colors[method], markersize=3, linewidth=1)\n            # plt.plot(1000, means[task][method][episodes.index(1000), step[task]], 'x', color=colors[method], markersize=4, linewidth=1)\n\n        for method in methods:\n            if divide_by is None:\n                plt.plot(episodes, f(means[task][method][:, step[task]]), '.-' if method != 'odom' else '--', color=colors[method], label=labels[method], markersize=2, linewidth=1)\n            else:\n                plt.plot(episodes, f(means[task][method][:, step[task]]) / f(means[task][divide_by][:, step[task]]), '.-' if method != 'odom' else '--', color=colors[method], label=labels[method], markersize=2, linewidth=1)\n\n\n\n        plt.gca().set_xscale(\"log\", nonposx='clip')\n        if ylim_func is not None:\n            plt.ylim(ylim_func(task))\n        # plt.ylim([0, max_1])\n        # plt.ylim([0, 1.0])\n        plt.xticks(episodes)\n        plt.gca().get_xaxis().set_major_formatter(matplotlib.ticker.ScalarFormatter())\n        plt.gca().get_xaxis().set_tick_params(which='minor', size=0)\n        plt.gca().get_xaxis().set_tick_params(which='minor', width=0)\n        plt.xlabel('Training episodes (log. scale)')\n        plt.ylabel(ylabel_func(step[task]))\n        # plt.tight_layout()\n        if show_legend is None or show_legend[task]:\n            plt.legend(loc='upper right')\n\n        # plt.figure(task + \" \" + str(step) + \" steps\")\n        extra = 'lc' + save_extra\n        plt.savefig('../plots/' + exp + '/'+exp+'_'+task+'_'+extra+'.pdf', bbox_inches=\"tight\", transparent=True, dpi=600, frameon=True, facecolor='w', pad_inches=0.01)\n        plt.savefig('../plots/' + exp + '/'+exp+'_'+task+'_'+extra+'.png', bbox_inches=\"tight\", transparent=True, dpi=600, frameon=True, facecolor='w', pad_inches=0.01)\n\n\n# PLOT FILTER CONVERGENCE\ndef plot_filter_convergence(means, ses, step, ylabel_func, ylim_func=None, save_extra=''):\n\n    for task in tasks:\n        max_2 = {n: 0 for n in episodes}\n\n        for num_episodes in [1000]:\n            i = episodes.index(num_episodes)\n            plt.figure(task + \" \" + str(num_episodes) + \" training episodes \" + save_extra, [2,2.5])\n            for method in methods:\n                n = means[task][method].shape[1]\n                # if num_episodes == 125:\n                #     plt.plot([step[task]], means[task][method][i, step[task]], 'o', color=colors[method], markersize=3, linewidth=1) # label=labels[method]\n                # elif num_episodes == 1000:\n                #     plt.plot([step[task]], means[task][method][i, step[task]], 'x', color=colors[method], markersize=4, linewidth=1) # label=labels[method]\n                plt.fill_between(np.arange(n), (np.array(means[task][method])-np.array(ses[task][method]))[i,:], (np.array(means[task][method])+np.array(ses[task][method]))[i,:], color=colors[method], alpha=0.3, linewidth=0.0)\n                if method is not 'pf_ind':\n                    max_2[num_episodes] = max(means[task][method][i, -1], max_2[num_episodes])\n            for method in methods:\n                n = means[task][method].shape[1]\n                # plt.plot(np.arange(1, 20+1), means[task][method][i, :20], '--', color=colors[method], markersize=3, linewidth=1)\n                # plt.plot(np.arange(20, max_steps+1), means[task][method][i, 19:], '-', color=colors[method], markersize=3, linewidth=1) # label=labels[method]\n                plt.plot(means[task][method][i, :], '-' if method != 'odom' else '--', color=colors[method], markersize=3, linewidth=1) # label=labels[method]\n\n            # plt.plot([step], [0], 'w', label=' ', linewidth=0)\n            # plt.plot([step[task]], [0], '--', color='gray', label='Steps optimized during training', linewidth=1)\n            # plt.plot([step[task], step[task]], [0, 5*max_2[num_episodes]], ':', color='gray', label='Step tested in learning curve', linewidth=1)\n            if ylim_func is not None:\n                plt.ylim(ylim_func(task))\n            plt.xticks([0, 20, 40])\n            plt.xlabel('Tested at step')\n            plt.ylabel(ylabel_func(step[task]))\n            plt.ylabel('MSE ({} episodes)'.format(num_episodes))\n            plt.tight_layout()\n            # if task == 'nav01':\n            #     plt.legend()\n\n            # plt.figure(task + \" \" + str(num_episodes) + \" training episodes\")\n            extra = 'convrg' + save_extra\n            plt.savefig('../plots/'+exp+'/'+exp+'_'+task+'_steps_'+str(num_episodes)+'_'+extra+'.pdf', bbox_inches=\"tight\", transparent=True, dpi=600, frameon=True, facecolor='w', pad_inches=0.01)\n            plt.savefig('../plots/'+exp+'/'+exp+'_'+task+'_steps_'+str(num_episodes)+'_'+extra+'.png', bbox_inches=\"tight\", transparent=True, dpi=600, frameon=True, facecolor='w', pad_inches=0.01)\n\n\nmethods = ['lstm', 'pf_ind', 'pf_e2e', 'pf_ind_e2e']\nepisodes = [16, 32, 64, 125, 250, 500, 1000]\nexp = 'lc'; tasks = ['nav01', 'nav02', 'nav03']; max_steps = 50\n# exp = 'pl'; tasks = ['nav02']; max_steps = 50; #methods = ['lstm', 'pf_ind_e2e']\n# exp = 'mx'; tasks = ['nav02']; max_steps = 50; #methods = ['lstm', 'pf_ind_e2e']\n\n# exp = 'tr'; tasks = ['nav02']; methods = ['lstm', 'pf_ind', 'pf_e2e', 'pf_ind_e2e', 'odom']; max_steps = 50\n\nplot_path = '../plots/' + exp\nif not os.path.exists(plot_path):\n    os.makedirs(plot_path)\n\nresults = load_results(exp=exp)\nsqe_means, sqe_ses, acc_means, acc_ses = compute_statistics(results)\n\n# print(acc_means['nav01']['lstm'].shape)\n\nif exp == 'lc':\n\n    plot_learning_curve(sqe_means, sqe_ses, step = {'nav01': 20, 'nav02': 20, 'nav03': 30},\n                        ylabel_func=lambda step: 'MSE (at step {})'.format(step),\n                        ylim_func=lambda task: {'nav01':[0,25], 'nav02':[0,55], 'nav03':[0,110]}[task],\n                        save_extra='_mse', show_legend={'nav01': True, 'nav02': False, 'nav03': False})\n\n    plot_learning_curve(acc_means, acc_ses, f=lambda x: x, step = {'nav01': 2, 'nav02': 2, 'nav03': 3},\n                        ylabel_func=lambda step: 'Error rate',\n                        ylim_func=lambda task: [0.0,1.0],\n                        save_extra='_er', show_legend={'nav01': True, 'nav02': False, 'nav03': False})\n\n    plot_learning_curve(sqe_means, sqe_ses, step = {'nav01': 20, 'nav02': 20, 'nav03': 30},\n                        ylabel_func=lambda step: 'MSE relative to LSTM',\n                        ylim_func=lambda task: [0.0, 1.2],\n                        save_extra='_mse_div', show_legend={'nav01': False, 'nav02': False, 'nav03': False}, divide_by='lstm')\n\n    plot_learning_curve(acc_means, acc_ses, f=lambda x: x, step = {'nav01': 2, 'nav02': 2, 'nav03': 3},\n                        ylabel_func=lambda step: 'Error rate relative to LSTM',\n                        ylim_func=lambda task: [0.0, 1.2],\n                        save_extra='_er_div', show_legend={'nav01': False, 'nav02': False, 'nav03': False}, divide_by='lstm')\n\n    # plot_filter_convergence(sqe_means, sqe_ses, step = {'nav01': 40, 'nav02': 40, 'nav03': 40},\n    #                     ylabel_func=lambda step: 'Test MSE (at step {})'.format(step),\n    #                     ylim_func=lambda task: {'nav01':[0,25], 'nav02':[0,55], 'nav03':[0,110]}[task],\n    #                     save_extra='_mse')\n    #\n    # plot_filter_convergence(acc_means, acc_ses, step = {'nav01': 2, 'nav02': 2, 'nav03': 3},\n    #                     ylabel_func=lambda step: 'Test Accuracy (at step {})'.format(step*10),\n    #                     ylim_func=lambda task: [0.0,1.0],\n    #                     save_extra='_acc')\n\n    # plot_filter_convergence(sqe_means, sqe_ses, step = {'nav01': 40, 'nav02': 40, 'nav03': 40},\n    #                     ylabel_func=lambda step: 'Test MSE (at step {})'.format(step),\n    #                     ylim_func=lambda task: {'nav01':[0,25], 'nav02':[0,1.0], 'nav03':[0,110]}[task],\n    #                     save_extra='_mse')\n\nelif exp == 'tr':\n\n    plot_learning_curve(sqe_means, sqe_ses, step = {'nav01': 20, 'nav02': 20, 'nav03': 30},\n                        ylabel_func=lambda step: 'MSE (at step {})'.format(step),\n                        ylim_func=lambda task: {'nav01':[0,25], 'nav02':[0,55], 'nav03':[0,110]}[task],\n                        save_extra='_mse', show_legend={'nav01': True, 'nav02': True, 'nav03': False})\n\n    plot_learning_curve(acc_means, acc_ses, f=lambda x: x, step = {'nav01': 2, 'nav02': 2, 'nav03': 3},\n                        ylabel_func=lambda step: 'Error rate',\n                        ylim_func=lambda task: [0.0,1.0],\n                        save_extra='_er', show_legend={'nav01': True, 'nav02': True, 'nav03': False})\n    #\n    # plot_learning_curve(sqe_means, sqe_ses, step = {'nav01': 20, 'nav02': 20, 'nav03': 30},\n    #                     ylabel_func=lambda step: 'MSE relative to LSTM',\n    #                     ylim_func=lambda task: [0.0, 1.2],\n    #                     save_extra='_mse_div', show_legend={'nav01': False, 'nav02': False, 'nav03': False}, divide_by='lstm')\n    #\n    # plot_learning_curve(acc_means, acc_ses, f=lambda x: x, step = {'nav01': 2, 'nav02': 2, 'nav03': 3},\n    #                     ylabel_func=lambda step: 'Error rate relative to LSTM',\n    #                     ylim_func=lambda task: [0.0, 1.2],\n    #                     save_extra='_er_div', show_legend={'nav01': False, 'nav02': False, 'nav03': False}, divide_by='lstm')\n\n\n    plot_filter_convergence(sqe_means, sqe_ses, step = {'nav01': 40, 'nav02': 40, 'nav03': 40},\n                        ylabel_func=lambda step: 'MSE (at step {})'.format(step),\n                        ylim_func=lambda task: {'nav01':[0,25], 'nav02':[0,1.0], 'nav03':[0,110]}[task],\n                        save_extra='_mse')\n    #\n    # plot_filter_convergence(acc_means, acc_ses, step = {'nav01': 2, 'nav02': 2, 'nav03': 3},\n    #                     ylabel_func=lambda step: 'Test Accuracy (at step {})'.format(step*10),\n    #                     ylim_func=lambda task: [0.0,1.0],\n    #                     save_extra='_acc')\n\nelse:\n\n    plot_learning_curve(sqe_means, sqe_ses, step = {'nav01': 20, 'nav02': 20, 'nav03': 30},\n                        ylabel_func=lambda step: 'MSE (at step {})'.format(step),\n                        ylim_func=lambda task: {'nav01':[0,25], 'nav02':[0,55], 'nav03':[0,110]}[task],\n                        save_extra='_mse', show_legend={'nav01': True, 'nav02': False, 'nav03': False})\n\n    plot_learning_curve(acc_means, acc_ses, f=lambda x: x, step = {'nav01': 2, 'nav02': 2, 'nav03': 3},\n                        ylabel_func=lambda step: 'Error rate',\n                        ylim_func=lambda task: [0.0,1.0],\n                        save_extra='_er', show_legend={'nav01': True, 'nav02': False, 'nav03': False})\n\n    plot_learning_curve(sqe_means, sqe_ses, step = {'nav01': 20, 'nav02': 20, 'nav03': 30},\n                        ylabel_func=lambda step: 'MSE relative to LSTM',\n                        ylim_func=lambda task: [0.0, 1.2],\n                        save_extra='_mse_div', show_legend={'nav01': False, 'nav02': False, 'nav03': False}, divide_by='lstm')\n\n    plot_learning_curve(acc_means, acc_ses, f=lambda x: x, step = {'nav01': 2, 'nav02': 2, 'nav03': 3},\n                        ylabel_func=lambda step: 'Error rate relative to LSTM',\n                        ylim_func=lambda task: [0.0, 1.2],\n                        save_extra='_er_div', show_legend={'nav01': False, 'nav02': False, 'nav03': False}, divide_by='lstm')\n\n    # plot_filter_convergence(sqe_means, sqe_ses, step = {'nav01': 40, 'nav02': 40, 'nav03': 40},\n    #                     ylabel_func=lambda step: 'Test MSE (at step {})'.format(step),\n    #                     ylim_func=lambda task: {'nav01':[0,25], 'nav02':[0,55], 'nav03':[0,110]}[task],\n    #                     save_extra='_mse')\n    #\n    # plot_filter_convergence(acc_means, acc_ses, step = {'nav01': 2, 'nav02': 2, 'nav03': 3},\n    #                     ylabel_func=lambda step: 'Test Accuracy (at step {})'.format(step*10),\n    #                     ylim_func=lambda task: [0.0,1.0],\n    #                     save_extra='_acc')\n\n    # plot_filter_convergence(sqe_means, sqe_ses, step = {'nav01': 40, 'nav02': 40, 'nav03': 40},\n    #                     ylabel_func=lambda step: 'Test MSE (at step {})'.format(step),\n    #                     ylim_func=lambda task: {'nav01':[0,25], 'nav02':[0,1.0], 'nav03':[0,110]}[task],\n    #                     save_extra='_mse')\n\nplt.show()\n\n\n"
  },
  {
    "path": "plotting/nt_plot.py",
    "content": "import pickle\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport os\n\nresults = None\n\n\n\n# matplotlib.rcParams.update({'font.size': 12})\n\ncolor_list = plt.cm.tab10(np.linspace(0, 1, 10))\ncolors = {'lstm': color_list[0], 'pf_e2e': color_list[1], 'pf_ind_e2e': color_list[2], 'pf_ind': color_list[3]}\nlabels = {'lstm': 'LSTM', 'pf_e2e': 'DPF (e2e)', 'pf_ind_e2e': 'DPF (ind+e2e)', 'pf_ind': 'DPF (ind)'}\ntasks = ['nav02']\nmethods = ['lstm', 'pf_ind', 'pf_e2e', 'pf_ind_e2e']\n\n# load results\nresults = dict()\n\ncount = 0\nfor task in tasks:\n    log_path = '../log/nt'\n    for filename in [f for f in os.listdir(log_path) if os.path.isfile(os.path.join(log_path, f))]:\n        full_filename = os.path.join(log_path, filename)\n        print('loading {}:'.format(count) + full_filename + ' ...')\n        try:\n            # if 'DeepThought' not in filename:\n            # if 'DeepThought' in filename:\n            with open(full_filename, 'rb') as f:\n                result = pickle.load(f)\n                result_name = result['exp_params'][0]['task'] + '/' + result['exp_params'][0]['method'] + '/' + str(result['exp_params'][0]['num_episodes']) + '/' + result['exp_params'][0]['noise_condition']\n                print(result_name)\n                if result_name not in results.keys():\n                    results[result_name] = result\n                else:\n                    for key in result.keys():\n                        if key in results[result_name].keys():\n                            results[result_name][key] += result[key]\n                        else:\n                            results[result_name][key] = result[key]\n                        # print(result_name, key)\n                count += 1\n        except Exception as e:\n            print('%r' % e)\n            # raise e\n            print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')\n\nfor result_name in results.keys():\n    print(result_name, len(results[result_name]['exp_params']))\n\nprint('Loaded {} results'.format(count))\n\n# print(results['test_errors'].shape, np.mean(results['test_errors']**2, axis=1))\n\n#print('SHAPE', results['test_mse'].shape)\n\n# plt.figure(1)\n# plt.gca().set_color_cycle(None)\n# for method in set(results['method']):\n\ntask = 'nav02'\nmetric = 'err'; step = 2 # err or mse\n# metric = 'mse'; step = 20 # err or mse\n\n# episodes = [125, 1000]\nepisodes = [1000]\nfig_names = []\n\nfor noise_in in ['odom', 'image']:\n\n    if noise_in == 'odom':\n        conditions = ['odom0_imgTG', 'odom5_imgTG', 'odom10_imgTG', 'odom20_imgTG', 'odomX_imgTG']; clabels = {'odom0_imgTG': '0', 'odom5_imgTG': '5', 'odom10_imgTG': '10', 'odom20_imgTG': '20', 'odomX_imgTG': 'X'}\n        # conditions = ['odom0_imgTG', 'odom5_imgTG', 'odom10_imgTG', 'odomX_imgTG']; clabels = {'odom0_imgTG': '0', 'odom5_imgTG': '5', 'odom10_imgTG': '10', 'odomX_imgTG': 'X'}\n    else:\n        conditions = ['odom10_imgC', 'odom10_imgG', 'odom10_imgT', 'odom10_imgTG', 'odom10_imgX']; clabels = {'odom10_imgC': 'N', 'odom10_imgG': 'G', 'odom10_imgT': 'S', 'odom10_imgTG': 'G+S', 'odom10_imgX': 'X'}\n\n    max_1 = 0\n    max_2 = {n: 0 for n in episodes}\n\n    means = dict()\n    ses = dict()\n\n    for num_episodes in episodes:\n\n        means[num_episodes] = dict()\n        ses[num_episodes] = dict()\n\n        for method in methods:\n\n            if metric == 'mse':\n                means[num_episodes][method] = np.zeros([len(conditions), len(conditions), 50])\n                ses[num_episodes][method] = np.zeros([len(conditions), len(conditions), 50])\n            elif metric == 'err':\n                means[num_episodes][method] = np.zeros([len(conditions), len(conditions), 5])\n                ses[num_episodes][method] = np.zeros([len(conditions), len(conditions), 5])\n\n            for c, condition in enumerate(conditions):\n\n                result_name = task + '/' + method + '/' + str(num_episodes) + '/' + condition\n                if result_name in results.keys():\n                    result = results[result_name]\n                    for ct, test_condition in enumerate(conditions):\n                        if 'test_'+test_condition+'_mse' not in result.keys():\n                            means[num_episodes][method][c, ct] *= np.nan\n                            ses[num_episodes][method][c, ct] *= np.nan\n                        else:\n                            if noise_in != \"odom\" and num_episodes == 1000 and c == 1 and ct == 1:\n                                print(method,\n                                    np.array(result['test_'+test_condition+'_mse'])[:, 30])\n\n                            if metric == 'mse':\n                                means[num_episodes][method][c, ct] = np.mean(result['test_'+test_condition+'_mse'], axis=0)\n                                std = np.std(result['test_'+test_condition+'_mse'], axis=0, ddof=1)\n                                ses[num_episodes][method][c, ct] = std / np.sqrt(len(result['test_'+test_condition+'_mse']))\n                            elif metric == 'err':\n                                hist = np.array([[h[i] for i in range(0, 50, 10)] for h in result['test_'+test_condition+'_hist']])  # result x time x sqe [.0, 0.1, .., 10.0]\n                                err = 1. - np.sum(hist[:,:,:10], axis=-1) # sqe < 1.0\n                                means[num_episodes][method][c, ct] = np.mean(err, axis=0)\n                                ses[num_episodes][method][c, ct] = np.std(err, axis=0, ddof=1) / np.sqrt(len(err))\n\n                else:\n                    # print(result_name, 0)\n                    for test_condition in conditions:\n                        means[num_episodes][method][c, :] *= np.nan\n\n\n        # if noise_in != 'odom':\n        #     print(means[1000]['pf_ind'][1,1,30])\n\n        means[num_episodes]['min'] = np.stack([means[num_episodes][method] for method in methods], axis=0).min(axis=0)\n\n        for method in methods:\n\n            fig_name = 'nt_{}_{}_{}'.format(noise_in,num_episodes,method)\n            plt.figure(fig_name, [3,2.5])\n            fig_names.append(fig_name)\n\n            min_val, max_val, diff = 0., len(conditions), 1.\n\n            #imshow portion\n            N_points = (max_val - min_val) / diff\n            m = means[num_episodes][method][:,:,step].T\n            is_min = m == means[num_episodes]['min'][:,:,step].T\n\n            # plt.imshow((means[:,:,30-1])**0.5, interpolation='nearest', vmin=0, vmax=15)\n            if metric == 'mse':\n                plt.imshow(np.log(m), interpolation='nearest', vmin=-3, vmax=6, cmap='viridis')\n            elif metric == 'err':\n                plt.imshow(m, interpolation='nearest', vmin=-0.33, vmax=1.0, cmap='viridis_r')\n            # data = np.reshape(np.arange(len(conditions)*len(conditions)), [len(conditions), len(conditions)])\n            # plt.imshow(data, interpolation='nearest', vmin=0, vmax=10)\n            plt.xticks(np.arange(len(conditions)), [clabels[c] for c in conditions])\n            plt.yticks(np.arange(len(conditions)), [clabels[c] for c in conditions])\n\n            plt.ylabel('Test {}{} noise'.format(noise_in, '.' if noise_in == 'odom' else ''))\n            plt.xlabel('Training {}{} noise'.format(noise_in, '.' if noise_in == 'odom' else ''))\n\n            # plt.colorbar()\n            #text portion\n            ind_array = np.arange(min_val, max_val, diff)\n            x, y = np.meshgrid(ind_array, ind_array)\n\n            for x_val, y_val in zip(x.flatten(), y.flatten()):\n                value = m[int(y_val),int(x_val)]\n                if metric == 'err':\n                    text = '{:.4s}'.format('{:.3f}'.format(value)[1:])\n                else:\n                    text = '{:.4s}'.format('{:.2f}'.format(value))\n                if x_val == y_val:\n                    if x_val == 3 and noise_in == 'image' or x_val == 2 and noise_in == 'odom':\n                        style = '-w'\n                    else:\n                        style = '--w'\n                    plt.plot(np.array([x_val, x_val+1, x_val+1, x_val, x_val])-0.5, np.array([y_val, y_val, y_val+1, y_val+1, y_val])-0.5, style, linewidth=1.5)\n                # if value > 0.9:\n                #     color = 'black'\n                # else:\n                #     color = 'white'\n                plt.text(x_val, y_val, text, va='center', ha='center', color='white', fontweight='bold' if is_min[int(y_val), int(x_val)] else 'normal')\n\n            fig_name = 'nt_diag_{}_{}'.format(noise_in,num_episodes)\n            # plt.figure(fig_name, [3,2.5])\n            plt.figure(fig_name, [2,2.5])\n            fig_names.append(fig_name)\n\n            if noise_in == 'odom':\n                x = np.array([int(clabels[c]) for c in conditions[:-1]])\n            else:\n                x = np.arange(len(conditions)-1)\n            m = means[num_episodes][method][:,:,step]\n            s = ses[num_episodes][method][:,:,step]\n            plt.plot(x, np.diag(m)[:-1], '.-', color=colors[method], label=labels[method], markersize=2)\n            ind = 2 if noise_in == 'odom' else 3\n            plt.plot(x[ind], np.diag(m)[ind], 'x', color=colors[method], markersize=4)\n            plt.fill_between(x, np.diag(m-s)[:-1], np.diag(m+s)[:-1], color=colors[method], alpha=0.5, linewidth=0.0)\n\n            plt.xticks(x, [clabels[c] for c in conditions[:-1]])\n            if noise_in == 'odom':\n                plt.xlabel('Odometry noise (%)')\n            else:\n                plt.xlabel('Image noise')\n                plt.legend()\n\n            if metric == 'mse':\n                plt.ylabel('MSE ({} episodes)'.format(num_episodes))\n            else:\n                plt.ylabel('Error rate ({} episodes)'.format(num_episodes))\n                plt.ylim([0, 0.3])\n                plt.yticks([0.0, 0.1, 0.2, 0.3])\n\n            fig_name = 'nt_shuffle_{}'.format(num_episodes)\n            # plt.figure(fig_name, [3,2.5])\n            plt.figure(fig_name, [2,2.5])\n            fig_names.append(fig_name)\n\n            if noise_in == 'odom':\n                plt.bar(0.0  - 0.5 + (methods.index(method)+1)/len(methods)*0.8,\n                        np.diag(m)[-3],\n                        0.8/len(methods),\n                        yerr=np.diag(s)[-3],\n                        color=colors[method], label=labels[method])\n            plt.bar((1.0 if noise_in == 'odom' else 2.0)  - 0.5 + (methods.index(method)+1)/len(methods)*0.8,\n                    np.diag(m)[-1], 0.8/len(methods),\n                    yerr=np.diag(s)[-1],\n                    color=colors[method])\n            if noise_in == 'odom':\n                relative = np.diag(m)[-1] / np.diag(m)[-3]\n            else:\n                relative = np.diag(m)[-1] / np.diag(m)[-2]\n            textpos = np.diag(m)[-1] + np.diag(s)[-1] + 2\n            if num_episodes == 1000:\n                if textpos > 10:\n                    textpos = 3\n            elif textpos > 80:\n                    textpos = 10\n            color = 'black'\n            if metric == 'err':\n                textpos = 0.05\n                color = 'white'\n            # plt.text((1.0 if noise_in == 'odom' else 2.0)  - 0.5 + (methods.index(method)+1)/len(methods)*0.8,\n            #          textpos, '×{:.0f}'.format(relative), va='bottom', ha='center',color=color, rotation=90)\n            plt.xticks([0, 1, 2], ['Both', 'Image', 'Odom.'])\n            plt.ylim([0,1.05])\n            plt.xlabel('Input')\n            if metric == 'mse':\n                plt.ylabel('MSE ({} episodes)'.format(num_episodes))\n            else:\n                plt.ylabel('Error rate ({} episodes)'.format(num_episodes))\n            # plt.legend()\n\n\nfor fn in fig_names:\n    plt.figure(fn)\n    # plt.tight_layout()\n    plt.savefig('../plots/nt/{}.pdf'.format(fn), bbox_inches=\"tight\", transparent=True, dpi=600, frameon=True, facecolor='w', pad_inches=0.01)\n    # plt.savefig('../plots/nt/{}.png'.format(fn), bbox_inches=\"tight\", transparent=True, dpi=600, frameon=True, facecolor='w', pad_inches=0.01)\n\n    # plt.bar(np.arange(len(conditions)) + (methods.index(method)+1)/len(methods)*0.8 - 0.5, np.diag(means[:,:,30-1])/np.diag(means[:,:,30-1])[0], 0.8/len(methods), color=colors[method])\n\n\nplt.show()\n"
  },
  {
    "path": "plotting/plot_models.py",
    "content": "import tensorflow as tf\nimport pickle\nfrom methods.dpf import DPF\nfrom methods.rnn import RNN\nfrom utils.data_utils import load_data, noisyfy_data, make_batch_iterator, reduce_data\nfrom utils.exp_utils import get_default_hyperparams\nfrom utils.method_utils import compute_sq_distance\nfrom utils.plotting_utils import plot_maze, plot_observations\nfrom methods.odom import OdometryBaseline\nimport numpy as np\nimport matplotlib.pyplot as plt\n# from mpl_toolkits.axes_grid1 import make_axes_locatable\n\nhead_scale = 1.5\nquiv_kwargs = {'scale_units':'xy', 'scale':1./40., 'width': 0.003, 'headlength': 5*head_scale, 'headwidth': 3*head_scale, 'headaxislength': 4.5*head_scale}\nmarker_kwargs = {'markersize': 4.5, 'markerfacecolor':'None', 'markeredgewidth':0.5}\n\ndef plot_measurement_model(session, method, statistics, batch, task, num_examples, variant):\n\n    batch_size = len(batch['o'])\n\n    x = np.linspace(100.0 / 4, 1000.0 - 100.0 / 4, 20)\n    y = np.linspace(100.0 / 4, 500.0 - 100.0 / 4, 10)\n    theta = np.linspace(-np.pi, np.pi, 12 + 1)[1:]\n    g = np.meshgrid(x, y, theta)\n\n    poses = np.vstack([np.ravel(x) for x in g]).transpose([1, 0])\n    test_poses = tf.tile(tf.constant(poses, dtype='float32')[None, :, :], [batch_size, 1, 1])\n    measurement_model_out = method.measurement_update(method.encodings[0, :], test_poses, statistics['means'],\n                                                      statistics['stds'])\n\n    # define the inputs and train/run the model\n    input_dict = {**{method.placeholders[key]: batch[key] for key in 'osa'},\n                  }\n\n    obs_likelihood = session.run(measurement_model_out, input_dict)\n    print(obs_likelihood.shape)\n\n    for i in range(num_examples):\n        # plt.figure(\"%s likelihood\" % i)\n        fig, (ax, cax) = plt.subplots(1, 2, figsize=(2.4 / 0.83 / 0.95 / 0.97, 1.29 / 0.9),\n                                      gridspec_kw={\"width_ratios\": [0.97, 0.03]}, num=\"%s %s likelihood\" % (variant, i))\n        # plt.gca().clear()\n        plot_maze(task, margin=5, linewidth=0.5, ax=ax)\n\n        idx = obs_likelihood[i,:] > 1*np.mean(obs_likelihood[i,:])\n        # idx = obs_likelihood[i, :] > 0 * np.mean(obs_likelihood[i, :])\n        max = np.max(obs_likelihood[i, :])\n\n        # ax.scatter([poses[:, 0]], [poses[:, 1]], s=[0.001], c=[(0.8, 0.8, 0.8)], marker='.')\n\n        quiv = ax.quiver(poses[idx, 0] + 0 * np.cos(poses[idx, 2]), poses[idx, 1] + 0* np.sin(poses[idx, 2]), np.cos(poses[idx, 2]),\n                         np.sin(poses[idx, 2]), obs_likelihood[i, idx],\n                         cmap='viridis_r',\n                         clim=[0.0, max],\n                         **quiv_kwargs\n                         )\n\n        ax.plot([batch['s'][0, i, 0]], [batch['s'][0, i, 1]], 'or', **marker_kwargs)\n\n        ax.quiver([batch['s'][0, i, 0]], [batch['s'][0, i, 1]], np.cos([batch['s'][0, i, 2]]),\n                  np.sin([batch['s'][0, i, 2]]), color='red',\n                  **quiv_kwargs\n                  )\n        ax.axis('off')\n        fig.colorbar(quiv, cax=cax, orientation=\"vertical\", label='Obs. likelihood', ticks=[0.0, 0.2, 0.4, 0.6, 0.8, 1.0])\n        plt.subplots_adjust(left=0.0, bottom=0.05, right=0.83, top=0.95, wspace=0.05, hspace=0.00)\n        plt.savefig('../plots/models/measurement_model{}.pdf'.format(i), transparent=True, dpi=600, frameon=False, facecolor='w', pad_inches=0.01)\n\n\ndef plot_proposer(session, method, statistics, batch, task, num_examples, variant):\n\n    num_particles = 1000\n    proposer_out = method.propose_particles(method.encodings[0, :], num_particles, statistics['state_mins'], statistics['state_maxs'])\n\n    # define the inputs and train/run the model\n    input_dict = {**{method.placeholders[key]: batch[key] for key in 'osa'},\n                  }\n    particles = session.run(proposer_out, input_dict)\n\n    for i in range(num_examples):\n        fig = plt.figure(figsize=(2.4, 1.29/0.9), num=\"%s %s proposer\" % (variant, i))\n        # plt.gca().clear()\n        plot_maze(task, margin=5, linewidth=0.5)\n\n        quiv = plt.quiver(particles[i, :, 0], particles[i, :, 1], np.cos(particles[i, :, 2]),\n                         np.sin(particles[i, :, 2]), np.ones([num_particles]), cmap='viridis_r', clim=[0, 2], alpha=1.0,\n                          **quiv_kwargs\n                          )\n\n        plt.quiver([batch['s'][0, i, 0]], [batch['s'][0, i, 1]], np.cos([batch['s'][0,i, 2]]),\n                  np.sin([batch['s'][0, i, 2]]), color='red',\n                **quiv_kwargs)  # width=0.01, scale=100\n        plt.plot([batch['s'][0, i, 0]], [batch['s'][0, i, 1]], 'or', **marker_kwargs)\n\n\n        plt.gca().axis('off')\n        plt.subplots_adjust(left=0.0, bottom=0.05, right=1.0, top=0.95, wspace=0.0, hspace=0.00)\n\n        # plt.subplots_adjust(left=0.0, bottom=0.0, right=1.0, top=1.0, wspace=0.001, hspace=0.1)\n        plt.savefig('../plots/models/prop{}.pdf'.format(i), transparent=True, dpi=600, frameon=False, facecolor='w', pad_inches=0.01)\n\n\ndef plot_motion_model(session, method, statistics, batch, task, num_examples, num_particles, variant):\n\n    motion_samples = method.motion_update(method.placeholders['a'][:, 1],\n                                        tf.tile(method.placeholders['s'][:, :1], [1, num_particles, 1]),\n                                        statistics['means'], statistics['stds'], statistics['state_step_sizes'])\n\n    # define the inputs and train/run the model\n    input_dict = {**{method.placeholders[key]: batch[key] for key in 'osa'},\n                  }\n    particles = session.run(motion_samples, input_dict)\n\n    fig = plt.figure(figsize=(2.4, 1.29), num=\"%s motion model\" % (variant))\n    # plt.gca().clear()\n    plot_maze(task, margin=5, linewidth=0.5)\n\n    for i in range(num_examples):\n\n        plt.quiver(particles[i, :, 0], particles[i, :, 1], np.cos(particles[i, :, 2]),\n                          np.sin(particles[i, :, 2]), np.ones([num_particles]), cmap='viridis_r',\n                   **quiv_kwargs,\n                   alpha=1.0, clim=[0, 2])  # width=0.01, scale=100\n\n        plt.quiver([batch['s'][i, 0, 0]], [batch['s'][i, 0, 1]], np.cos([batch['s'][i, 0, 2]]),\n                   np.sin([batch['s'][i, 0, 2]]), color='black',\n                   **quiv_kwargs,\n                   )  # width=0.01, scale=100\n\n        plt.plot(batch['s'][i, :2, 0], batch['s'][i, :2, 1], '--', color='black', linewidth=0.3)\n        plt.plot(batch['s'][i, :1, 0], batch['s'][i, :1, 1], 'o', color='black', linewidth=0.3, **marker_kwargs)\n        plt.plot(batch['s'][i, 1:2, 0], batch['s'][i, 1:2, 1], 'o', color='red', linewidth=0.3, **marker_kwargs)\n\n        plt.quiver([batch['s'][i, 1, 0]], [batch['s'][i, 1, 1]], np.cos([batch['s'][i, 1, 2]]),\n                   np.sin([batch['s'][i, 1, 2]]), color='red',\n                   **quiv_kwargs)  # width=0.01, scale=100\n\n    plt.gca().axis('off')\n\n    plt.subplots_adjust(left=0.0, bottom=0.0, right=1.0, top=1.0, wspace=0.001, hspace=0.1)\n    plt.savefig('../plots/models/motion_model{}.pdf'.format(i), transparent=True, dpi=600, frameon=False, facecolor='w', pad_inches=0.01)\n\n\ndef plot_particle_filter(session, method, statistics, batch, task, num_examples, num_particles, variant):\n    color_list = plt.cm.tab10(np.linspace(0, 1, 10))\n    colors = {'lstm': color_list[0], 'pf_e2e': color_list[1], 'pf_ind_e2e': color_list[2], 'pf_ind': color_list[3],\n              'ff': color_list[4], 'odom': color_list[4]}\n\n    pred, s_particle_list, s_particle_probs_list = method.predict(session, batch, num_particles, return_particles=True)\n\n    num_steps = 20 # s_particle_list.shape[1]\n\n    for s in range(num_examples):\n\n        plt.figure(\"example {}, vartiant: {}\".format(s, variant), figsize=[12, 5.15])\n        plt.gca().clear()\n\n        for i in range(num_steps):\n            ax = plt.subplot(4, 5, i + 1, frameon=False)\n            plt.gca().clear()\n\n            plot_maze(task, margin=5, linewidth=0.5)\n\n            if i < num_steps - 1:\n                ax.quiver(s_particle_list[s, i, :, 0], s_particle_list[s, i, :, 1],\n                           np.cos(s_particle_list[s, i, :, 2]), np.sin(s_particle_list[s, i, :, 2]),\n                           s_particle_probs_list[s, i, :], cmap='viridis_r', clim=[.0, 2.0/num_particles], alpha=1.0,\n                          **quiv_kwargs\n                          )\n\n                current_state = batch['s'][s, i, :]\n                plt.quiver(current_state[0], current_state[1], np.cos(current_state[2]),\n                           np.sin(current_state[2]), color=\"red\", **quiv_kwargs)\n\n                plt.plot(current_state[0], current_state[1], 'or', **marker_kwargs)\n            else:\n\n                ax.plot(batch['s'][s, :num_steps, 0], batch['s'][s, :num_steps, 1], '-', linewidth=0.6, color='red')\n                ax.plot(pred[s, :num_steps, 0], pred[s, :num_steps, 1], '-', linewidth=0.6,\n                        color=colors['pf_ind_e2e'])\n\n                ax.plot(batch['s'][s, :1, 0], batch['s'][s, :1, 1], '.', linewidth=0.6, color='red', markersize=3)\n                ax.plot(pred[s, :1, 0], pred[s, :1, 1], '.', linewidth=0.6, markersize=3,\n                        color=colors['pf_ind_e2e'])\n\n\n            plt.subplots_adjust(left=0.0, bottom=0.0, right=1.0, top=1.0, wspace=0.001, hspace=0.1)\n            plt.gca().set_aspect('equal')\n            plt.xticks([])\n            plt.yticks([])\n\n        plt.savefig('../plots/models/pf{}.pdf'.format(s), transparent=True, dpi=600, frameon=False, facecolor='w', pad_inches=0.01)\n\n        plt.figure('colorbar', (12, 0.6))\n        a = np.array([[0, 2.0/num_particles]])\n        img = plt.imshow(a, cmap=\"viridis_r\")\n        plt.gca().set_visible(False)\n        cax = plt.axes([0.25, 0.75, 0.50, 0.2])\n        plt.colorbar(orientation=\"horizontal\", cax=cax, label='Particle weight', ticks=[0, 0.001, 0.002])\n\n        plt.savefig('../plots/models/colorbar.pdf'.format(s), transparent=True, dpi=600, frameon=False, facecolor='w', pad_inches=0.01)\n\ndef plot_prediction(pred1, pred2, statistics, batch, task, num_examples, variant):\n    color_list = plt.cm.tab10(np.linspace(0, 1, 10))\n    colors = {'lstm': color_list[0], 'pf_e2e': color_list[1], 'pf_ind_e2e': color_list[2], 'pf_ind': color_list[3],\n                'ff': color_list[4], 'odom': color_list[4]}\n\n    num_steps = 50\n    init_steps = 20\n\n    for s in range(num_examples):\n\n        fig = plt.figure(figsize=(2.4, 1.29), num=\"%s prediction %s\" % (variant, s))\n\n        # plt.figure(\"example {}, vartiant: {}\".format(s, variant), figsize=[12, 5.15])\n        plt.gca().clear()\n        plot_maze(task, margin=5, linewidth=0.5)\n\n        plt.plot(batch['s'][s, :num_steps, 0], batch['s'][s, :num_steps, 1], '-', linewidth=0.3, color='gray')\n        plt.plot(pred1[s, :init_steps, 0], pred1[s, :init_steps, 1], '--', linewidth=0.3, color=colors['pf_ind_e2e'])\n        plt.plot(pred1[s, init_steps-1:num_steps, 0], pred1[s, init_steps-1:num_steps, 1], '-', linewidth=0.3, color=colors['pf_ind_e2e'])\n        plt.plot(pred2[s, :init_steps, 0], pred2[s, :init_steps, 1], '--', linewidth=0.3, color=colors['lstm'])\n        plt.plot(pred2[s, init_steps-1:num_steps, 0], pred2[s, init_steps-1:num_steps, 1], '-', color=colors['lstm'], linewidth=0.3)\n\n        # for i in range(init_steps, num_steps):\n        #\n        #     p = pred1[s, i, :]\n        #     plt.quiver(p[0], p[1], np.cos(p[2]),\n        #                np.sin(p[2]), color=colors['pf_ind_e2e'], **quiv_kwargs)\n        #     p = pred2[s, i, :]\n        #     plt.quiver(p[0], p[1], np.cos(p[2]),\n        #                np.sin(p[2]), color=colors['lstm'], **quiv_kwargs)\n        #     # plt.plot(p[0], p[1], 'og', **marker_kwargs)\n        #\n        #     current_state = batch['s'][s, i, :]\n        #     plt.quiver(current_state[0], current_state[1], np.cos(current_state[2]),\n        #                np.sin(current_state[2]), color=\"black\", **quiv_kwargs)\n        #     # plt.plot(current_state[0], current_state[1], 'or', **marker_kwargs)\n\n        plt.gca().set_aspect('equal')\n        plt.xticks([])\n        plt.yticks([])\n        plt.subplots_adjust(left=0.0, bottom=0.0, right=1.0, top=1.0, wspace=0.001, hspace=0.1)\n        plt.savefig('../plots/models/pred{}.pdf'.format(s), transparent=True, dpi=600, frameon=False, facecolor='w', pad_inches=0.01)\n\n\ndef plot_observation(batch, i, t=0):\n\n    plt.figure('%r obs' % i, (2, 2))\n    plt.imshow(np.clip(batch['o'][i, t, :, :, :] / 255.0, 0.0, 1.0), interpolation='nearest')\n    plt.axis('off')\n    # plt.subplots_adjust(left=0.0, bottom=0.0, right=1.0, top=1.0, wspace=0.001, hspace=0.001)\n    plt.subplots_adjust(left=0.0, bottom=0.15, right=1.0, top=0.85, wspace=0.0, hspace=0.00)\n    plt.savefig('../plots/models/obs{}.png'.format(i+t), transparent=True, dpi=600, frameon=False, facecolor='w', pad_inches=0.01)\n\n\ndef plot_measurement_statistics(session, method, statistics, batch_iterator, batch_size, variant):\n\n    color_list = plt.cm.tab10(np.linspace(0, 1, 10))\n    colors = {'lstm': color_list[0], 'e2e': color_list[1], 'ind_e2e': color_list[2], 'ind': color_list[3], 'ff': color_list[4], 'odom': color_list[4]}\n    labels = {'e2e': 'DPF(e2e)', 'ind_e2e': 'DPF(ind+e2e)', 'ind': 'DPF(ind)'}\n\n    x = np.linspace(100.0 / 4, 1500.0 - 100.0 / 4, 30)\n    y = np.linspace(100.0 / 4, 900.0 - 100.0 / 4, 18)\n    theta = np.linspace(-np.pi, np.pi, 12 + 1)[1:]\n    g = np.meshgrid(x, y, theta)\n\n    poses = np.vstack([np.ravel(x) for x in g]).transpose([1, 0])\n    test_poses = tf.tile(tf.constant(poses, dtype='float32')[None, :, :], [batch_size, 1, 1])\n    measurement_model_out = method.measurement_update(method.encodings[:, 0], test_poses, statistics['means'],\n                                                      statistics['stds'])\n    true_measurement_model_out = method.measurement_update(method.encodings[:, 0], method.placeholders['s'][:, 0, None, :], statistics['means'],\n                                                      statistics['stds'])\n\n    hist = 0.0\n    true_hist = 0.0\n\n    for i in range(1000000): # 1000000\n        # define the inputs and train/run the model\n        batch = next(batch_iterator)\n        input_dict = {**{method.placeholders[key]: batch[key] for key in 'osa'},\n                      }\n        if i < 100:\n            obs_likelihood, true_obs_likelihood = session.run([measurement_model_out, true_measurement_model_out], input_dict)\n            h, bins = np.histogram(obs_likelihood, 50, [0,1])\n            hist += h\n        else:\n            true_obs_likelihood = session.run(true_measurement_model_out, input_dict)\n        h, true_bins = np.histogram(true_obs_likelihood, 20, [0,1])\n        true_hist += h\n\n    true_hist = true_hist / np.sum(true_hist) * len(true_hist)\n    hist = hist / np.sum(hist) * len(hist)\n    plt.figure('Observation likelihood statistics', [3.3,2.5])\n    plt.plot(bins[1:] - (bins[1]-bins[0])/2, hist, '--', color=colors[variant])\n    plt.plot(true_bins[1:] - (true_bins[1]-true_bins[0])/2, true_hist, '-', color=colors[variant], label=labels[variant])\n    plt.legend(loc='upper center')\n    plt.yticks([0, 1, 2, 3])\n    plt.ylim([0,3])\n    plt.xlabel('Estimated observation likelihood')\n    plt.ylabel('Density')\n    # plt.gca().set_yscale(\"log\", nonposx='clip')\n    plt.tight_layout()\n    plt.savefig('../plots/models/measurement_statistics.pdf', transparent=True, dpi=600, frameon=False, facecolor='w', pad_inches=0.01)\n\n\n\ndef plot_motion_statistics(session, method, statistics, batch_iterator, task, variant):\n\n    color_list = plt.cm.tab10(np.linspace(0, 1, 10))\n    colors = {'lstm': color_list[0], 'e2e': color_list[1], 'ind_e2e': color_list[2], 'ind': color_list[3], 'ff': color_list[4], 'odom': color_list[4]}\n    labels = {'e2e': 'DPF(e2e)', 'ind_e2e': 'DPF(ind+e2e)', 'ind': 'DPF(ind)'}\n\n    num_particles = 100\n\n    motion_samples = method.motion_update(method.placeholders['a'][:, 1],\n                                        tf.tile(method.placeholders['s'][:, :1], [1, num_particles, 1]),\n                                        statistics['means'], statistics['stds'], statistics['state_step_sizes'])\n\n    odom = OdometryBaseline()\n    error_hist = 0.0\n    odom_error_hist = 0.0\n    for i in range(10000): # 100000\n        # define the inputs and train/run the model\n        batch = next(batch_iterator)\n        # define the inputs and train/run the model\n        input_dict = {**{method.placeholders[key]: batch[key] for key in 'osa'},\n                      }\n\n        # action_size = compute_sq_distance(batch['s'][:, 0, :], batch['s'][:, 1, :], state_step_sizes=statistics['state_step_sizes']) ** 0.5\n        action_size = abs(batch['s'][:, 0, 0] - batch['s'][:, 1, 0]) / statistics['state_step_sizes'][0]\n        action_size /= action_size\n\n        odom_pred = odom.predict(None, batch)\n        # odom_errors = compute_sq_distance(odom_pred[:, 1, :], batch['s'][:, 1, :], state_step_sizes=statistics['state_step_sizes']) ** 0.5\n        odom_errors = (odom_pred[:, 1, 0] - batch['s'][:, 1, 0]) / statistics['state_step_sizes'][0]\n        # odom_error_hist += np.histogram(odom_errors / action_size, 100, range=[0, 2])[0]\n        odom_error_hist += np.histogram(odom_errors / action_size, 101, range=[-1, 1])[0]\n\n        if i < 10000:\n            particles = session.run(motion_samples, input_dict)\n            # errors = compute_sq_distance(particles, batch['s'][:, 1, None, :], state_step_sizes=statistics['state_step_sizes']) ** 0.5\n            errors = (particles[:, :, 0] - odom_pred[:, 1, None, 0]) / statistics['state_step_sizes'][0]\n            # h, bins = np.histogram(errors / action_size[:, None], 100, range=[0, 2])\n            h, bins = np.histogram(errors / action_size[:, None], 101, range=[-1, 1])\n            error_hist += h\n        elif variant != 'e2e':\n            break\n\n    error_hist = error_hist / np.sum(error_hist) * 50.5\n    odom_error_hist = odom_error_hist / np.sum(odom_error_hist) * 50.5\n\n    plt.figure('motion statistics', [2.5,2.5])\n    plt.plot(bins[1:] - (bins[1]-bins[0])/2, error_hist, color=colors[variant], label=labels[variant])\n    if variant == 'e2e':\n        plt.plot(bins[1:] - (bins[1]-bins[0])/2, odom_error_hist, ':', color='k', label='True noise')\n    plt.xlabel('Predicted pos. relative to odom.')\n    plt.ylabel('Density')\n    # plt.legend()\n    plt.tight_layout()\n    plt.savefig('../plots/models/motion_statistics.pdf', transparent=True, dpi=600, frameon=False, facecolor='w', pad_inches=0.01)\n\n\n\ndef plot_models():\n    task = 'nav01'\n    data_path = '../data/100s'\n    test_data = load_data(data_path=data_path, filename=task + '_test')\n    noisy_test_data = noisyfy_data(reduce_data(test_data, 10))\n    num_examples = 10\n    # same seqlen and batchsize needed here!\n    # test_batch_iterator = make_batch_iterator(noisy_test_data, seq_len=50, batch_size=50)\n    test_batch_iterator = make_batch_iterator(noisy_test_data, seq_len=50, batch_size=num_examples)\n    batch = next(test_batch_iterator)\n\n    # for i in range(num_examples):\n    #     plot_observation(batch, i=0, t=i)\n\n    predictions = dict()\n\n    for variant, file_name in {\n                               'ind_e2e': '2017-12-23_03:32:47_compute-0-9_nav01_pf_ind_e2e_1000',\n                               # 'ind_e2e': '2017-12-22_18:30:30_compute-0-1_nav02_pf_ind_e2e_1000',\n                               # 'lstm': '2017-12-24_13:25:53_compute-0-1_nav01_lstm_1000',\n                               # 'lstm': '2017-12-22_18:29:21_compute-1-2_nav02_lstm_1000',\n                               # 'ind': '2017-12-23_00:48:08_compute-0-74_nav01_pf_ind_500',\n                               # 'e2e': '2017-12-22_18:29:49_compute-0-15_nav01_pf_e2e_500',\n                               }.items():\n\n        with open('../log/lc/'+file_name, 'rb') as f:\n            log = pickle.load(f)\n        hyper_params = log['hyper_params'][0]\n        model_path = '../models/' + log['exp_params'][0]['model_path'].split('/models/')[-1] # ['exp_params']['model_path]\n\n        # reset tensorflow graph\n        tf.reset_default_graph()\n\n        # instantiate method\n        if 'lstm' in variant:\n            method = RNN(**hyper_params['global'])\n        else:\n            method = DPF(**hyper_params['global'])\n\n        with tf.Session() as session:\n            # load method and apply to new data\n            statistics = method.load(session, model_path)\n            # print('predicting now')\n            # predictions[variant] = method.predict(session, batch, num_particles=1000, return_particles=False)\n            # print('prediction done')\n            # plot_measurement_model(session, method, statistics, batch, task, num_examples, variant)\n            # plot_proposer(session, method, statistics, batch, task, num_examples, variant)\n            # plot_motion_model(session, method, statistics, batch, task, 10, 50, variant)\n            plot_particle_filter(session, method, statistics, batch, task, num_examples, 1000, variant)\n\n    print(predictions.keys())\n    # plot_prediction(predictions['ind_e2e'], predictions['lstm'], statistics, batch, task, num_examples, variant)\n\n    plt.pause(10000.0)\n\ndef plot_statistics():\n    task = 'nav02'\n    data_path = '../data/100s'\n    test_data = load_data(data_path=data_path, filename=task + '_test')\n    noisy_test_data = noisyfy_data(test_data)\n    # noisy_test_data = noisyfy_data(test_data)\n    batch_size = 32\n    test_batch_iterator = make_batch_iterator(noisy_test_data, seq_len=2, batch_size=batch_size)\n\n\n    filenames = {              'ind_e2e': '2017-12-22_18:30:30_compute-0-1_nav02_pf_ind_e2e_1000',\n                               'ind': '2017-12-23_06:56:07_compute-0-26_nav02_pf_ind_1000',\n                               'e2e': '2017-12-24_00:51:18_compute-1-0_nav02_pf_e2e_1000',\n                               }\n\n    for variant in ['ind', 'e2e', 'ind_e2e']:\n        file_name = filenames[variant]\n\n        with open('../log/lc/'+file_name, 'rb') as f:\n            log = pickle.load(f)\n        hyper_params = log['hyper_params'][0]\n        model_path = '../models/' + log['exp_params'][0]['model_path'].split('/models/')[-1] # ['exp_params']['model_path]\n\n        # reset tensorflow graph\n        tf.reset_default_graph()\n\n        # instantiate method\n        method = DPF(**hyper_params['global'])\n\n        with tf.Session() as session:\n            # load method and apply to new data\n            statistics = method.load(session, model_path)\n            plot_measurement_statistics(session, method, statistics, test_batch_iterator, batch_size, variant)\n            plot_motion_statistics(session, method, statistics, test_batch_iterator, task, variant)\n\n    plt.pause(10000.0)\n\nplot_models()\n# plot_statistics()\n"
  },
  {
    "path": "plotting/plotting_kitti.py",
    "content": "import tensorflow as tf\n\nfrom methods.dpf_kitti import DPF\nfrom methods.odom import OdometryBaseline\nfrom utils.data_utils_kitti import load_data, noisyfy_data, make_batch_iterator, remove_state, split_data, load_kitti_sequences, make_batch_iterator_for_evaluation, wrap_angle, plot_video\nfrom utils.exp_utils_kitti import get_default_hyperparams\nimport matplotlib.pyplot as plt\nimport numpy as np\n\ndef get_evaluation_stats(model_path='../models/tmp/', test_trajectories=[9], seq_lengths = [100], plot_results=True):\n\n    data = load_kitti_sequences(test_trajectories)\n    # data = load_all_data(test_trajectories, train=False)\n    # data['o'] = data['o-m']  # flip ops to apply models that were trained on inverted data\n    # plot_video(data)\n\n    # reset tensorflow graph\n    tf.reset_default_graph()\n\n    # instantiate method\n    hyperparams = get_default_hyperparams()\n    method = DPF(**hyperparams['global'])\n\n    with tf.Session() as session:\n\n        # load method and apply to new data\n        # method.load(session, model_path)\n\n        errors = dict()\n\n        for i, test_traj in enumerate(test_trajectories):\n            # pick statest for traj\n            s_test_traj = data['s'][0:data['seq_num'][i*2]]  # take care of duplicated trajectories (left and right camera)\n            distance = compute_distance_for_trajectory(s_test_traj)\n            errors[test_traj] = dict()\n\n            for seq_len in seq_lengths:\n\n                errors[test_traj][seq_len] = {'trans': [], 'rot': []}\n\n                for start_step in range(0, 1):\n\n                    # print('start_step:', start_step)\n\n                    # end_step, dist = find_end_step(distance, start_step, seq_len, use_meters=False)\n                    # print('!!!', start_step, seq_len[seq_len], end_step, dist)\n                    end_step = distance.shape[0]\n                    dist = distance[-1]\n\n                    if end_step == -1:\n                        continue\n\n                    # test_batch_iterator = make_batch_iterator(test_data, seq_len=50)\n                    test_batch_iterator = make_batch_iterator_for_evaluation(data, start_step, trajectory=0, batch_size=1, seq_len=end_step-start_step)\n\n                    batch = next(test_batch_iterator)\n                    # batch_input = remove_state(batch, provide_initial_state=True)\n\n                    # prediction, particle_list, particle_prob_list = method.predict(session, batch_input, return_particles=True)\n                    # np.savez('./plot_results_traj_9', prediction, particle_list, particle_prob_list)\n                    npzfile = np.load('plot_results_traj_9.npz')\n                    prediction = npzfile['arr_0']\n                    particle_list = npzfile['arr_1']\n                    particle_prob_list = npzfile['arr_2']\n                    error_x = batch['s'][0, -1, 0] - prediction[0, -1, 0]\n                    error_y = batch['s'][0, -1, 1] - prediction[0, -1, 1]\n                    error_trans = np.sqrt(error_x ** 2 + error_y ** 2) / dist\n                    error_rot = abs(wrap_angle(batch['s'][0, -1, 2] - prediction[0, -1, 2]))/dist * 180 / np.pi\n\n                    errors[test_traj][seq_len]['trans'].append(error_trans)\n                    errors[test_traj][seq_len]['rot'].append(error_rot)\n\n                    if plot_results:\n\n                        dim_names = ['pos']\n                        fig1 = plt.figure(figsize=[3,3])\n                        fig2 = plt.figure(figsize=[3,3])\n                        grid = plt.GridSpec(3, 6)\n                        ax1 = fig2.add_subplot(111)\n                        ax2 = fig1.add_subplot(grid[0, :3])\n                        ax3 = fig1.add_subplot(grid[0, 3:6])\n                        ax4 = fig1.add_subplot(grid[1, :3])\n                        ax5 = fig1.add_subplot(grid[1, 3:6])\n                        ax6 = fig1.add_subplot(grid[2, :3])\n                        ax7 = fig1.add_subplot(grid[2, 3:6])\n                        # ax4 = fig.add_subplot(224)\n                        # ax6 = fig.add_subplot(326)\n                        # for t in range(particle_list.shape[1]):\n                        dim = 0\n                            # ax1.scatter(particle_list[0, t, :, dim], particle_list[0, t, :, dim+1], c=particle_prob_list[0, t, :], cmap='viridis_r', marker='o', s=1, alpha=0.1,\n                            #                     linewidths=0.05,\n                            #                     vmin=0.0,\n                            #                     vmax=0.02)\n\n                        ax1.plot(prediction[0, :, dim], prediction[0, :, dim+1], 'b')\n\n                        ax1.plot(batch['s'][0, :, dim], batch['s'][0, :, dim+1], 'r')\n                        ax1.plot(batch['s'][0, 100:350:100, dim], batch['s'][0, 100:350:100, dim+1], 'ok', markersize=3, markerfacecolor='None')\n                        ax1.plot(batch['s'][0, 0, dim], batch['s'][0, 0, dim+1], 'xk', markersize=5, markerfacecolor='None')\n\n                        ax1.set_aspect('equal')\n                        ax1.set_ylim([-450, 320])\n                            # ax2.scatter(particle_list[0, t, :, dim], particle_list[0, t, :, dim+1], c=particle_prob_list[0, t, :], cmap='viridis_r', marker='o', s=1, alpha=0.1,\n                            #                     linewidths=0.05,\n                            #                     vmin=0.0,\n                            #                     vmax=0.02)\n                            #\n                            # ax1.plot([prediction[0, t, dim]], [prediction[0, t, dim+1]], 'o', markerfacecolor='None', markeredgecolor='b',\n                            #                  markersize=0.5)\n                            #\n                            # ax1.plot([batch['s'][0, t, dim]], [batch['s'][0, t, dim+1]], '+', markerfacecolor='None', markeredgecolor='r',\n                            #                  markersize=0.5)\n\n                        ax2.imshow(np.clip(batch['o'][0, 100, :, :, 0:3]/255.0, 0.0, 1.0), interpolation='nearest')\n                        ax3.imshow(np.clip(batch['o'][0, 100, :, :, 3:6]/255.0 + 0.5, 0.0, 1.0), interpolation='nearest')\n                        ax4.imshow(np.clip(batch['o'][0, 200, :, :, 0:3]/255.0, 0.0, 1.0), interpolation='nearest')\n                        ax5.imshow(np.clip(batch['o'][0, 200, :, :, 3:6]/255.0 + 0.5, 0.0, 1.0), interpolation='nearest')\n                        ax6.imshow(np.clip(batch['o'][0, 300, :, :, 0:3]/255.0, 0.0, 1.0), interpolation='nearest')\n                        ax7.imshow(np.clip(batch['o'][0, 300, :, :, 3:6]/255.0 + 0.5, 0.0, 1.0), interpolation='nearest')\n                        ax2.set_axis_off()\n                        ax3.set_axis_off()\n                        ax4.set_axis_off()\n                        ax5.set_axis_off()\n                        ax6.set_axis_off()\n                        ax7.set_axis_off()\n                        # ax2.set_axis_off()\n                        ax1.set_xlabel('x (m)')\n                        ax1.set_ylabel('y (m)')\n                        ax1.legend(['Predicted pose','Ground truth'])\n                        # ax1.set_title(dim_names[0])\n                        # ax2.set_title(dim_names[1])\n                        # ax3.set_title(dim_names[2])\n                        # ax4.set_title(dim_names[3])\n                        fig1.savefig('{}.pdf'.format('test'), bbox_inches='tight')\n                        fig2.savefig('{}.pdf'.format('test2'), bbox_inches='tight')\n                        # plt.savefig('../plots/800_{}'.format(start_step))\n\n    return errors\n\n\n\ndef compute_distance_for_trajectory(s):\n\n    # for ii in range(len(output_oxts_file)):\n    distance = [0]\n    for i in range(1, s.shape[0]):\n        diff_x = s[i, 0, 0] - s[i-1, 0, 0]\n        diff_y = s[i, 0, 1] - s[i-1, 0, 1]\n        dist = distance[-1] + np.sqrt(diff_x ** 2 + diff_y ** 2)\n        distance.append(dist)\n    distance = np.asarray(distance)\n    return distance\n\ndef find_end_step(distance, start_step, length, use_meters=True):\n\n    for i in range(start_step, distance.shape[0]):\n        if (use_meters and distance[i] > (distance[start_step] + length)) or \\\n            (not use_meters and (i - start_step) >= length):\n            return i, distance[i] - distance[start_step]\n    return -1, 0\n\ndef find_all_cross_val_models(model_path):\n    import os\n    models = ([name for name in os.listdir(model_path) if not os.path.isfile(os.path.join(model_path, name))])\n    trajs = [int(name.split('_')[3]) for name in models]\n    return zip(models, trajs)\n\nif __name__ == '__main__':\n    plt.ion()\n\n    # errors = dict()\n    # average_errors = {'trans': {i: [] for i in [100, 200, 400, 800]},\n    #                   'rot': {i: [] for i in [100, 200, 400, 800]}}\n    # model_path = '../models/tmp/Cross_validation_plot/'\n    # for model, traj in find_all_cross_val_models(model_path):\n    #     print('!!! Evaluatng model {} on trajectory {}'.format(model, traj))\n    new_errors = get_evaluation_stats()\n        # errors.update(new_errors)\n        # print('')\n        # print('Trajectory {}'.format(traj))\n        # for seq_len in sorted(errors[traj].keys()):\n        #     for measure in ['trans', 'rot']:\n        #         e = errors[traj][seq_len][measure]\n        #         mean_error = np.mean(e)\n        #         se_error = np.std(e, ddof=1) / np.sqrt(len(e))\n        #         average_errors[measure][seq_len].append(mean_error)\n        #         print('{:>5} error for seq_len {}: {:.4f}+-{:.4f}'.format(measure, seq_len, mean_error, se_error))\n        #\n        # print('Averaged errors:')\n        # for measure in ['trans', 'rot']:\n        #     mean_error_over_all_subsequences = []\n        #     for seq_len in sorted(average_errors[measure].keys()):\n        #         e = np.array(average_errors[measure][seq_len])\n        #         e = e[~np.isnan(e)]\n        #         mean_error = np.mean(e)\n        #         se_error = np.std(e, ddof=1) / np.sqrt(len(e))\n        #         mean_error_over_all_subsequences.append(mean_error)\n        #         print('{:>5} error for seq_len {}: {:.4f}+-{:.4f}'.format(measure, seq_len, mean_error, se_error))\n        #     print('{:>5} mean error over all sequence_lengths: {:.4f}'.format(measure, np.mean(np.asarray(mean_error_over_all_subsequences))))\n"
  },
  {
    "path": "plotting/swap_plot.py",
    "content": "import pickle\nimport numpy as np\nimport matplotlib\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D\nimport itertools\nimport os\n\nresults = None\n\n# matplotlib.rcParams.update({'font.size': 12})\n\ncolor_list = plt.cm.tab10(np.linspace(0, 1, 10))\ncolors = {'lstm': color_list[0], 'pf_e2e': color_list[1], 'pf_ind_e2e': color_list[2], 'pf_ind': color_list[3]}\nlabels = {'lstm': 'LSTM', 'pf_e2e': 'DPF (e2e)', 'pf_ind_e2e': 'DPF (ind+e2e)', 'pf_ind': 'DPF (ind)', 'ff': 'FF', 'odom': 'Odom. baseline'}\n# conditions = ['normal', 'no_motion_likelihood', 'learn_odom', 'no_proposer']\n# conditions = ['normal', 'learn_odom', 'no_inject']\n# clabels = {'normal': 'Default', 'no_motion_likelihood': 'W/o motion likelihood', 'learn_odom': 'Learned odometry', 'no_proposer': 'W/o particle proposer', 'no_inject': \"No inject\"}\n\nif False:\n    test_conditions = ['odom5_imgTG', 'odom10_imgTG']\n    \n    conditions = ['orig_odom5_imgTG', 'odom5_imgTG_odom5_imgTG',\n                  'odom10_imgTG_odom5_imgTG', 'odom5_imgTG_odom10_imgTG',\n                    'odom10_imgTG_odom10_imgTG', 'orig_odom10_imgTG']\n\n    clabels = {'orig_odom5_imgTG':'5',\n               'odom5_imgTG_odom5_imgTG':'mo5,me5',\n               'odom5_imgTG_odom10_imgTG':'mo5,me10',\n               'orig_odom10_imgTG':'10',\n               'odom10_imgTG_odom10_imgTG': 'mo10,me10',\n               'odom10_imgTG_odom5_imgTG': 'mo10,me5',\n               'odom5_imgTG': '5',\n               'odom10_imgTG': '10',\n                }\n    xlabels = ['A', 'B']\n    exp = 'swapmo'\nelse:\n\n    test_conditions = ['odom10_imgG', 'odom10_imgTG']\n\n    conditions = ['orig_odom10_imgG', 'odom10_imgG_odom10_imgG', 'odom10_imgG_odom10_imgTG',\n                  'odom10_imgTG_odom10_imgG', 'odom10_imgTG_odom10_imgTG', 'orig_odom10_imgTG']\n\n    # clabels = {'orig_odom10_imgG':'G',\n    #            'odom10_imgG_odom10_imgG':'meG,moG',\n    #            'odom10_imgG_odom10_imgTG':'meG,moTG',\n    #            'orig_odom10_imgTG':'TG',\n    #            'odom10_imgTG_odom10_imgTG': 'meTG,moTG',\n    #            'odom10_imgTG_odom10_imgG': 'meTG,moG',\n    #            'odom10_imgG': 'G',\n    #            'odom10_imgTG': 'TG',\n    #             }\n\n    clabels = {'orig_odom10_imgG':'A(A)*',\n           'odom10_imgG_odom10_imgG':'A(A)',\n           'odom10_imgG_odom10_imgTG':'A(B)',\n           'orig_odom10_imgTG':'B(B)*',\n           'odom10_imgTG_odom10_imgTG': 'B(B)',\n           'odom10_imgTG_odom10_imgG': 'B(A)',\n           'odom10_imgG': 'A',\n           'odom10_imgTG': 'B',\n            }\n    xlabels = ['C', 'D']\n    exp = 'swapme'\n\nvmax = 0.4\n\ntask = 'nav02'\n# methods = ['pf_ind', 'pf_e2e', 'pf_ind_e2e', 'lstm']\nmethods = ['pf_ind_e2e']\n\n# load results\nresults = dict()\n\ncount = 0\nfor cond in conditions:\n    # log_path = '/home/rbo/Desktop/log/'+task+'_ab1'\n    log_path = '../log/'+exp+'/'+cond\n    for filename in [f for f in os.listdir(log_path) if os.path.isfile(os.path.join(log_path, f))]:\n        full_filename = os.path.join(log_path, filename)\n        print('loading {}:'.format(count) + full_filename + ' ...')\n        try:\n            # if 'DeepThought' not in filename:\n            # if 'DeepThought' in filename:\n            with open(full_filename, 'rb') as f:\n                result = pickle.load(f)\n                # result_name = result['task'][0] + '/' + result['method'][0] + '/' + str(result['num_episodes'][0]) + '/' + result['condition'][0]\n                result_name = cond #+ '_' + result['exp_params'][0]['file_ending'] #result['exp_params'][0]['task'] + '/' + result['exp_params'][0]['method'] + '/' + str(result['exp_params'][0]['num_episodes']) + '/' + result['exp_params'][0]['ab_cond']\n                print(result_name)\n                if result_name not in results.keys():\n                    results[result_name] = result\n                else:\n                    for key in result.keys():\n                        if key in results[result_name].keys():\n                            results[result_name][key] += result[key]\n                        else:\n                            results[result_name][key] = result[key]\n                        # print(result_name, key)\n                count += 1\n        except Exception as e:\n            print(e)\n            print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')\n\nprint()\nfor result_name, r in results.items():\n    print(result_name, len(r['exp_params']))\n    # print(result_name, len(r['test_odom5_imgTG_mse']))\n\nprint('Loaded {} results'.format(count))\n\n\ntask = 'nav02'\nstep = 3\n\nepisodes = [1000]\n# episodes = [1000]\n\n\nmeans = []\nses = []\n\nfor c, condition in enumerate(conditions):\n\n    means.append(np.zeros([len(test_conditions), 5]))\n    ses.append(np.zeros([len(test_conditions), 5]))\n\n    for tc, test_condition in enumerate(test_conditions):\n\n        result_name = condition\n        if result_name in results.keys():\n            result = results[result_name]\n\n            hist = np.array([[h[i] for i in range(0, 50, 10)] for h in result['test_'+test_condition+'_hist' ]])  # result x time x sqe [.0, 0.1, .., 10.0]\n            err = 1. - np.sum(hist[:,:,:10], axis=-1) # sqe < 1.0\n            # err = np.sum(hist[:,:,:10], axis=-1) # sqe < 1.0\n            print(result_name, err)\n            means[c][tc] = np.mean(err, axis=0)\n            ses[c][tc] = np.std(err, axis=0, ddof=1) / np.sqrt(len(err))\n\n        else:\n            print(result_name + 'not found')\n            means[tc] *= np.nan\n            ses[tc] *= np.nan\n\n\n# means[num_episodes]['min'] = np.stack([means[num_episodes][method] for method in methods], axis=0).min(axis=1)\n\n# ax = fig.add_subplot(111)\n# # Turn off axis lines and ticks of the big subplot\n# ax.spines['top'].set_color('none')\n# ax.spines['bottom'].set_color('none')\n# ax.spines['left'].set_color('none')\n# ax.spines['right'].set_color('none')\n# ax.tick_params(labelcolor='w', top='off', bottom='off', left='off', right='off')\n\nm = np.array(means)[:, :, step].T\ns = np.array(ses)[:, :, step].T\n\nfor i in range(2):\n    z = m[i, :]\n    y = s[i, :]\n    x = [[z[0]] * 2 + [z[2]] * 2,\n         [z[1]] * 2 + [z[2]] * 2,\n         [z[3]] * 2 + [z[5]] * 2,\n         [z[3]] * 2 + [z[4]] * 2]\n    plt.figure(i, [1.35,1.35])\n    plt.imshow(x, interpolation='nearest', vmin=-0.33*vmax, vmax=vmax, cmap='viridis_r')\n    plt.plot([-0.5, 3.5],[1.5, 1.5], '-w', linewidth=0.5)\n    plt.plot([1.5, 1.5],[-0.5, 3.5], '-w', linewidth=0.5)\n\n    for j, x_coord, y_coord, value, s_value in [\n            (0, 0.5, 0, z[0], y[0]),\n            (1, 0.5, 1, z[1], y[1]),\n            (2, 2.5, 0.5, z[2], y[2]),\n            (3, 0.5, 2.5, z[3], y[3]),\n            (5, 2.5, 2, z[5], y[5]),\n            (4, 2.5, 3, z[4], y[4])]:\n        if j == 0 or j == 5:\n            # text = '{:.4s}*\\n+-{:.4s}'.format('{:.3f}'.format(value)[1:],'{:.2f}'.format(s_value)[1:])\n            text = ' {:.4s}*'.format('{:.3f}'.format(value)[1:],'{:.2f}'.format(s_value)[1:])\n        else:\n            # text = '{:.4s}\\n+-{:.4s}'.format('{:.3f}'.format(value)[1:],'{:.2f}'.format(s_value)[1:])\n            text = '{:.4s}'.format('{:.3f}'.format(value)[1:],'{:.2f}'.format(s_value)[1:])\n\n        plt.text(x_coord, y_coord, text, va='center', ha='center', color='white', fontweight='normal')\n\n    plt.gca().set_aspect('equal')\n\n    plt.xlabel('Motion model')\n    plt.xticks([0.5, 2.5], xlabels)\n    plt.ylabel('Measurem. model')\n    plt.yticks([0.5, 2.5], xlabels)\n    plt.tight_layout(0.0, 0.0, 0.0)\n    print('saving')\n    plt.savefig('../plots/cr/'+exp+'%s.pdf' % i, bbox_inches=\"tight\", transparent=True, dpi=600, frameon=True, facecolor='w', pad_inches=0.01)\n\nplt.figure('colorbar', [0.6, 1.35])\na = np.array([[0.0, 0.3]])\nimg = plt.imshow(a, cmap=\"viridis_r\", vmin=-0.33*vmax, vmax=vmax)\nplt.gca().set_visible(False)\ncax = plt.axes([0.0, 0.2, 0.1, 0.65])\nplt.colorbar(orientation=\"vertical\", cax=cax, label='Error rate', boundaries=np.linspace(0,0.4,100), ticks=np.linspace(0.0, 0.4, 5))\n\nplt.savefig('../plots/cr/colorbar.pdf'.format(s), transparent=True, dpi=600, frameon=False, facecolor='w', pad_inches=0.01)\n\n\nplt.show()\n\n"
  },
  {
    "path": "setup.sh",
    "content": "#!/bin/bash\necho\necho 'Creating additional folders .. '\necho \nmkdir models\nmkdir log\nmkdir plots\necho 'Downloading data (2.5GB, this might take a bit) .. '\necho\nwget -N 'https://depositonce.tu-berlin.de/bitstreams/fe02c1e0-64d9-4a92-ac4d-a8a0ef455c8f/download'\necho 'Unpacking data .. '\necho\nunzip download\nrm download\n"
  },
  {
    "path": "utils/__init__.py",
    "content": ""
  },
  {
    "path": "utils/data_utils.py",
    "content": "import numpy as np\nimport matplotlib.pyplot as plt\nimport os\n\nfrom utils.plotting_utils import plot_trajectories, plot_maze, plot_observations, plot_trajectory\n\ndef wrap_angle(angle):\n    return ((angle - np.pi) % (2 * np.pi)) - np.pi\n\ndef mix_data(file_in1, file_in2, file_out, steps_per_episode=100, num_episodes=1000):\n    data1 = dict(np.load(file_in1))\n    data2 = dict(np.load(file_in2))\n    data_mix = dict()\n    for key in data1.keys():\n        d1 = data1[key][:steps_per_episode*num_episodes//2]\n        d2 = data2[key][:steps_per_episode*num_episodes//2]\n        data_mix[key] = np.concatenate((d1, d2), axis=0)\n    np.savez(file_out, **data_mix)\n\ndef average_nn(states_from, states_to, step_sizes, num_from=10, num_to=100):\n\n    states_from = np.reshape(states_from, [-1, 3])\n    states_to = np.reshape(states_to, [-1, 3])\n\n    idx_from = np.random.choice(len(states_from), num_from)\n    idx_to = np.random.choice(len(states_to), num_to)\n\n    sum = 0.0\n    for i in range(3):\n        diff = states_from[idx_from, None, i] - states_to[None, idx_to, i]\n        if i == 2:\n            diff = wrap_angle(diff)\n        sum += (diff / step_sizes[i])**2\n    average_dist = np.mean(np.min(sum, axis=1) > 0.5)\n    return average_dist\n\ndef load_data(data_path='../data/100s', filename='nav01_train', steps_per_episode=100, num_episodes=None):\n\n    # data = dict(np.load(os.path.join(data_path, '100s', filename + '.npz')))\n    data = dict(np.load(os.path.join(data_path, filename + '.npz')))\n\n    # reshape data\n    for key in data.keys():\n        # 'vel': (100, 1000, 3), 'rgbd': (100, 1000, 32, 32, 4), 'pose': (100, 1000, 3)\n        if num_episodes is not None:\n            data[key] = data[key][:num_episodes*steps_per_episode]\n        data[key] = np.reshape(data[key], [-1, steps_per_episode] + list(data[key].shape[1:])).astype('float32')\n\n    # convert degrees into radients and\n    for key in ['pose', 'vel']:\n        data[key][:, :, 2] *= np.pi / 180\n    # angles should be between -pi and pi\n    data['pose'][:, :, 2] = wrap_angle(data['pose'][:, :, 2])\n\n    abs_d_x = (data['pose'][:, 1:, 0:1] - data['pose'][:, :-1, 0:1])\n    abs_d_y = (data['pose'][:, 1:, 1:2] - data['pose'][:, :-1, 1:2])\n    d_theta = wrap_angle(data['pose'][:, 1:, 2:3] - data['pose'][:, :-1, 2:3])\n    s = np.sin(data['pose'][:, :-1, 2:3])\n    c = np.cos(data['pose'][:, :-1, 2:3])\n    rel_d_x = c * abs_d_x + s * abs_d_y\n    rel_d_y = s * abs_d_x - c * abs_d_y\n\n\n    # define observations, states, and actions for the filter, use current and previous velocity measurement as action\n    # and ignore the 0th timestep because we don't have the previous velocity of that step\n    return {'o': data['rgbd'][:, 1:, :, :, :3],\n            's': data['pose'][:, 1:, :],\n            'a': np.concatenate([rel_d_x, rel_d_y, d_theta], axis=-1)}\n            # 'a': np.concatenate([data['vel'][:, :-1, None, :], data['vel'][:, 1:, None, :]], axis=-2)}\n\n\ndef compute_staticstics(data):\n\n    means = dict()\n    stds = dict()\n    state_step_sizes = []\n    state_mins = []\n    state_maxs = []\n\n    for key in 'osa':\n        # compute means\n        means[key] = np.mean(data[key], axis=(0, 1), keepdims=True)\n        if key == 's':\n            means[key][:, :, 2] = 0  # don't touch orientation because we'll feed this into cos/sin functions\n        if key == 'a':\n            means[key][:, :, :] = 0  # don't change means of velocities, 0.0, positive and negative values have semantics\n\n        # compute stds\n        axis = tuple(range(len(data[key].shape) - 1))  # compute std by averaging over all but the last dimension\n        stds[key] = np.std(data[key] - means[key], axis=axis, keepdims=True)\n        if key == 's':\n            stds[key][:, :, :2] = np.mean(stds[key][:, :, :2])  # scale x and by by the same amount\n        if key == 'a':\n            stds[key][:, :, :2] = np.mean(stds[key][:, :, :2])  # scale x and by by the same amount\n\n    # compute average step size in x, y, and theta for the distance metric\n    for i in range(3):\n        steps = np.reshape(data['s'][:, 1:, i] - data['s'][:, :-1, i], [-1])\n        if i == 2:\n            steps = wrap_angle(steps)\n        state_step_sizes.append(np.mean(abs(steps)))\n    state_step_sizes[0] = state_step_sizes[1] = (state_step_sizes[0] + state_step_sizes[1]) / 2\n    state_step_sizes = np.array(state_step_sizes)\n\n    # compute min and max in x, y and theta\n    for i in range(3):\n        state_mins.append(np.min(data['s'][:, :, i]))\n        state_maxs.append(np.max(data['s'][:, :, i]))\n    state_mins = np.array(state_mins)\n    state_maxs = np.array(state_maxs)\n\n    return means, stds, state_step_sizes, state_mins, state_maxs\n\n\ndef split_data(data, ratio=0.8, categories=['train', 'val']):\n    print('SPLIT {}'.format(data['s'].shape))\n    split_data = {categories[0]: dict(), categories[1]: dict()}\n    for key in data.keys():\n        split_point = int(data[key].shape[0] * ratio)\n        split_data[categories[0]][key] = data[key][:split_point]\n        split_data[categories[1]][key] = data[key][split_point:]\n    for key in split_data.keys():\n        print('SPLIT --> {}: {}'.format(key, len(split_data[key]['s'])))\n    return split_data\n\n\ndef reduce_data(data, num_episodes):\n    new_data = dict()\n    for key in 'osa':\n        new_data[key] = data[key][:num_episodes]\n    return new_data\n\ndef shuffle_data(data):\n    new_data = dict()\n    shuffled_indices = np.random.permutation(len(data['o']))\n    for key in 'osa':\n        new_data[key] = data[key][shuffled_indices]\n    return new_data\n\ndef remove_state(data, provide_initial_state=False):\n    new_data = dict()\n    new_data['o'] = data['o']\n    new_data['a'] = data['a']\n    if provide_initial_state:\n        new_data['s'] = data['s'][..., :1, :]\n    return new_data\n\ndef noisify_data_condition(data, condition):\n    print('condition', condition)\n    if condition == 'odom0_imgTG':\n        return noisyfy_data(data, odom_noise_factor=0.0)\n    elif condition == 'odom5_imgTG':\n        return noisyfy_data(data, odom_noise_factor=0.5)\n    elif condition == 'odom10_imgTG':\n        return noisyfy_data(data)\n    elif condition == 'odom20_imgTG':\n        return noisyfy_data(data, odom_noise_factor=2.0)\n    elif condition == 'odomX_imgTG':\n        data = noisyfy_data(data, odom_noise_factor=0.0)\n        # shuffle actions to basically make them meaningless\n        shape = data['a'].shape\n        a = np.reshape(data['a'], [-1, shape[-1]])\n        np.random.shuffle(a)\n        data['a'] = np.reshape(a, shape)\n        return data\n    elif condition == 'odom10_imgC':\n        return noisyfy_data(data, img_noise_factor=0.0, img_random_shift=False)\n    elif condition == 'odom10_imgG':\n        return noisyfy_data(data, img_noise_factor=1.0, img_random_shift=False)\n    elif condition == 'odom10_imgT':\n        return noisyfy_data(data, img_noise_factor=0.0, img_random_shift=True)\n    elif condition == 'odom10_imgX':\n        data = noisyfy_data(data, img_noise_factor=0.0, img_random_shift=False)\n        # shuffle observations to basically make them meaningless\n        shape = data['o'].shape\n        o = np.reshape(data['o'], [-1, shape[-1]])\n        np.random.shuffle(o)\n        data['o'] = np.reshape(o, shape)\n        return data\n\ndef noisyfy_data(data, odom_noise_factor=1.0, img_noise_factor=1.0, img_random_shift=True):\n    print(\"noisyfying data ... \")\n    new_data = dict()\n    new_data['s'] = data['s']\n    new_data['a'] = data['a'] * np.random.normal(1.0, 0.1 * odom_noise_factor, data['a'].shape)\n    new_o = np.zeros([data['o'].shape[0], data['o'].shape[1], 24, 24, 3])\n    for i in range(data['o'].shape[0]):\n        for j in range(data['o'].shape[1]):\n            if img_random_shift:\n                offsets = np.random.random_integers(0, 8, 2)\n            else:\n                offsets = (4, 4)\n            new_o[i, j] = data['o'][i, j, offsets[0]:offsets[0]+24, offsets[1]:offsets[1]+24, :]\n    new_o += np.random.normal(0.0, 20 * img_noise_factor, new_o.shape)\n    # for i in range(data['o'].shape[0]):\n    #     for j in range(data['o'].shape[1]):\n    #         plt.figure()\n    #         plt.imshow(new_o[i,j]/255, interpolation='nearest')\n    #         plt.figure()\n    #         plt.imshow(data['o'][i,j]/255, interpolation='nearest')\n    #         plt.show()\n    new_data['o'] = new_o\n    return new_data\n\n\ndef make_batch_iterator(data, batch_size=32, seq_len=10):\n    # go through data and select a subsequence from each sequence\n    while True:\n        episodes = np.random.random_integers(0, len(data['s']) - 1, size=batch_size)\n        start_steps = np.random.random_integers(0, len(data['s'][0]) - seq_len - 1, size=batch_size)\n        batches = {k: np.concatenate([data[k][i:i + 1, j:j + seq_len] for i, j in zip(episodes, start_steps)]) for k in data.keys()}\n        yield batches\n\ndef make_repeating_batch_iterator(data, epoch_len, batch_size=32, seq_len=10):\n    # go through data and select a subsequence from each sequence\n    repeating_episodes = np.random.random_integers(0, len(data['s']) - 1, size=[epoch_len, batch_size])\n    repeating_start_steps = np.random.random_integers(0, len(data['s'][0]) - seq_len - 1, size=[epoch_len, batch_size])\n    while True:\n        for episodes, start_steps in zip(repeating_episodes, repeating_start_steps):\n            batches = {k: np.concatenate([data[k][i:i + 1, j:j + seq_len] for i, j in zip(episodes, start_steps)]) for k in data.keys()}\n            yield batches\n\ndef make_complete_batch_iterator(data, batch_size=1000, seq_len=10):\n    num_episodes = len(data['s'])\n    num_start_steps = len(data['s'][0]) - seq_len\n    batch_indices = [(i, j) for i in range(num_episodes) for j in range(num_start_steps)]\n    while batch_indices != []:\n        batches = {k: np.concatenate([data[k][i:i + 1, j:j + seq_len] for (i, j) in batch_indices[:batch_size]]) for k in data.keys}\n        batch_indices = batch_indices[batch_size:]\n        yield batches\n\n\ndef compare_data_coverage():\n\n    task = 'nav02'\n\n    data = load_data(filename=task + '_train', data_path='../data/100s_mix', steps_per_episode=100, num_episodes=100)\n    means, stds, state_step_sizes, state_mins, state_maxs = compute_staticstics(data)\n    states = dict()\n    states['ab'] = data['s']\n    data = load_data(filename=task + '_train', data_path='../data/100s_astar', steps_per_episode=100, num_episodes=100)\n    states['b'] = data['s']\n    data = load_data(filename=task + '_train', data_path='../data/100s', steps_per_episode=100, num_episodes=100)\n    states['a'] = data['s']\n    # plt.figure()\n    # h, b = np.histogram(states['a'][:,:,2], bins=100)\n    # plt.plot(b[1:], h)\n    # h, b = np.histogram(states['b'][:,:,2], bins=100)\n    # plt.plot(b[1:], h)\n    # plt.show()\n    for f in ['a', 'b']:\n        for t in ['a', 'b', 'ab']:\n            d = average_nn(states_from=states[f], states_to=states[t], step_sizes=state_step_sizes, num_from=10000, num_to=10000)\n            print('{} <- {}: {}'.format(f, t, d))\n            plt.pause(0.01)\n\nif __name__ == '__main__':\n\n    # mix_data('../data/100s/nav02_test.npz',\n    #          '../data/100s_astar/nav02_test.npz',\n    #          '../data/100s_mix/nav02_test')\n    #\n    # compare_data_coverage()\n\n    task = 'nav03'\n\n    # data = load_data(filename=task + '_train')\n    data = load_data(filename=task + '_train', data_path='../data/100s', steps_per_episode=100, num_episodes=1000)\n    # data = noisyfy_data(data)\n\n    data = split_data(data, ratio=0.5)\n    # means, stds, state_step_sizes, state_mins, state_maxs = compute_staticstics(data)\n\n    # batch_iterator = make_batch_iterator(data['train'])\n    scaling = 0.5  # 0.5\n    if task == 'nav01':\n        plt.figure(figsize=[10*scaling,5*scaling])\n    elif task == 'nav02':\n        plt.figure(figsize=[15*scaling,9*scaling])\n    elif task == 'nav03':\n        plt.figure(figsize=[20*scaling,13*scaling])\n    # plot_trajectories(noisy_data, emphasize=2, mincolor=0.3)\n\n    # np.random.seed(11)\n    # nav02: i=108\n    i = 108\n    # for i in range(100, 120):\n    np.random.seed(i)\n    dat = shuffle_data(data['train'])\n    dat = reduce_data(dat, 1)\n    dat = noisyfy_data(dat)\n\n    plot_trajectory(dat, figure_name=None, emphasize=0, mincolor=0.0, linewidth=0.5)\n    plot_maze(task)\n    # plot_trajectories(data['val'], figure_name='2', emphasize=None, mincolor=0.0, linewidth=0.5)\n    # plot_maze(task)\n    plt.tick_params(top='off', bottom='off', left='off', right='off', labelleft='off', labelbottom='off')\n\n    plt.tight_layout()\n    # plt.savefig(\"../plots/\"+task +\".png\",\n    #            bbox_inches='tight',\n    #            transparent=False,\n    #            pad_inches=0,\n    #            dpi=200)\n    plt.savefig(\"../plots/\"+task +\".pdf\",\n               bbox_inches='tight',\n               transparent=False,\n               pad_inches=0)\n\n    plt.figure()\n    # plot_observations(data)\n    # plt.savefig(\"../plots/\"+ task +\"_obs.png\",\n    #            bbox_inches='tight',\n    #            transparent=False,\n    #            pad_inches=0,\n    #            dpi=200)\n    plot_observations(dat, n=5)\n    plt.savefig(\"../plots/\"+ task +\"_noisy_obs.pdf\",\n               bbox_inches='tight',\n               transparent=False,\n               pad_inches=0,\n               dpi=200)\n\n    plt.show()\n"
  },
  {
    "path": "utils/data_utils_kitti.py",
    "content": "import numpy as np\nimport matplotlib.pyplot as plt\nimport os\nimport math\nimport glob\nfrom time import time\nfrom PIL import Image\nfrom utils.plotting_utils import plot_trajectories, plot_maze, plot_observations\n\ndef wrap_angle(angle):\n    return ((angle - np.pi) % (2 * np.pi)) - np.pi\n\ndef rotation_matrix(x):\n    rot_psi = np.array([[math.cos(x[2]), -math.sin(x[2]), 0], [math.sin(x[2]), math.cos(x[2]), 0], [0, 0, 1]])\n    rot_theta = np.array([[math.cos(x[1]), 0, math.sin(x[1])], [0, 1, 0], [-math.sin(x[1]), 0, math.cos(x[1])]])\n    rot_phi = np.array([[1, 0, 0], [0, math.cos(x[0]), -math.sin(x[0])], [0, math.sin(x[0]), math.cos(x[0])]])\n    R = np.dot(rot_psi,np.dot(rot_theta,rot_phi))\n    return R\n\ndef read_oxts_data(oxts, oxts_prev, oxts_init):\n\n    with open(oxts, 'r') as f:\n        oxts_data = np.loadtxt(f)\n\n    with open(oxts_init, 'r') as f:\n        oxts_init = np.loadtxt(f)\n\n    with open(oxts_prev, 'r') as f:\n        oxts_prev = np.loadtxt(f)\n\n    north  = (oxts_data[0] - oxts_init[0]) * 6378137 * math.pi / 180\n    east = (oxts_data[1] - oxts_init[1]) * 6378137 * math.pi / 180 * math.cos(oxts_init[0] * math.pi / 180)\n    alpha = (oxts_data[22] - oxts_prev[22])/0.103\n    state = np.array([east, north, -oxts_data[5], oxts_data[8], -oxts_data[22]])\n    action = np.array([oxts_data[14], oxts_data[15], alpha])\n\n    return state, action\n\ndef load_image(img_file):\n    return np.asarray(Image.open(img_file), 'float32')\n\ndef image_input(img1, img2):\n    return np.concatenate((img1, img1-img2), axis=2)\n\ndef load_data_for_stats(oxts_data, images, diff_images, seq_num, base_frame):\n\n    state = np.zeros((len(oxts_data), 6))\n    action = np.zeros((len(oxts_data), 3))\n    with open(base_frame, 'r') as f:\n        data = np.loadtxt(f)\n        base_lat = data[0]\n        base_long = data[1]\n\n    for ii in range(len(oxts_data)):\n        with open(oxts_data[ii], 'r') as f:\n            data = np.loadtxt(f)\n        # if ii==0: #or ii in seq_num[:-1]:\n        #     base_lat = data[0]\n        #     base_long = data[1]\n        north = (data[0] - base_lat) * 6378137 * math.pi / 180\n        east = (data[1] - base_long) * 6378137 * math.pi / 180 * math.cos(base_lat * math.pi / 180)\n        state[ii,:] = np.array([north, east, data[5], data[6], data[7], data[22]])\n        action[ii,:] = np.array([data[8], data[14], data[15]])\n\n    images_per_seq = 100\n    obs = np.zeros((len(seq_num) * images_per_seq, 50, 150, 6))\n    for ii in range(1, len(seq_num)-1):\n        for jj in range(images_per_seq):\n            img1 = load_image(images[seq_num[ii - 1] + jj])\n            obs[images_per_seq*(ii-1)+jj,:,:,:3] = img1\n            img2 = load_image(diff_images[seq_num[ii - 1] + jj])\n            obs[images_per_seq*(ii-1)+jj,:,:,3:6] = img2\n\n\n    data_for_stats = {'s': state, 'a': action, 'o': obs}\n\n    return data_for_stats\n\n# loading all sequences for KITTI\ndef load_kitti_sequences(sequence_list=None):\n\n    print('Loading KITTI DATA')\n    t1 = time()\n    try:\n        if sequence_list is None:\n            print('Trying to load from cache ... ')\n            data = dict(np.load('../data/kitti.npz'))\n            t2 = time()\n            print('Done! ({:.2f}s)'.format(t2-t1))\n        else:\n            raise Exception\n\n    except:\n\n        if sequence_list is None:\n            sequence_list = list(range(11))\n\n        print('Cache not found, loading from KITTI_dataset')\n        path = \"../data/kitti\"\n\n        image_seq_1_full_path = [\"{}/{:02d}/image_2\".format(path, x) for x in sequence_list]\n        image_seq_2_full_path = [\"{}/{:02d}/image_3\".format(path, x) for x in sequence_list]\n\n        # Extract original image and difference image\n        input_image_file = []\n        seq_num = []\n        for ii in range(len(sequence_list)):\n            for name in glob.glob('{}/image*.png'.format(image_seq_1_full_path[ii])):\n                input_image_file = input_image_file + [name]\n            for name in glob.glob('{}/image*.png'.format(image_seq_2_full_path[ii])):\n                input_image_file = input_image_file + [name]\n\n        input_image_file.sort()\n        # print(len(input_image_file))\n\n        oxts_seq_1 = [\"%.2d_image1.txt\" % i for i in sequence_list]\n        oxts_seq_1 = oxts_seq_1 + [\"%.2d_image2.txt\" % i for i in sequence_list]\n        oxts_seq_1.sort()\n        oxts_seq_1_full_path = [\"{}/{}\".format(path, x) for x in oxts_seq_1]\n        output_oxts_file = oxts_seq_1_full_path\n\n        sequence_starts_ends = [[0, 4540], [0, 1100], [0, 4660], [0, 800], [0, 270], [0, 2760], [0, 1100], [0, 1100], [1100, 5170], [0, 1590],\n         [0, 1200]]\n        data_values = np.array([sequence_starts_ends[i] for i in sequence_list])\n        seq_num = np.zeros((2*data_values.shape[0],))\n        weights = np.zeros((2*data_values.shape[0],))\n\n        for ii in range(data_values.shape[0]):\n            if ii == 0:\n                seq_num[0] = data_values[ii,1] - data_values[ii,0]\n                seq_num[1] = seq_num[0] + data_values[ii,1] - data_values[ii,0]\n                weights[0] = weights[1] = data_values[ii,1] - data_values[ii,0]\n            else:\n                seq_num[2*ii] = seq_num[2*ii-1] + data_values[ii, 1] - data_values[ii, 0]\n                seq_num[2*ii+1] = seq_num[2*ii] + data_values[ii, 1] - data_values[ii, 0]\n                weights[2*ii] = weights[2*ii+1] = data_values[ii, 1] - data_values[ii, 0]\n\n        # seq_num is an array of the cumulative sequence lengths, e.g. [100, 300, 350] for sequences of length 100, 200, 50\n        seq_num = seq_num.astype(int)\n        weights = weights/seq_num[-1]\n        print(seq_num, weights)\n\n        o = np.zeros((seq_num[-1], 50, 150, 6))\n        count = 0\n        # for all sequences\n        for ii in range(len(seq_num)):\n            # find out the start and end of the current sequence\n            if ii == 0:\n                start = 1\n            else:\n                start = seq_num[ii-1]+ii+1\n\n            # load first image\n            prev_image = load_image(input_image_file[start-1])\n            # for all time steps\n            for jj in range(start, seq_num[ii]+ii+1):\n                # load next image\n                cur_image = load_image(input_image_file[jj])\n                # observation from current and last image\n                o[count, :, :, :] = image_input(cur_image, prev_image)\n                prev_image = cur_image\n                count += 1\n\n        a = np.zeros((seq_num[-1], 3))\n        s = np.zeros((seq_num[-1], 5))\n        for ii in range(len(output_oxts_file)):\n\n            # load text file\n            with open(output_oxts_file[ii], 'r') as f:\n                tmp = np.loadtxt(f)\n\n            start = 0 if ii == 0 else seq_num[ii-1]\n\n            x = tmp[:, 11]\n            y = -tmp[:, 3]\n            theta = -np.arctan2(-tmp[:, 8], tmp[:, 10])\n            s[start:seq_num[ii], 0] = x[1:]  # x\n            s[start:seq_num[ii], 1] = y[1:]  # y\n            s[start:seq_num[ii], 2] = theta[1:]  # angle\n            s[start:seq_num[ii], 3] = np.sqrt((y[1:] - y[:-1]) ** 2 + (x[1:] - x[:-1]) ** 2) / 0.103  # forward vel\n            s[start:seq_num[ii], 4] = wrap_angle(theta[1:] - theta[:-1])/0.103  # angular vel\n\n        t2 = time()\n        print('Done! ({:.2f}s)'.format(t2 - t1))\n        print('By default not saving data to cache ... ')\n        # if len(sequence_list) == 11:\n        #     print('Saving data to cache in ../data/kitti')\n        #     np.savez('../data/kitti', s=s, a=a, o=o, seq_num=seq_num, weights=weights)\n\n        print(s.shape, a.shape, o.shape, seq_num.shape, weights.shape)\n\n        data = {'s': s,\n                'a': a,\n                'o': o,\n                'seq_num': seq_num,\n                'weights': weights\n                }\n\n    for key in 'osa':\n        # add dimension to be consistent with the batch x seq x dim convention\n        data[key] = data[key][:, np.newaxis, :]\n\n    return add_mirrored_data(data)\n\n\ndef load_data(data_path='data/100s', filename='nav01_train', steps_per_episode=100, num_episodes=None):\n\n    data = dict(np.load(os.path.join(data_path, '100s', filename + '.npz')))\n    data = dict(np.load(os.path.join(data_path, filename + '.npz')))\n\n    # reshape data\n    for key in data.keys():\n        # 'vel': (100, 1000, 3), 'rgbd': (100, 1000, 32, 32, 4), 'pose': (100, 1000, 3)\n        if num_episodes is not None:\n            data[key] = data[key][:num_episodes*steps_per_episode]\n        data[key] = np.reshape(data[key], [-1, steps_per_episode] + list(data[key].shape[1:])).astype('float32')\n\n    # convert degrees into radients and\n    for key in ['pose', 'vel']:\n        data[key][:, :, 2] *= np.pi / 180\n    # angles should be between -pi and pi\n    data['pose'][:, :, 2] = wrap_angle(data['pose'][:, :, 2])\n\n    abs_d_x = (data['pose'][:, 1:, 0:1] - data['pose'][:, :-1, 0:1])\n    abs_d_y = (data['pose'][:, 1:, 1:2] - data['pose'][:, :-1, 1:2])\n    d_theta = wrap_angle(data['pose'][:, 1:, 2:3] - data['pose'][:, :-1, 2:3])\n    s = np.sin(data['pose'][:, :-1, 2:3])\n    c = np.cos(data['pose'][:, :-1, 2:3])\n    rel_d_x = c * abs_d_x + s * abs_d_y\n    rel_d_y = s * abs_d_x - c * abs_d_y\n\n    # define observations, states, and actions for the filter, use current and previous velocity measurement as action\n    # and ignore the 0th timestep because we don't have the previous velocity of that step\n    return {'o': data['rgbd'][:, 1:, :, :, :3],\n            's': data['pose'][:, 1:, :],\n            'a': np.concatenate([rel_d_x, rel_d_y, d_theta], axis=-1)}\n            # 'a': np.concatenate([data['vel'][:, :-1, None, :], data['vel'][:, 1:, None, :]], axis=-2)}\n\ndef compute_statistics(data):\n    means = dict()\n    stds = dict()\n    state_step_sizes = []\n    state_mins = []\n    state_maxs = []\n\n    for key in 'osa':\n        # compute means\n        axis = tuple(range(len(data[key].shape) - 1))  # means std by averaging over all but the last dimension\n        means[key] = np.mean(data[key], axis=axis, keepdims=True)\n\n        # compute stds\n        axis = tuple(range(len(data[key].shape) - 1))  # compute std by averaging over all but the last dimension\n        stds[key] = np.std(data[key] - means[key], axis=axis, keepdims=True)\n\n    # compute average step size in x, y, and theta for the distance metric\n    for i in range(5):\n        for j in range(len(data['seq_num'])):\n            if j == 0:\n                steps = np.reshape(data['s'][1:data['seq_num'][j], :, i] - data['s'][0:data['seq_num'][j]-1, :, i], [-1])\n            else:\n                steps = np.append(steps, np.reshape(data['s'][data['seq_num'][j-1]+1:data['seq_num'][j], :, i] - data['s'][data['seq_num'][j-1]:data['seq_num'][j]-1, :, i], [-1]))\n            if i == 2:\n                steps = wrap_angle(steps)\n        state_step_sizes.append(np.mean(abs(steps)))\n    state_step_sizes[0] = state_step_sizes[1] = (state_step_sizes[0] + state_step_sizes[1]) / 2\n    state_step_sizes = np.array(state_step_sizes)\n\n    # compute min and max in x, y and theta\n    for i in range(5):\n        state_mins.append(np.min(data['s'][:, :, i]))\n        state_maxs.append(np.max(data['s'][:, :, i]))\n    state_mins = np.array(state_mins)\n    state_maxs = np.array(state_maxs)\n\n    return means, stds, state_step_sizes, state_mins, state_maxs\n\n\ndef split_data(data, ratio=0.8, categories=['train', 'val']):\n    split_data = {categories[0]: dict(), categories[1]: dict()}\n    split_point_seq = math.floor(data['seq_num'].shape[0] * ratio)\n    split_point_data = data['seq_num'][split_point_seq-1]\n    for key in data.keys():\n        if key == 'seq_num':\n            split_data[categories[0]][key] = data[key][:split_point_seq]\n            split_data[categories[1]][key] = data[key][split_point_seq:] - data[key][split_point_seq-1]\n        elif key == 'weights':\n            split_data[categories[0]][key] = data[key][:split_point_seq]\n            split_data[categories[0]][key] = split_data[categories[0]][key]/np.sum(split_data[categories[0]][key])\n            split_data[categories[1]][key] = data[key][split_point_seq:]\n            split_data[categories[1]][key] = split_data[categories[1]][key]/np.sum(split_data[categories[1]][key])\n        else:\n            split_data[categories[0]][key] = data[key][:split_point_data]\n            split_data[categories[1]][key] = data[key][split_point_data:]\n    for key in split_data.keys():\n        print('SPLIT --> {}: {}'.format(key, len(split_data[key]['seq_num'])))\n    return split_data\n\n\ndef reduce_data(data, num_episodes):\n    new_data = dict()\n    for key in 'osa':\n        new_data[key] = data[key][:num_episodes]\n    return new_data\n\ndef shuffle_data(data):\n    new_data = dict()\n    shuffled_indices = np.random.permutation(len(data['o']))\n    for key in 'osa':\n        new_data[key] = data[key][shuffled_indices]\n    return new_data\n\ndef remove_state(data, provide_initial_state=False):\n    new_data = dict()\n    new_data['o'] = data['o']\n    new_data['a'] = data['a']\n    if provide_initial_state:\n        new_data['s'] = data['s'][..., :1, :]\n    return new_data\n\n\ndef noisify_data_condition(data, condition):\n    print('condition', condition)\n    if condition == 'odom0_imgTG':\n        return noisyfy_data(data, odom_noise_factor=0.0)\n    elif condition == 'odom5_imgTG':\n        return noisyfy_data(data, odom_noise_factor=0.5)\n    elif condition == 'odom10_imgTG':\n        return noisyfy_data(data)\n    elif condition == 'odom20_imgTG':\n        return noisyfy_data(data, odom_noise_factor=2.0)\n    elif condition == 'odomX_imgTG':\n        data = noisyfy_data(data, odom_noise_factor=0.0)\n        # shuffle actions to basically make them meaningless\n        shape = data['a'].shape\n        a = np.reshape(data['a'], [-1, shape[-1]])\n        np.random.shuffle(a)\n        data['a'] = np.reshape(a, shape)\n        return data\n    elif condition == 'odom10_imgC':\n        return noisyfy_data(data, img_noise_factor=0.0, img_random_shift=False)\n    elif condition == 'odom10_imgG':\n        return noisyfy_data(data, img_noise_factor=1.0, img_random_shift=False)\n    elif condition == 'odom10_imgT':\n        return noisyfy_data(data, img_noise_factor=0.0, img_random_shift=True)\n    elif condition == 'odom10_imgX':\n        data = noisyfy_data(data, img_noise_factor=0.0, img_random_shift=False)\n        # shuffle observations to basically make them meaningless\n        shape = data['o'].shape\n        o = np.reshape(data['o'], [-1, shape[-1]])\n        np.random.shuffle(o)\n        data['o'] = np.reshape(o, shape)\n        return data\n\ndef noisyfy_data(data, odom_noise_factor=1.0, img_noise_factor=1.0, img_random_shift=True):\n    print(\"noisyfying data ... \")\n    new_data = dict()\n    new_data['s'] = data['s']\n    new_data['a'] = data['a'] * np.random.normal(1.0, 0.1 * odom_noise_factor, data['a'].shape)\n    new_data['seq_num'] = data['seq_num']\n    new_data['o'] = data['o']\n    return new_data\n\ndef make_batch_iterator(data, batch_size=32, seq_len=10, use_mirrored_data=True):\n\n    while True:\n        o = np.zeros((batch_size, seq_len, 50, 150, 6))\n        a = np.zeros((batch_size, seq_len, 3))\n        s = np.zeros((batch_size, seq_len, 5))\n        for ii in range(batch_size):\n            trajectory = np.random.choice(len(data['seq_num']), p = data['weights'])\n\n            start = 0 if trajectory == 0 else data['seq_num'][trajectory-1]\n            start_steps = np.random.random_integers(start, data['seq_num'][trajectory] - seq_len - 1)\n            key_append = '-m' if use_mirrored_data and ii >= batch_size / 2 else ''\n            o[ii, :, :, :, :] = data['o'+key_append][start_steps:start_steps + seq_len, 0]\n            a[ii, :, :] = data['a'][start_steps:start_steps + seq_len, 0]\n            s[ii, :, :] = data['s'+key_append][start_steps:start_steps + seq_len, 0]\n\n        batches =  {'o': o, 'a': a, 's': s}\n        yield batches\n\ndef make_repeating_batch_iterator(data, epoch_len, batch_size=32, seq_len=10, use_mirrored_data=True):\n\n    o = np.zeros((batch_size, seq_len, 50, 150, 6))\n    a = np.zeros((batch_size, seq_len, 3))\n    s = np.zeros((batch_size, seq_len, 5))\n    start_steps = np.zeros((epoch_len, batch_size))\n    trajectory = np.random.random_integers(0, len(data['seq_num']) - 1, size=[epoch_len, batch_size])\n    for kk in range(epoch_len):\n        for ii in range(batch_size):\n            start = 0 if trajectory[kk, ii] == 0 else data['seq_num'][trajectory[kk, ii] - 1]\n            start_steps[kk, ii] = np.random.random_integers(start, data['seq_num'][trajectory[kk, ii]] - seq_len - 1)\n\n    start_steps = start_steps.astype(int)\n\n    while True:\n        for kk in range(epoch_len):\n            for ii in range(batch_size):\n                ssteps = start_steps[kk, ii]\n                key_append = '-m' if use_mirrored_data and ii >= batch_size / 2 else ''\n                o[ii, :, :, :, :] = data['o' + key_append][ssteps:ssteps + seq_len, 0]\n                a[ii, :, :] = data['a'][ssteps:ssteps + seq_len, 0]\n                s[ii, :, :] = data['s' + key_append][ssteps:ssteps + seq_len, 0]\n            batches =  {'o': o, 'a': a, 's': s}\n            yield batches\n\n\ndef make_complete_batch_iterator(data, batch_size=1000, seq_len=10):\n    num_episodes = len(data['s'])\n    num_start_steps = len(data['s'][0]) - seq_len\n    batch_indices = [(i, j) for i in range(num_episodes) for j in range(num_start_steps)]\n    while batch_indices != []:\n        batches = {k: np.concatenate([data[k][i:i + 1, j:j + seq_len] for (i, j) in batch_indices[:batch_size]]) for k in data.keys}\n        batch_indices = batch_indices[batch_size:]\n        yield batches\n\n\ndef make_batch_iterator_for_evaluation(data, start_step, trajectory, batch_size = 1, seq_len=10):\n    while True:\n        o = np.zeros((batch_size,seq_len, 50, 150, 6))\n        a = np.zeros((batch_size, seq_len, 3))\n        s = np.zeros((batch_size, seq_len, 5))\n        for ii in range(batch_size):\n\n            # shift start step to where the sequence begins\n            if trajectory != 0:\n                start_step = data['seq_num'][trajectory-1] + start_step\n\n            for jj in range(seq_len):\n                o[ii, jj, :, :, :] = data['o'][start_step+jj, :, :, :]\n                a[ii, jj, :] = data['a'][start_step+jj, :]\n                s[ii, jj, :] = data['s'][start_step+jj, :]\n\n        batches =  {'o': o, 'a': a, 's': s}\n        yield batches\n\ndef plot_observation_check(data, means, stds):\n\n    observations = data['o']\n    plt.ion()\n    for o in observations:\n        # shape(o): (1, 50, 150, 6)\n        # shape(means['o']) = (1, 1, 50, 150, 6)\n\n        norm_o = (o - means['o'][0]) / stds['o'][0]\n\n        for d in range(o.shape[-1]):\n            plt.figure(d)\n            plt.clf()\n\n            plt.imshow(norm_o[0, :, :, d], interpolation='nearest', cmap='coolwarm', vmin=-3, vmax=3)\n            print('dimension {}: ({}-{})'.format(d, np.min(o[:, :, d]), np.max(o[:, :, d])))\n\n        for d in range(2):\n            plt.figure(10 + d)\n            plt.clf()\n            if d == 0:\n                plt.imshow(np.clip(o[0, :, :, 3*d:3*(d+1)]/255.0, 0.0, 1.0), interpolation='nearest')\n            else:\n                plt.imshow(o[0, :, :, 3*d:3*(d+1)]/255.0/2 + 0.5, interpolation='nearest')\n\n        d = 2\n        plt.figure('means')\n        plt.clf()\n        plt.imshow(means['o'][0, 0, :, :, d], interpolation='nearest', cmap='coolwarm', vmin=0, vmax=255)\n\n        plt.figure('stds')\n        plt.clf()\n        plt.imshow(stds['o'][0, 0, :, :, d], interpolation='nearest', cmap='coolwarm', vmin=0, vmax=255)\n\n        plt.pause(10)\n\ndef plot_video(data):\n    observations = data['o']\n    plt.ion()\n    for i, o in enumerate(observations):\n        # shape(o): (1, 50, 150, 6)\n        # shape(means['o']) = (1, 1, 50, 150, 6)\n\n        d = 0\n        plt.figure(10 + d)\n        plt.clf()\n        plt.imshow(np.clip(o[0, :, :, 3 * d:3 * (d + 1)] / 255.0, 0.0, 1.0), interpolation='nearest')\n\n        plt.pause(0.05)\n        print(i)\n\ndef plot_sequences(data, means, stds, state_step_sizes):\n\n    for k, i in enumerate(data['seq_num']):\n        print(i)\n        if k < len(data['seq_num'])-1:\n            plt.figure(k)\n            for j in range(4):\n                plt.subplot(4,1,j+1)\n                plt.imshow(np.clip(data['o'][i-2+j, 0, :, :, :3]/255.0, 0.0, 1.0), interpolation='nearest')\n                plt.xticks([])\n                plt.yticks([])\n                plt.ylabel(i-2+j)\n\n    plt.figure('trajectories')\n    last_seq_num = 0\n    for k, i in enumerate(data['seq_num']):\n        if k % 2 == 0:\n            plt.plot(data['s'][last_seq_num:i, 0, 0], data['s'][last_seq_num:i, 0, 1], label=\"trajectory {}/{}\".format(k, k+1))\n            plt.quiver(data['s'][last_seq_num:i, 0, 0], data['s'][last_seq_num:i, 0, 1],\n                       np.cos(data['s'][last_seq_num:i, 0, 2]), np.sin(data['s'][last_seq_num:i, 0, 2]), color='k')\n        last_seq_num = i\n\n    plt.gca().set_aspect('equal')\n    plt.legend()\n\n    plt.figure('normalized state')\n    norm_states = (data['s'] - means['s'][0]) / stds['s'][0]\n    for d in range(data['s'].shape[-1]):\n        plt.plot(norm_states[:, 0, d], label='state dim {}'.format(d))\n\n    plt.figure('state')\n    for d in range(data['s'].shape[-1]):\n        plt.plot(data['s'][:, 0, d], label='state dim {}'.format(d))\n\n    plt.figure('scaled state')\n    for d in range(data['s'].shape[-1]):\n        plt.plot(data['s'][:, 0, d] / state_step_sizes[d], label='state dim {}'.format(d))\n        print('dim {}: state step size: {}'.format(d, state_step_sizes[d]))\n    for i in data['seq_num']:\n        plt.plot([i, i], [-3, 3], 'k')\n\n    plt.legend()\n    plt.show()\n\n\n\ndef add_mirrored_data(data):\n    data['o-m'] = data['o'][..., ::-1, :]\n    data['s-m'] = np.concatenate([ data['s'][..., 0:1], # keep x\n                                         -data['s'][..., 1:2], # invert y\n                                         -data['s'][..., 2:3], # invert angle\n                                          data['s'][..., 3:4], # keep foward vel\n                                         -data['s'][..., 4:5], # invert angular vel\n                                         ], axis=-1)\n    return data\n\nif __name__ == '__main__':\n\n    data = load_kitti_sequences()\n\n    print(data['o'].shape)\n    means, stds, state_step_sizes, state_mins, state_maxs = compute_statistics(data)\n    print(data['o'].shape)\n\n\n"
  },
  {
    "path": "utils/exp_utils.py",
    "content": "import os\nimport itertools\nimport numpy as np\n\ndef get_default_hyperparams():\n\n    # default hyperparams\n    return {\n        'global': {\n            'init_with_true_state': False,\n            'learn_odom': False,\n            'use_proposer': True,\n            'propose_ratio': 0.7,\n            'proposer_keep_ratio': 0.15,\n            'min_obs_likelihood': 0.004,\n        },\n        'train': {\n            'train_individually': True,\n            'train_e2e': True,\n            'split_ratio': 0.9,\n            'seq_len': 20,\n            'batch_size': 32,\n            'epoch_length': 50,\n            'num_epochs': 10000,\n            'patience': 200,\n            'learning_rate': 0.0003,\n            'dropout_keep_ratio': 0.3,\n            'num_particles': 100,\n            'particle_std': 0.2,\n        },\n        'test' : {\n            'num_particles': 1000,\n        }\n    }\n\ndef exp_variables_to_name(x):\n    return '_'.join(map(str, x))\n\n\ndef sample_exp_variables(path, exp_variables):\n\n    # compute all combinations of the experiment variables\n    product = list(itertools.product(*exp_variables))\n    n = len(product)\n    # turn them into filename endings\n    file_endings = list(map(exp_variables_to_name, product))\n    # count how often each ending appears, i.e. how often each experiment has been run\n    counts = [0] * n\n    try:\n        for filename in os.listdir(path):\n            if os.path.isfile(os.path.join(path, filename)):\n                for i in range(n):\n                    if filename.endswith(file_endings[i]):\n                        counts[i] += 1\n    except FileNotFoundError:\n        pass\n\n    # compute a sample list with samples according to which experimental variables need more examples\n    min_count = np.min(counts)\n    sample_list = []\n    for i in range(n):\n        sample_list += [product[i]] * max(0, (min_count + 2) - counts[i])\n    if sample_list == []:\n        sample_list = product\n\n    # sample from this list\n    print('sampling from:', sample_list)\n    sample = sample_list[np.random.choice(len(sample_list))]\n    print('--> ', sample)\n    return sample, min_count\n\n\ndef print_msg_and_dict(msg, d):\n    keys = sorted(list(d.keys()))\n    msg += ' '\n    for k in keys:\n        msg += '\\n{}: {}'.format(k, d[k])\n    print('########################################################')\n    print(msg)\n    print('########################################################')\n\n\ndef add_to_log(log, d):\n    for k in d.keys():\n        if k not in log.keys():\n            log[k] = []\n        log[k].append(d[k])\n    return log\n"
  },
  {
    "path": "utils/exp_utils_kitti.py",
    "content": "def get_default_hyperparams():\n\n    # default hyperparams\n    return {\n        'global': {\n            'init_with_true_state': True,\n            'learn_odom': False,\n            'use_proposer': False,\n            'propose_ratio': 0.7,\n            'proposer_keep_ratio': 0.15,\n            'min_obs_likelihood': 0.004,\n            'learn_gaussian_mle': False\n\n        },\n        'train': {\n            'train_individually': True,\n            'train_e2e': True,\n            'split_ratio': 0.95,\n            'seq_len': 50,\n            'batch_size': 32,\n            'epoch_length': 50,\n            'num_epochs': 10000,\n            'patience': 200,\n            'learning_rate': 0.0003,\n            'dropout_keep_ratio': 0.3,\n            'num_particles': 100,\n            'particle_std': 2.0,\n            'learn_gaussian_mle': False\n        },\n        'test' : {\n            'num_particles': 400\n        }\n    }"
  },
  {
    "path": "utils/method_utils.py",
    "content": "import numpy as np\nimport tensorflow as tf\n\nfrom utils.data_utils_kitti import wrap_angle\n\n\ndef compute_sq_distance(a, b, state_step_sizes):\n    result = 0.0\n    for i in range(a.shape[-1]):\n        # compute difference\n        diff = a[..., i] - b[..., i]\n        # wrap angle for theta\n        if i == 2:\n            diff = wrap_angle(diff)\n        # add up scaled squared distance\n        result += (diff / state_step_sizes[i]) ** 2\n    return result\n\n\ndef atan2(x, y, epsilon=1.0e-12):\n    \"\"\"\n    A hack until the tf developers implement a function that can find the angle from an x and y co-\n    ordinate.\n    :param x:\n    :param epsilon:\n    :return:\n    \"\"\"\n    # Add a small number to all zeros, to avoid division by zero:\n    x = tf.where(tf.equal(x, 0.0), x + epsilon, x)\n    y = tf.where(tf.equal(y, 0.0), y + epsilon, y)\n\n    angle = tf.where(tf.greater(x, 0.0), tf.atan(y / x), tf.zeros_like(x))\n    angle = tf.where(tf.logical_and(tf.less(x, 0.0), tf.greater_equal(y, 0.0)), tf.atan(y / x) + np.pi, angle)\n    angle = tf.where(tf.logical_and(tf.less(x, 0.0), tf.less(y, 0.0)), tf.atan(y / x) - np.pi, angle)\n    angle = tf.where(tf.logical_and(tf.equal(x, 0.0), tf.greater(y, 0.0)), 0.5 * np.pi * tf.ones_like(x), angle)\n    angle = tf.where(tf.logical_and(tf.equal(x, 0.0), tf.less(y, 0.0)), -0.5 * np.pi * tf.ones_like(x), angle)\n    angle = tf.where(tf.logical_and(tf.equal(x, 0.0), tf.equal(y, 0.0)), tf.zeros_like(x), angle)\n    return angle\n"
  },
  {
    "path": "utils/plotting_utils.py",
    "content": "import matplotlib.pyplot as plt\nfrom matplotlib.patches import Rectangle\nimport numpy as np\n\nhead_scale = 3.0  # 1.5\n# head_scale = 2.0  # 1.5\nquiv_kwargs = {'scale_units':'xy', 'scale':1./80., 'width': 0.003, 'headlength': 5*head_scale, 'headwidth': 3*head_scale, 'headaxislength': 4.5*head_scale}\n\n\ndef show_pause(show=False, pause=0.0):\n    '''Shows a plot by either blocking permanently using show or temporarily using pause.'''\n    if show:\n        plt.ioff()\n        plt.show()\n    elif pause:\n        plt.ion()\n        plt.pause(pause)\n\n\ndef plot_maze(maze='nav01', margin=1, means=None, stds=None, figure_name=None, show=False, pause=False, ax=None, linewidth=1.0):\n    if ax is None:\n        ax = plt.gca()\n    if figure_name is not None:\n        plt.figure(figure_name)\n\n    if 'nav01' in maze:\n        walls = np.array([\n            # horizontal\n            [[0, 500], [1000, 500]],\n            [[400, 400], [500, 400]],\n            [[600, 400], [700, 400]],\n            [[800, 400], [1000, 400]],\n            [[200, 300], [400, 300]],\n            [[100, 200], [200, 200]],\n            [[400, 200], [700, 200]],\n            [[200, 100], [300, 100]],\n            [[600, 100], [900, 100]],\n            [[0, 0], [1000, 0]],\n            # vertical\n            [[0, 0], [0, 500]],\n            [[100, 100], [100, 200]],\n            [[100, 300], [100, 500]],\n            [[200, 200], [200, 400]],\n            [[200, 0], [200, 100]],\n            [[300, 100], [300, 200]],\n            [[300, 400], [300, 500]],\n            [[400, 100], [400, 400]],\n            [[500, 0], [500, 200]],\n            [[600, 100], [600, 200]],\n            [[700, 200], [700, 300]],\n            [[800, 200], [800, 400]],\n            [[900, 100], [900, 300]],\n            [[1000, 0], [1000, 500]],\n        ])\n        rooms = [\n            # [[400, 200], 300, 200]\n            ]\n        ax.set_xlim([-margin, 1000+margin])\n        ax.set_ylim([-margin, 500+margin])\n\n    if 'nav02' in maze:\n        walls = np.array([\n            # horizontal\n            [[0, 900], [1500, 900]],\n            [[100, 800], [400, 800]],\n            [[500, 800], [600, 800]],\n            [[800, 800], [1000, 800]],\n            [[1100, 800], [1200, 800]],\n            [[1300, 800], [1400, 800]],\n            [[100, 700], [600, 700]],\n            [[700, 700], [800, 700]],\n            [[1000, 700], [1100, 700]],\n            [[1200, 700], [1400, 700]],\n            [[900, 600], [1200, 600]],\n            [[1300, 600], [1500, 600]],\n            [[0, 500], [100, 500]],\n            [[1300, 500], [1400, 500]],\n            [[100, 400], [200, 400]],\n            [[1200, 400], [1400, 400]],\n            [[300, 300], [800, 300]],\n            [[900, 300], [1200, 300]],\n            [[400, 200], [600, 200]],\n            [[700, 200], [800, 200]],\n            [[1200, 200], [1500, 200]],\n            [[200, 100], [300, 100]],\n            [[500, 100], [700, 100]],\n            [[800, 100], [900, 100]],\n            [[1100, 100], [1400, 100]],\n            [[0, 0], [1500, 0]],\n            # vertical\n            [[0, 0], [0, 900]],\n            [[100, 0], [100, 300]],\n            [[100, 500], [100, 600]],\n            [[100, 700], [100, 800]],\n            [[200, 100], [200, 200]],\n            [[200, 300], [200, 400]],\n            [[200, 500], [200, 700]],\n            [[300, 100], [300, 300]],\n            [[400, 0], [400, 200]],\n            [[500, 800], [500, 900]],\n            [[700, 100], [700, 200]],\n            [[700, 700], [700, 800]],\n            [[800, 200], [800, 800]],\n            [[900, 100], [900, 700]],\n            [[1000, 0], [1000, 200]],\n            [[1000, 700], [1000, 800]],\n            [[1100, 700], [1100, 800]],\n            [[1100, 100], [1100, 300]],\n            [[1200, 800], [1200, 900]],\n            [[1200, 400], [1200, 700]],\n            [[1300, 200], [1300, 300]],\n            [[1300, 500], [1300, 600]],\n            [[1400, 300], [1400, 500]],\n            [[1400, 700], [1400, 800]],\n            [[1500, 0], [1500, 900]],\n        ])\n        rooms = [\n            # [[900, 300], 300, 300]\n            ]\n        ax.set_xlim([-margin, 1500+margin])\n        ax.set_ylim([-margin, 900+margin])\n\n    if 'nav03' in maze:\n        walls = np.array([\n            # horizontal\n            [[0, 1300], [2000, 1300]],\n            [[100, 1200], [500, 1200]],\n            [[600, 1200], [1400, 1200]],\n            [[1600, 1200], [1700, 1200]],\n            [[0, 1100], [600, 1100]],\n            [[1500, 1100], [1600, 1100]],\n            [[1600, 1000], [1800, 1000]],\n            [[800, 1000], [900, 1000]],\n            [[100, 1000], [200, 1000]],\n            [[700, 900], [800, 900]],\n            [[1600, 900], [1800, 900]],\n            [[200, 800], [300, 800]],\n            [[800, 800], [1200, 800]],\n            [[1300, 800], [1500, 800]],\n            [[1600, 800], [1900, 800]],\n            [[900, 700], [1400, 700]],\n            [[1500, 700], [1600, 700]],\n            [[1700, 700], [1900, 700]],\n            [[700, 600], [800, 600]],\n            [[1400, 600], [1500, 600]],\n            [[1600, 600], [1700, 600]],\n            [[100, 500], [200, 500]],\n            [[300, 500], [500, 500]],\n            [[600, 500], [700, 500]],\n            [[1400, 500], [1900, 500]],\n            [[100, 400], [200, 400]],\n            [[400, 400], [600, 400]],\n            [[1500, 400], [1600, 400]],\n            [[1700, 400], [1800, 400]],\n            [[200, 300], [300, 300]],\n            [[400, 300], [500, 300]],\n            [[600, 300], [800, 300]],\n            [[900, 300], [1100, 300]],\n            [[1300, 300], [1500, 300]],\n            [[1600, 300], [1700, 300]],\n            [[100, 200], [200, 200]],\n            [[500, 200], [600, 200]],\n            [[800, 200], [1100, 200]],\n            [[1200, 200], [1400, 200]],\n            [[1500, 200], [1600, 200]],\n            [[200, 100], [300, 100]],\n            [[500, 100], [800, 100]],\n            [[1000, 100], [1200, 100]],\n            [[1400, 100], [1600, 100]],\n            [[1800, 100], [1900, 100]],\n            [[0, 0], [2000, 0]],\n            # vertical\n            [[0, 0], [0, 1300]],\n            [[100, 0], [100, 300]],\n            [[100, 400], [100, 1000]],\n            [[200, 300], [200, 400]],\n            [[200, 600], [200, 800]],\n            [[200, 900], [200, 1000]],\n            [[300, 100], [300, 600]],\n            [[300, 800], [300, 1100]],\n            [[400, 0], [400, 300]],\n            [[400, 1200], [400, 1300]],\n            [[500, 100], [500, 200]],\n            [[600, 200], [600, 400]],\n            [[600, 1100], [600, 1200]],\n            [[700, 200], [700, 300]],\n            [[700, 400], [700, 1100]],\n            [[800, 100], [800, 200]],\n            [[800, 300], [800, 500]],\n            [[800, 600], [800, 700]],\n            [[800, 1000], [800, 1100]],\n            [[900, 0], [900, 100]],\n            [[900, 300], [900, 600]],\n            [[900, 900], [900, 1200]],\n            [[1000, 100], [1000, 200]],\n            [[1200, 100], [1200, 200]],\n            [[1300, 0], [1300, 100]],\n            [[1400, 100], [1400, 700]],\n            [[1500, 700], [1500, 1000]],\n            [[1500, 1100], [1500, 1200]],\n            [[1600, 200], [1600, 400]],\n            [[1600, 600], [1600, 700]],\n            [[1600, 1000], [1600, 1100]],\n            [[1600, 1200], [1600, 1300]],\n            [[1700, 1100], [1700, 1200]],\n            [[1700, 700], [1700, 800]],\n            [[1700, 500], [1700, 600]],\n            [[1700, 0], [1700, 300]],\n            [[1800, 100], [1800, 400]],\n            [[1800, 600], [1800, 700]],\n            [[1800, 900], [1800, 1200]],\n            [[1900, 800], [1900, 1300]],\n            [[1900, 100], [1900, 600]],\n            [[2000, 0], [2000, 1300]],\n        ])\n        rooms = [\n                # [[300, 500], 400, 600],\n                #  [[900, 800], 600, 400],\n                #  [[900, 300], 500, 400],\n                 ]\n        ax.set_xlim([-margin, 2000 + margin])\n        ax.set_ylim([-margin, 1300 + margin])\n\n    if means is not None:\n        walls -= means['pose'][:, :, :2]\n    if stds is not None:\n        walls /= stds['pose'][:, :, :2]\n    # color = (0.8, 0.8, 0.8)\n    color = (0, 0, 0)\n\n    ax.plot(walls[:, :, 0].T, walls[:, :, 1].T, color=color, linewidth=linewidth)\n    for room in rooms:\n        ax.add_patch(Rectangle(*room, facecolor=(0.85, 0.85, 0.85), linewidth=0))\n    ax.set_aspect('equal')\n\n\ndef plot_trajectories(data, figure_name=None, show=False, pause=False, emphasize=None, odom=False, mincolor=0.0, linewidth=0.3):\n    from methods.odom import OdometryBaseline\n    if figure_name is not None:\n        plt.figure(figure_name)\n    for i, trajectories in enumerate(data['s']):\n        color = np.random.uniform(low=mincolor, high=1.0, size=3)\n        plt.plot(trajectories[:, 0], trajectories[:, 1], color=color, linewidth=linewidth, zorder=0)\n    if emphasize is not None:\n        true_traj = data['s'][emphasize, :20, :]\n        odom = OdometryBaseline()\n        odom_traj = odom.predict(None, {k:data[k][emphasize:emphasize+1, :20] for k in data.keys()})[0]\n        print(true_traj)\n        print(odom_traj)\n\n        traj = odom_traj\n        plt.plot(traj[:, 0], traj[:, 1], '--', color=[0.0, 0.0, 1.0], linewidth=0.8, zorder=0)\n        # plt.plot(traj[:, 0], traj[:, 1], 'o', markerfacecolor='None',\n        #     markeredgecolor=[0.0, 0.0, 0.0],\n        #     markersize=5)\n        # plt.quiver(traj[:, 0], traj[:, 1], np.cos(traj[:, 2]), np.sin(traj[:, 2]),\n        #            color=[0.0, 0.0, 0.0], zorder=1, headlength=0, headaxislength=0, scale=10, width=0.02, units='inches', scale_units='inches')\n\n        traj = true_traj\n        plt.plot(traj[:, 0], traj[:, 1], '-', color=[0.0, 0.0, 1.0], linewidth=0.8, zorder=0)\n        plt.plot(traj[:, 0], traj[:, 1], 'o', markerfacecolor='None',\n            markeredgecolor=[0.0, 0.0, 0.0],\n            markersize=5)\n        plt.quiver(traj[:, 0], traj[:, 1], np.cos(traj[:, 2]), np.sin(traj[:, 2]),\n                   color=[0.0, 0.0, 0.0], zorder=1, headlength=0, headaxislength=0, scale=10, width=0.02, units='inches', scale_units='inches')\n\n    plt.gca().set_aspect('equal')\n    show_pause(show, pause)\n\n\ndef plot_trajectory(data, figure_name=None, show=False, pause=False, emphasize=None, odom=False, mincolor=0.0, linewidth=0.3):\n    from methods.odom import OdometryBaseline\n    if figure_name is not None:\n        plt.figure(figure_name)\n    for i, trajectories in enumerate(data['s']):\n        plt.plot(trajectories[:, 0], trajectories[:, 1], '-', color='red', linewidth=linewidth, zorder=0, markersize=4)\n        plt.plot(trajectories[:5, 0], trajectories[:5, 1], '.', color='blue', linewidth=linewidth, zorder=0, markersize=8)\n        plt.plot(trajectories[0, 0], trajectories[0, 1], '.', color='blue', linewidth=linewidth, zorder=0, markersize=16)\n\n        # plt.quiver(trajectories[:5, 0], trajectories[:5, 1],\n        #        np.cos(trajectories[:5, 2]), np.sin(trajectories[:5, 2]),\n        #            # np.arange(len(trajectories)), cmap='viridis', alpha=1.0,\n        #            color='red', alpha=1.0,\n        #        **quiv_kwargs\n        #        )\n\n    plt.gca().set_aspect('equal')\n    show_pause(show, pause)\n\n\ndef plot_observations(data, n=20, figure_name=None, show=False, pause=False):\n\n    plt.figure(figsize=(10,2.5))\n    for i in range(n):\n        # plt.figure('Normalized image')\n        # plt.gca().clear()\n        # plt.imshow(0.5 + rgbds[i, :, :, :3]/10, interpolation='nearest')\n        # plt.pause(0.001)\n        #\n        # plt.figure('Depth image')\n        # plt.gca().clear()\n        # plt.imshow(0.5 + rgbds[i, :, :, 3] / 10, interpolation='nearest', cmap='coolwarm', vmin=0.0, vmax=1.0)\n        # plt.pause(0.001)\n\n\n        # plt.gca().clear()\n        # plt.subplot(2, 10, i+1)\n        plt.subplot(1, n, i+1)\n        plt.imshow(np.clip(data['o'][0, i, :, :, :]/255.0, 0.0, 1.0), interpolation='nearest')\n        plt.axis('off')\n        # plt.tight_layout(pad=0.1)\n        # plt.pause(0.1)\n    show_pause(show, pause)\n\n\ndef view_data(data):\n    # overview plot\n    for poses in data['s']:\n        plt.figure('Overview')\n        plt.plot(poses[:, 0], poses[:, 1])\n\n        # # sample plot\n        # for poses, velocities, rgbds in zip(data['pose'], data['vel'], data['rgbd']):\n        #     # for poses in data['pose']:\n        #     plt.ioff()\n        #     plt.figure('Sample')\n        #     # plt.plot(poses[:, 0], 'r-')\n        #     # plt.plot(poses[:, 1], 'g-')\n        #     plt.plot(poses[:, 2], 'b-')\n        #     # plt.plot(velocities[:, 0], 'r--')\n        #     # plt.plot(velocities[:, 1], 'g--')\n        #     plt.plot(velocities[:, 2], 'b--')\n        #     plt.show()\n        #\n        #     # for i in range(100):\n        #     #     plt.figure('Normalized image')\n        #     #     plt.gca().clear()\n        #     #     plt.imshow(0.5 + rgbds[i, :, :, :3]/10, interpolation='nearest')\n        #     #     plt.pause(0.001)\n        #     #\n        #     #     plt.figure('Depth image')\n        #     #     plt.gca().clear()\n        #     #     plt.imshow(0.5 + rgbds[i, :, :, 3] / 10, interpolation='nearest', cmap='coolwarm', vmin=0.0, vmax=1.0)\n        #     #     plt.pause(0.001)\n        #     #\n        #     #     plt.figure('Real image')\n        #     #     plt.gca().clear()\n        #     #     plt.imshow((rgbds*stds['rgbd'][0] + means['rgbd'][0])[i, :, :, :3]/255.0, interpolation='nearest')\n        #     #     plt.pause(0.1)\n"
  }
]