diff --git a/.gitignore b/.gitignore index af684323..8b624483 100644 --- a/.gitignore +++ b/.gitignore @@ -168,3 +168,4 @@ Temporary Items # exlcluded files plotting_scripts/plots +/experiments/1d_visualization/plot_dicts/ diff --git a/experiments/bimodal_exp/bimodal_exp.py b/experiments/bimodal_exp/bimodal_exp.py index fc1910c7..6d95e04c 100644 --- a/experiments/bimodal_exp/bimodal_exp.py +++ b/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() diff --git a/experiments/bimodal_exp/bimodal_launcher.py b/experiments/bimodal_exp/bimodal_launcher.py index 74347b63..04af3c81 100644 --- a/experiments/bimodal_exp/bimodal_launcher.py +++ b/experiments/bimodal_exp/bimodal_launcher.py @@ -1,11 +1,12 @@ 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] } @@ -13,17 +14,20 @@ 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__': diff --git a/experiments/data_provider.py b/experiments/data_provider.py index 73286d6e..57b785b7 100644 --- a/experiments/data_provider.py +++ b/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')] diff --git a/experiments/meta_learning_exp/sweep_regression_exp_num_data.py b/experiments/meta_learning_exp/sweep_regression_exp_num_data.py new file mode 100644 index 00000000..5bda246e --- /dev/null +++ b/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) diff --git a/experiments/offline_rl_from_recorded_data/exp.py b/experiments/offline_rl_from_recorded_data/exp.py index 442ec24e..af2e2762 100644 --- a/experiments/offline_rl_from_recorded_data/exp.py +++ b/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) diff --git a/experiments/offline_rl_from_recorded_data/launcher.py b/experiments/offline_rl_from_recorded_data/launcher.py index 10dca1ca..8f2bb8ab 100644 --- a/experiments/offline_rl_from_recorded_data/launcher.py +++ b/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], diff --git a/experiments/regression_exp/plots_num_data.py b/experiments/regression_exp/plots_num_data.py index ed6ffc93..4d7745cf 100644 --- a/experiments/regression_exp/plots_num_data.py +++ b/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) \ No newline at end of file + + main(args, fig=subfigs[1]) + plt.tight_layout(rect=[0, 0.2, 1, 0.9]) + plt.savefig('regression_exp.pdf') + plt.show() diff --git a/lenart_internal/analysis_simulator_data.py b/lenart_internal/analysis_simulator_data.py index 12efe5d1..9e4e046f 100644 --- a/lenart_internal/analysis_simulator_data.py +++ b/lenart_internal/analysis_simulator_data.py @@ -38,10 +38,8 @@ def join_group_names_and_prepare_statistics_mean(data: pd.DataFrame): def join_group_names_and_prepare_statistics(data: pd.DataFrame): summary_rewards = data.groupby('Group')[reward_column].agg(['median', - lambda x: x.quantile(0.2), - # 0.2 quantile - lambda x: x.quantile(0.8) - # 0.8 quantile + lambda x: x.quantile(0.2), # 0.2 quantile + lambda x: x.quantile(0.8) # 0.8 quantile ]) summary = pd.concat([summary_rewards], axis=1) diff --git a/sim_transfer/models/abstract_model.py b/sim_transfer/models/abstract_model.py index 6992d4b8..8670073b 100644 --- a/sim_transfer/models/abstract_model.py +++ b/sim_transfer/models/abstract_model.py @@ -1,18 +1,22 @@ import copy import logging +import os +import pickle import time from typing import Optional, Tuple, Callable, Dict, List, Union -from jaxtyping import PyTree + import jax import jax.numpy as jnp import matplotlib.pyplot as plt import numpy as np +import optax import tensorflow as tf import tensorflow_datasets as tfds import tensorflow_probability.substrates.jax.distributions as tfd import wandb +from jaxtyping import PyTree from tensorflow_probability.substrates import jax as tfp -import optax + from sim_transfer.modules.distribution import AffineTransform from sim_transfer.modules.metrics import _check_dist_and_sample_shapes, _get_mean_std_from_dist from sim_transfer.modules.nn_modules import BatchedMLP @@ -246,12 +250,21 @@ def eval(self, x: jnp.ndarray, y: np.ndarray, prefix: str = '', per_dim_metrics: eval_stats = {name: float(val) for name, val in eval_stats.items()} return eval_stats - def plot_1d(self, x_train: jnp.ndarray, y_train: jnp.ndarray, - domain_l: Optional[float] = None, domain_u: Optional[float] = None, - true_fun: Optional[Callable] = None, title: Optional[str] = '', show: bool = True, - plot_posterior_samples: bool = False, log_to_wandb: bool = False): + def plot_1d(self, + x_train: jnp.ndarray, + y_train: jnp.ndarray, + domain_l: Optional[float] = None, + domain_u: Optional[float] = None, + true_fun: Optional[Callable] = None, + title: Optional[str] = '', + show: bool = True, + plot_posterior_samples: bool = False, + log_to_wandb: bool = False, + save_plot_dict: bool = False): assert self.input_size == 1, 'Can only plot if input_size = 1' + plot_data = dict() + # determine plotting domain x_min, x_max = jnp.min(x_train, axis=0), jnp.max(x_train, axis=0) width = x_max - x_min @@ -269,24 +282,40 @@ def plot_1d(self, x_train: jnp.ndarray, y_train: jnp.ndarray, if self.output_size == 1: ax = [ax] for i in range(self.output_size): - ax[i].scatter(x_train.flatten(), y_train[:, i], label='train points') + plot_data_dim = dict() + plot_data_dim['Train points'] = (x_train.flatten(), y_train[:, i]) + ax[i].scatter(*plot_data_dim['Train points'], label='train points') if true_fun is not None: - ax[i].plot(x_plot, true_fun(x_plot)[:, i], label='true fun') - ax[i].plot(x_plot.flatten(), pred_mean[:, i], label='pred mean') - ax[i].fill_between(x_plot.flatten(), pred_mean[:, i] - pred_std[:, i], - pred_mean[:, i] + pred_std[:, i], alpha=0.3) + plot_data_dim['True function'] = (x_plot, true_fun(x_plot)[:, i]) + ax[i].plot(*plot_data_dim['True function'], label='true fun') + + plot_data_dim['Mean prediction'] = (x_plot.flatten(), pred_mean[:, i]) + plot_data_dim['Fill between std'] = (x_plot.flatten(), pred_mean[:, i] - pred_std[:, i], + pred_mean[:, i] + pred_std[:, i]) + ax[i].plot(*plot_data_dim['Mean prediction'], label='pred mean') + ax[i].fill_between(*plot_data_dim['Fill between std'], alpha=0.3) if hasattr(self, 'predict_post_samples') and plot_posterior_samples: y_post_samples = self.predict_post_samples(x_plot) + plot_data_dim['Sample predictions'] = (x_plot, y_post_samples[..., i]) for y in y_post_samples: ax[i].plot(x_plot, y[:, i], linewidth=0.2, color='green') if title is not None: ax[i].set_title(f'Output dimension {i}') ax[i].legend() + plot_data[f'Dimension_{i}'] = plot_data_dim fig.suptitle(title) if log_to_wandb: wandb.log({title: wandb.Image(fig)}) + if save_plot_dict: + directory = os.path.join(wandb.run.dir, 'data') + if not os.path.exists(directory): + os.makedirs(directory) + data_path = os.path.join('data', 'plot_data.pkl') + with open(os.path.join(wandb.run.dir, data_path), 'wb') as handle: + pickle.dump(plot_data, handle) + wandb.save(os.path.join(wandb.run.dir, data_path), wandb.run.dir) if show: fig.show()