Merge branch 'main' of github.com:lasgroup/simulation_transfer

.gitignore

@@ -168,3 +168,4 @@ Temporary Items
 
 # exlcluded files
 plotting_scripts/plots
+/experiments/1d_visualization/plot_dicts/

experiments/bimodal_exp/bimodal_exp.py

@@ -1,16 +1,20 @@
+import argparse
+
 import jax
 import jax.numpy as jnp
+import jax.random as jr
+import wandb
 from jax import vmap
 from jax.lax import cond
-import argparse
-
 
-import wandb
 from sim_transfer.models.bnn_fsvgd_sim_prior import BNN_FSVGD_SimPrior
 from sim_transfer.sims import LinearBimodalSim
 
 
-def experiment(score_estimator: str, seed: int, project_name: str):
+def experiment(score_estimator: str,
+               seed: int,
+               data_seed: int,
+               project_name: str):
     def key_iter():
         key = jax.random.PRNGKey(seed)
         while True:
@@ -36,7 +40,7 @@ def negative(x):
     domain = sim.domain
 
     num_train_points = 3
-    x_train = jax.random.uniform(key=next(key_iter), shape=(num_train_points,),
+    x_train = jax.random.uniform(key=jr.PRNGKey(data_seed), shape=(num_train_points,),
                                  minval=domain.l, maxval=jnp.array([0.0])).reshape(-1, 1)
     y_train = v_fun(x_train)
 
@@ -62,18 +66,38 @@ def negative(x):
     for i in range(10):
         bnn.fit(x_train, y_train, x_eval=x_test, y_eval=y_test, num_steps=2000, log_to_wandb=True)
         if NUM_DIM_X == 1:
-            bnn.plot_1d(x_train, y_train, true_fun=v_fun,
+            bnn.plot_1d(x_train, y_train,
+                        true_fun=v_fun,
                         title=f'FSVGD SimPrior {score_estimator}, iter {(i + 1) * 2000}',
-                        domain_l=domain.l[0], domain_u=domain.u[0], log_to_wandb=True)
+                        domain_l=domain.l[0],
+                        domain_u=domain.u[0],
+                        plot_posterior_samples=True,
+                        log_to_wandb=True,
+                        )
+
+    # Final plot, save data as well
+    bnn.plot_1d(x_train, y_train,
+                true_fun=v_fun,
+                title=f'FSVGD SimPrior {score_estimator}, final plot',
+                domain_l=domain.l[0],
+                domain_u=domain.u[0],
+                plot_posterior_samples=True,
+                log_to_wandb=True,
+                save_plot_dict=True,
+                )
 
 
 def main(args):
-    experiment(args.score_estimator, args.seed, args.project_name)
+    experiment(score_estimator=args.score_estimator,
+               seed=args.seed,
+               data_seed=args.data_seed,
+               project_name=args.project_name)
 
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--data_seed', type=int, default=0)
     parser.add_argument('--score_estimator', type=str, default='gp')
     parser.add_argument('--project_name', type=str, default='LinearBimodal')
     args = parser.parse_args()

experiments/bimodal_exp/bimodal_launcher.py

@@ -1,29 +1,33 @@
 import bimodal_exp
 from experiments.util import generate_run_commands, generate_base_command
 
-PROJECT_NAME = 'LinearBimodalSimPrior'
+PROJECT_NAME = 'LinearBimodalSimPrior_19Jan_N1'
 
 applicable_configs = {
     'score_estimator': ['gp', 'ssge', 'nu-method', 'kde'],
-    'seed': [i for i in range(10)],
+    'seed': [i for i in range(5)],
+    'data_seed': [0]
 }
 
 
 def main():
     command_list = []
     for score_estimator in applicable_configs['score_estimator']:
         for seed in applicable_configs['seed']:
-            flags = {
-                'seed': seed,
-                'score_estimator': score_estimator,
-                'project_name': PROJECT_NAME
-            }
+            for data_seed in applicable_configs['data_seed']:
+                flags = {
+                    'data_seed': data_seed,
+                    'seed': seed,
+                    'score_estimator': score_estimator,
+                    'project_name': PROJECT_NAME
+                }
 
-            cmd = generate_base_command(bimodal_exp, flags=flags)
-            command_list.append(cmd)
+                cmd = generate_base_command(bimodal_exp, flags=flags)
+                command_list.append(cmd)
 
     # submit jobs
-    generate_run_commands(command_list, num_cpus=1, num_gpus=1, mode='euler', duration='3:59:00', prompt=True, mem=16000)
+    generate_run_commands(command_list, num_cpus=1, num_gpus=1, mode='euler', duration='3:59:00', prompt=True,
+                          mem=16000)
 
 
 if __name__ == '__main__':

experiments/data_provider.py

@@ -191,7 +191,7 @@ def get_rccar_recorded_data_new(encode_angle: bool = True, skip_first_n_points:
                                 action_delay: int = 3, action_stacking: bool = False,
                                 car_id: int = 2):
 
-    assert car_id in [1, 2]
+    assert car_id in [1, 2, 3]
     if car_id == 1:
         assert dataset in ['all', 'v1']
         recordings_dir = [os.path.join(DATA_DIR, 'recordings_rc_car_v1')]

experiments/meta_learning_exp/sweep_regression_exp_num_data.py

@@ -0,0 +1,107 @@
+from experiments.util import (generate_run_commands, generate_base_command, RESULT_DIR, sample_param_flags, hash_dict)
+from experiments.data_provider import DATASET_CONFIGS
+
+import experiments.meta_learning_exp.run_meta_learning_exp
+import numpy as np
+import datetime
+import itertools
+import argparse
+import os
+
+MODEL_SPECIFIC_CONFIG = {
+    'PACOH': {
+        'prior_weight': {'values': [0.1]},
+        'hyper_prior_weight': {'values': [1e-3]},
+        'num_iter_meta_train': {'values': [100_000]},
+        'meta_batch_size': {'values': [4]},
+        'bandwidth': {'values': [10.]},
+        'lr': {'values': [5e-4]},
+    },
+    'NP': {
+        'num_iter_meta_train': {'values': [100_000]},
+        'latent_dim': {'values': [256]},
+        'hidden_dim': {'values': [128]},
+        'lr': {'values': [5e-4]},
+    },
+}
+
+
+def main(args):
+    # setup random seeds
+    rds = np.random.RandomState(args.seed)
+    model_seeds = list(rds.randint(0, 10**6, size=(100,)))
+    data_seeds = list(rds.randint(0, 10**6, size=(100,)))
+
+    sweep_config = {
+        'learn_likelihood_std': {'value': args.learn_likelihood_std},
+        'data_source': {'value': args.data_source},
+        'pred_diff': {'value': args.pred_diff},
+        'model': {'value': args.model},
+    }
+    # update with model specific sweep ranges
+    assert args.model in MODEL_SPECIFIC_CONFIG
+    sweep_config.update(MODEL_SPECIFIC_CONFIG[args.model])
+
+    # determine name of experiment
+    exp_base_path = os.path.join(RESULT_DIR, args.exp_name)
+    exp_path = os.path.join(exp_base_path, f'{args.data_source}_{args.model}')
+
+    if args.data_source == 'racecar':
+        N_SAMPLES_LIST = [50, 100, 200, 400, 800, 1600, 3200]
+    elif args.data_source == 'pendulum':
+        N_SAMPLES_LIST = [10, 20, 40, 80, 160, 320, 640]
+    else:
+        raise NotImplementedError(f'Unknown data source {args.data_source}.')
+
+    command_list = []
+    output_file_list = []
+    for _ in range(args.num_hparam_samples):
+        flags = sample_param_flags(sweep_config)
+        exp_hash = hash_dict(flags)
+
+        for num_samples_train in N_SAMPLES_LIST:
+
+            exp_result_folder = os.path.join(exp_path, f'{exp_hash}_{num_samples_train}')
+            flags['exp_result_folder'] = exp_result_folder
+
+            for model_seed, data_seed in itertools.product(model_seeds[:args.num_model_seeds],
+                                                           data_seeds[:args.num_data_seeds]):
+                cmd = generate_base_command(experiments.meta_learning_exp.run_meta_learning_exp,
+                                            flags=dict(**flags, **{'model_seed': model_seed, 'data_seed': data_seed,
+                                                                   'num_samples_train': num_samples_train}))
+                command_list.append(cmd)
+                output_file_list.append(os.path.join(exp_result_folder, f'{model_seed}_{data_seed}.out'))
+
+    generate_run_commands(command_list, output_file_list, num_cpus=args.num_cpus, mem=8*1024,
+                          duration='11:59:00' if args.long else '3:59:00',
+                          num_gpus=1 if args.gpu else 0, mode=args.run_mode, prompt=not args.yes)
+
+
+if __name__ == '__main__':
+    current_date = datetime.datetime.now().strftime("%b%d").lower()
+    parser = argparse.ArgumentParser(description='Meta-BO run')
+
+    # sweep args
+    parser.add_argument('--num_hparam_samples', type=int, default=1)
+    parser.add_argument('--num_model_seeds', type=int, default=5, help='number of model seeds per hparam')
+    parser.add_argument('--num_data_seeds', type=int, default=5, help='number of model seeds per hparam')
+    parser.add_argument('--num_cpus', type=int, default=1, help='number of cpus to use')
+    parser.add_argument('--run_mode', type=str, default='euler')
+
+    # general args
+    parser.add_argument('--exp_name', type=str, default=f'test_{current_date}')
+    parser.add_argument('--seed', type=int, default=94563)
+    parser.add_argument('--gpu', default=False, action='store_true')
+    parser.add_argument('--yes', default=False, action='store_true')
+    parser.add_argument('--long', default=False, action='store_true')
+
+    # data parameters
+    parser.add_argument('--data_source', type=str, default='racecar')
+    parser.add_argument('--pred_diff', type=int, default=1)
+
+    # # standard BNN parameters
+    parser.add_argument('--model', type=str, default='BNN_SVGD')
+    parser.add_argument('--learn_likelihood_std', type=int, default=0)
+
+    args = parser.parse_args()
+    main(args)

experiments/offline_rl_from_recorded_data/exp.py

@@ -152,7 +152,7 @@ def experiment(horizon_len: int,
     model_key = jr.PRNGKey(model_seed)
     data_key = jr.PRNGKey(data_seed)
 
-    int_data_seed = jr.randint(key_data_seed, (), minval=0, maxval=2 ** 13 - 1)
+    int_data_seed = jr.randint(data_key, (), minval=0, maxval=2 ** 13 - 1)
     assert num_offline_collected_transitions <= 20_000, "Cannot have more than 20_000 points for training"
     if bool(obtain_consecutive_data):
         if bool(data_from_simulation):
@@ -201,9 +201,7 @@ def experiment(horizon_len: int,
                            },
                 data_seed=int(int_data_seed), )
 
-    # Deal with randomness
-    key = jr.PRNGKey(seed)
-    key_bnn, key_offline_rl, key_evaluation_trained_bnn, key_evaluation_pretrained_bnn = jr.split(key, 4)
+    key_bnn, key_offline_rl, key_evaluation_trained_bnn, key_evaluation_pretrained_bnn = jr.split(model_key, 4)
 
     standard_params = {
         'input_size': sim.input_size,
@@ -315,7 +313,8 @@ def experiment(horizon_len: int,
 
 def main(args):
     experiment(
-        seed=args.seed,
+        model_seed=args.model_seed,
+        data_seed=args.data_seed,
         project_name=args.project_name,
         horizon_len=args.horizon_len,
         sac_num_env_steps=args.sac_num_env_steps,
@@ -356,10 +355,11 @@ def main(args):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--model_seed', type=int, default=0)
+    parser.add_argument('--data_seed', type=int, default=0)
     parser.add_argument('--horizon_len', type=int, default=200)
     parser.add_argument('--sac_num_env_steps', type=int, default=10_000)
-    parser.add_argument('--project_name', type=str, default='RaceCarPPO')
+    parser.add_argument('--project_name', type=str, default='RaceCarTestExperiments')
     parser.add_argument('--learnable_likelihood_std', type=str, default='yes')
     parser.add_argument('--include_aleatoric_noise', type=int, default=1)
     parser.add_argument('--best_bnn_model', type=int, default=1)
@@ -388,7 +388,7 @@ def main(args):
     parser.add_argument('--max_train_steps', type=int, default=10_000)
     parser.add_argument('--min_train_steps', type=int, default=10_000)
     parser.add_argument('--length_scale_aditive_sim_gp', type=float, default=1.0)
-    parser.add_argument('--input_from_recorded_data', type=int, default=1)
+    parser.add_argument('--input_from_recorded_data', type=int, default=0)
     parser.add_argument('--obtain_consecutive_data', type=int, default=1)
     parser.add_argument('--lr', type=float, default=3e-4)
     parser.add_argument('--car_id', type=int, default=3)

experiments/offline_rl_from_recorded_data/launcher.py

@@ -5,7 +5,8 @@
 
 _applicable_configs = {
     'horizon_len': [100],
-    'seed': list(range(10)),
+    'model_seed': list(range(5)),
+    'data_seed': list(range(5)),
     'project_name': [PROJECT_NAME],
     'sac_num_env_steps': [2_000_000],
     'num_epochs': [50],

experiments/regression_exp/plots_num_data.py

@@ -40,7 +40,7 @@ def different_method_plot(df_agg: pd.DataFrame, metric: str = 'nll', display: bo
     return best_rows_df, fig
 
 
-def main(args, drop_nan=False):
+def main(args, drop_nan=False, fig=None, show_legend=True):
     df_full, param_names = collect_exp_results(exp_name=args.exp_name)
     df_full = df_full[df_full['data_source'] == args.data_source]
     #df_full = df_full[df_full['model'] == args.data_source]
@@ -80,7 +80,7 @@ def main(args, drop_nan=False):
         df_full[column] = df_full[column].fillna('N/A')
 
     # first take mean over the data seeds
-    df_mean = df_full.groupby(by=groupby_names + ['model_seed'], axis=0).mean()
+    df_mean = df_full.groupby(by=groupby_names + ['model_seed'], axis=0).mean(numeric_only=True)
     df_mean.reset_index(drop=False, inplace=True)
 
     # then compute the stats over the model seeds
@@ -91,35 +91,61 @@ def main(args, drop_nan=False):
     # filter all the rows where the count is less than 3
     df_agg = df_agg[df_agg['rmse']['count'] >= 3]
 
-    available_models = set(df_agg['model'])
-    for metric in ['nll', 'rmse']:
-        fig, ax = plt.subplots()
+    available_models = sorted(list(set(df_agg['model'])))
+    if fig is None:
+        fig, axs = plt.subplots(1, 2)
+    else:
+        axs = fig.subplots(1, 2)
+    for idx, metric in enumerate(['nll', 'rmse']):
         for model in available_models:
             df_model = df_agg[df_agg['model'] == model].sort_values(by=[('num_samples_train', '')], ascending=True)
             if args.quantile_cis:
-                ax.plot(df_model[('num_samples_train', '')], df_model[(metric, 'median')], label=model)
+                axs[idx].plot(df_model[('num_samples_train', '')], df_model[(metric, 'median')], label=model)
                 lower_ci = df_model[(metric, 'lcb')]
                 upper_ci = df_model[(metric, 'ucb')]
             else:
-                ax.plot(df_model[('num_samples_train', '')], df_model[(metric, 'mean')], label=model)
+                axs[idx].plot(df_model[('num_samples_train', '')], df_model[(metric, 'mean')], label=model)
                 CI_width = 2 * np.sqrt(df_model[(metric, 'count')])
                 lower_ci = df_model[(metric, 'mean')] - CI_width * df_model[(metric, 'std')]
                 upper_ci = df_model[(metric, 'mean')] + CI_width * df_model[(metric, 'std')]
-            ax.fill_between(df_model[('num_samples_train', '')], lower_ci, upper_ci, alpha=0.3)
-        ax.set_title(f'{args.data_source} - {metric}')
-        #ax.set_xscale('log')
+            axs[idx].fill_between(df_model[('num_samples_train', '')], lower_ci, upper_ci, alpha=0.3)
+        axs[idx].set_title(f'{args.data_source} - {metric}')
+        # ax.set_xscale('log')
         #ax.set_ylim((-3.8, -2.))
 
-        ax.legend()
-        fig.show()
+    handles, labels = plt.gca().get_legend_handles_labels()
+    by_label = dict(zip(labels, handles))
+    if show_legend:
+        fig.legend(by_label.values(), by_label.keys(), ncols=4, loc='lower center',
+                   bbox_to_anchor=(0.5, 0), fontsize=10)
+        # fig.tight_layout(rect=[0, 0.1, 1, 1])
+
+    # plt.show()
     print('Models:', set(df_agg['model']))
 
 
 
 if __name__ == '__main__':
+    figure = plt.figure(figsize=(10, 7))
+    subfigs = figure.subfigures(2, 1, hspace=-0.1)
+
+
+    parser = argparse.ArgumentParser(description='Inspect results of a regression experiment.')
+    parser.add_argument('--exp_name', type=str, default='jan10_num_data')
+    parser.add_argument('--quantile_cis', type=int, default=1)
+    parser.add_argument('--data_source', type=str, default='racecar')
+    args = parser.parse_args()
+
+
+    main(args, fig=subfigs[0], show_legend=False)
+
     parser = argparse.ArgumentParser(description='Inspect results of a regression experiment.')
     parser.add_argument('--exp_name', type=str, default='jan10_num_data')
     parser.add_argument('--quantile_cis', type=int, default=1)
     parser.add_argument('--data_source', type=str, default='pendulum')
     args = parser.parse_args()
-    main(args)
+
+    main(args, fig=subfigs[1])
+    plt.tight_layout(rect=[0, 0.2, 1, 0.9])
+    plt.savefig('regression_exp.pdf')
+    plt.show()