Merge remote-tracking branch 'origin/main' into main

baselines/pacoh_nn/bnn/bnn_svgd.py

@@ -14,8 +14,9 @@
 class BayesianNeuralNetworkSVGD(RegressionModel):
 
     def __init__(self, x_train, y_train, hidden_layer_sizes=(32, 32, 32, 32), activation='relu',
-                 likelihood_std=0.1, learn_likelihood=True, prior_std=0.1, prior_weight=1e-4,
-                 likelihood_prior_mean=tf.math.log(0.1), likelihood_prior_std=1.0, sqrt_mode=False,
+                 likelihood_std=0.1, learn_likelihood=True, normalize_likelihood_std: bool = True,
+                 prior_std=0.1, prior_weight=1e-4, likelihood_prior_mean=tf.math.log(0.1),
+                 likelihood_prior_std=1.0, sqrt_mode=False,
                  n_particles=10, batch_size=8, bandwidth=100., lr=1e-3, meta_learned_prior=None,
                  normalization_stats=None):
 
@@ -53,6 +54,8 @@ def __init__(self, x_train, y_train, hidden_layer_sizes=(32, 32, 32, 32), activa
             self.meta_learned_prior_mode = True
 
         # Likelihood
+        if normalize_likelihood_std:
+            likelihood_std = tf.convert_to_tensor(likelihood_std) / tf.reshape(self.y_std, (self.output_dim,))
         self.likelihood = GaussianLikelihood(self.output_dim, n_particles,
                                              trainable=learn_likelihood, std=likelihood_std)
 

baselines/pacoh_nn/pacoh_nn_regression.py

@@ -35,6 +35,8 @@ def __init__(self, meta_train_data, lr=2e-3,  hidden_layer_sizes=(32, 32, 32, 32
         # prepare config for meta eval model
         self.eval_model_config = {'hidden_layer_sizes': hidden_layer_sizes, 'activation': activation,
                                   'learn_likelihood': learn_likelihood,
+                                  'likelihood_std': likelihood_std,
+                                  'normalize_likelihood_std': normalize_likelihood_std,
                                   'n_particles': num_posterior_particles * self.n_batched_priors,
                                   'prior_weight': prior_weight, 'bandwidth': bandwidth,
                                   'batch_size': batch_size}

experiments/data_provider.py

@@ -1,11 +1,13 @@
 import os
 import pickle
-from functools import partial
-from typing import Dict, Any
-
+import glob
 import jax
 import jax.numpy as jnp
 
+from functools import partial
+from typing import Dict, Any, List
+from brax.training.types import Transition
+
 from experiments.util import load_csv_recordings
 from sim_transfer.sims.car_sim_config import OBS_NOISE_STD_SIM_CAR
 from sim_transfer.sims.simulators import PredictStateChangeWrapper, StackedActionSimWrapper
@@ -132,88 +134,88 @@ def prepare_rccar_data(df, encode_angles: bool = False, change_signs: bool = Tru
     return x_train, y_train, x_test, y_test
 
 
-def get_rccar_recorded_data_new(encode_angle: bool = True, skip_first_n_points: int = 10,
-                                action_delay: int = 3, action_stacking: bool = False,
-                                car_id: int = 2, key_data: jax.random.PRNGKeyArray = jax.random.PRNGKey(0)):
-    from brax.training.types import Transition
-
-    assert car_id in [1, 2]
-    if car_id == 1:
-        num_train_traj = 8
-        recordings_dir = os.path.join(DATA_DIR, 'recordings_rc_car_v1')
-        import glob
-        file_name = glob.glob(recordings_dir + '/*.pickle')
-    elif car_id == 2:
-        num_train_traj = 12
-        recordings_dir = os.path.join(DATA_DIR, 'recordings_rc_car_v2')
-        import glob
-        file_name = glob.glob(recordings_dir + '/*.pickle')
-
-    else:
-        raise ValueError(f"Unknown car id {car_id}")
-
+def _load_transitions(file_names: List[str]) -> List[Transition]:
     transitions = []
-    for fn in file_name:
+    for fn in file_names:
         with open(fn, 'rb') as f:
             data = pickle.load(f)
             assert (data.observation.shape[-1] == 13 or data.observation.shape[-1] == 12), "state must be 12D or 13D"
             if data.observation.shape[-1] == 13:
-                x = decode_angles_fn(data.observation, angle_idx=2)
-                x_next = decode_angles_fn(data.next_observation, angle_idx=2)
                 data = Transition(
-                    observation=x,
-                    action=data.action,
-                    reward=data.reward,
-                    discount=data.discount,
-                    next_observation=x_next,
-                    extras=data.extras,
-                )
+                    observation=decode_angles_fn(data.observation, angle_idx=2),
+                    action=data.action, reward=data.reward, discount=data.discount,
+                    next_observation=decode_angles_fn(data.next_observation, angle_idx=2), extras=data.extras,)
             transitions.append(data)
+    return transitions
 
-    indices = jnp.arange(0, len(transitions))
-    indices = jax.random.shuffle(key=key_data, x=indices)
-    transitions = [transitions[idx] for idx in indices]
 
-    def prepare_rccar_data(transitions: Transition, encode_angles: bool = False, skip_first_n: int = 30,
-                           action_delay: int = 3, action_stacking: bool = False):
-        assert 0 <= action_delay <= 3, "Only recorded the last 3 actions and the current action"
-        assert action_delay >= 0, "Action delay must be non-negative"
+def _rccar_transitions_to_dataset(transitions: Transition, encode_angles: bool = False, skip_first_n: int = 30,
+                                  action_delay: int = 3, action_stacking: bool = False):
+    assert 0 <= action_delay <= 3, "Only recorded the last 3 actions and the current action"
+    assert action_delay >= 0, "Action delay must be non-negative"
 
-        if action_stacking:
-            if action_delay == 0:
-                u = transitions.action
-            else:
-                u = jnp.concatenate([transitions.observation[:, -2 * action_delay:], transitions.action], axis=-1)
-            assert u.shape[-1] == 2 * (action_delay + 1)
+    if action_stacking:
+        if action_delay == 0:
+            u = transitions.action
         else:
-            if action_delay == 0:
-                u = transitions.action
-            elif action_delay == 1:
-                u = transitions.observation[:, -2:]
-            else:
-                u = transitions.observation[:, -2 * action_delay: -2 * (action_delay - 1)]
-            assert u.shape[-1] == 2
+            u = jnp.concatenate([transitions.observation[:, -2 * action_delay:], transitions.action], axis=-1)
+        assert u.shape[-1] == 2 * (action_delay + 1)
+    else:
+        if action_delay == 0:
+            u = transitions.action
+        elif action_delay == 1:
+            u = transitions.observation[:, -2:]
+        else:
+            u = transitions.observation[:, -2 * action_delay: -2 * (action_delay - 1)]
+        assert u.shape[-1] == 2
 
-        x = transitions.observation[:, :6]
-        y = transitions.next_observation[: , :6]
+    # make sure that we only take the actual state without the stacked actions
+    x = transitions.observation[:, :6]
+    y = transitions.next_observation[:, :6]
 
-        # project theta into [-\pi, \pi]
-        x[:, 2] = (x[:, 2] + jnp.pi) % (2 * jnp.pi) - jnp.pi
-        y[:, 2] = (y[:, 2] + jnp.pi) % (2 * jnp.pi) - jnp.pi
-        if encode_angles:
-            x = encode_angles_fn(x, angle_idx=2)
-            y = encode_angles_fn(y, angle_idx=2)
+    # project theta into [-\pi, \pi]
+    x[:, 2] = (x[:, 2] + jnp.pi) % (2 * jnp.pi) - jnp.pi
+    y[:, 2] = (y[:, 2] + jnp.pi) % (2 * jnp.pi) - jnp.pi
+    if encode_angles:
+        x = encode_angles_fn(x, angle_idx=2)
+        y = encode_angles_fn(y, angle_idx=2)
 
-        # remove first n steps (since often not much is happening)
-        x, u, y = x[skip_first_n:], u[skip_first_n:], y[skip_first_n:]
-        x_data = jnp.concatenate([x, u], axis=-1)  # current state + action
-        y_data = y # next state
-        assert x_data.shape[0] == y_data.shape[0]
-        assert x_data.shape[1] - (2 * (1 + int(action_stacking) * action_delay)) == y_data.shape[1]
-        return x_data, y_data
+    # remove first n steps (since often not much is happening)
+    x, u, y = x[skip_first_n:], u[skip_first_n:], y[skip_first_n:]
 
+    # concatenate state and action
+    x_data = jnp.concatenate([x, u], axis=-1)  # current state + action
+    y_data = y  # next state
+
+    # check shapes
+    assert x_data.shape[0] == y_data.shape[0]
+    assert x_data.shape[1] - (2 * (1 + int(action_stacking) * action_delay)) == y_data.shape[1]
+
+    return x_data, y_data
+
+
+def get_rccar_recorded_data_new(encode_angle: bool = True, skip_first_n_points: int = 10,
+                                action_delay: int = 3, action_stacking: bool = False,
+                                car_id: int = 2):
+
+    assert car_id in [1, 2]
+    if car_id == 1:
+        num_train_traj = 8
+        recordings_dir = os.path.join(DATA_DIR, 'recordings_rc_car_v1')
+    elif car_id == 2:
+        num_train_traj = 10
+        recordings_dir = os.path.join(DATA_DIR, 'recordings_rc_car_v2')
+    else:
+        raise ValueError(f"Unknown car id {car_id}")
+    file_names = sorted(glob.glob(recordings_dir + '/*.pickle'))
+
+    # load and shuffle transitions
+    transitions = _load_transitions(file_names)
+    indices = jax.random.permutation(key=jax.random.PRNGKey(9345), x=jnp.arange(0, len(transitions)))
+    transitions = [transitions[idx] for idx in indices]
 
-    prep_fn = partial(prepare_rccar_data, encode_angles=encode_angle, skip_first_n=skip_first_n_points,
+    # transform transitions into supervised learning datasets
+    prep_fn = partial(_rccar_transitions_to_dataset, encode_angles=encode_angle, skip_first_n=skip_first_n_points,
                       action_delay=action_delay, action_stacking=action_stacking)
     x_train, y_train = map(lambda x: jnp.concatenate(x, axis=0), zip(*map(prep_fn, transitions[:num_train_traj])))
     x_test, y_test = map(lambda x: jnp.concatenate(x, axis=0), zip(*map(prep_fn, transitions[num_train_traj:])))

experiments/lf_hf_transfer_exp/inspect_results.py

@@ -71,7 +71,12 @@ def main(args, drop_nan=False):
         df_full[column] = df_full[column].fillna('N/A')
 
     groupby_names.remove('likelihood_std')
-    groupby_names.remove('added_gp_outputscale')
+    # groupby_names.remove('added_gp_outputscale')
+
+    # replace the added_gp_outputscale with the first element of the list
+    if not type(df_full['added_gp_outputscale'].iloc[0]) == np.float_:
+
+        df_full['added_gp_outputscale'] = df_full['added_gp_outputscale'].apply(lambda x: x[0])
 
     # first take mean over the data seeds
     df_mean = df_full.groupby(by=groupby_names + ['model_seed'], axis=0).mean()
@@ -101,62 +106,16 @@ def main(args, drop_nan=False):
     for param in ['num_measurement_points', 'added_gp_lengthscale', 'added_gp_outputscale']:# ['num_f_samples', 'bandwidth_score_estim', 'bandwidth_svgd', 'num_measurement_points']:
         plt.scatter(df_method[param], df_method[(metric, 'mean')])
         plt.xlabel(param)
-        plt.xscale('log')
+        #plt.xscale('log')
         plt.ylabel(metric)
-        #plt.scatter(0.2, -7, color='red')
-        #plt.ylim(-8, 0)
+        # plt.scatter(1.5, -3.5, color='red')
+        # plt.ylim(-3.6, -3.)
         plt.show()
 
-    QUANTILE_BASED_CI = True
-    METHODS = ['random_search', 'hill_search']
-    METRICS = ['x_diff', 'f_diff']
-    PLOT_N_BEST = 2
-    n_metrics = len(METRICS)
-
-    fig, axes = plt.subplots(ncols=n_metrics, nrows=1, figsize=(n_metrics * 4, 4))
-    for k, method in enumerate(METHODS):
-        df_plot = df_agg.loc[df_agg['method'] == method]
-        df_plot.sort_values(by=('x_diff', 'mean'), ascending=True, inplace=True)
-
-        if df_plot.empty:
-            continue
-
-        for i, metric in enumerate(METRICS):
-            for j in range(PLOT_N_BEST):
-                row = df_plot.iloc[j]
-                num_seeds = row[(metric, 'count')]
-                ci_factor = 1 / np.sqrt(num_seeds)
-
-                if QUANTILE_BASED_CI:
-                    metric_median = row[(metric, 'median')]
-                    axes[i].scatter(k, metric_median, label=f'{method}_{j}')
-                    lower_err = - (row[(metric, 'lcb')] - metric_median) * ci_factor
-                    upper_err = (row[(metric, 'ucb')] - metric_median) * ci_factor
-                    axes[i].errorbar(k, metric_median, yerr=np.array([[lower_err, upper_err]]).T,
-                                        capsize=5)
-                else:
-                    metric_mean = row[(metric, 'mean')]
-                    metric_std = row[(metric, 'std')]
-                    axes[i].scatter(k, metric_mean, label=f'{method}_{j}')
-                    axes[i].errorbar(k, metric_mean, yerr=2 * metric_std * ci_factor)
-
-                if i == 0:
-                    print(f'{method}_{j}', row['exp_result_folder'][0])
-            if k == 0:
-                axes[i].set_ylabel(metric)
-            axes[i].set_xticks(np.arange(len(METHODS)))
-            axes[i].set_xticklabels(METHODS)
-            axes[i].set_xlim((-0.5, len(METHODS) - 0.5))
-            axes[i].set_yscale('log')
-
-    plt.legend()
-    plt.tight_layout()
-    plt.show()
-
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='Inspect results of a regression experiment.')
-    parser.add_argument('--exp_name', type=str, default='aug21_lfhf_v2')
-    parser.add_argument('--data_source', type=str, default='racecar_hf_no_angvel')
+    parser.add_argument('--exp_name', type=str, default='aug24_lfhf')
+    parser.add_argument('--data_source', type=str, default='racecar_hf')
     args = parser.parse_args()
     main(args)

experiments/lf_hf_transfer_exp/run_regression_exp.py

@@ -35,6 +35,7 @@ def regression_experiment(
                           data_source: str,
                           num_samples_train: int,
                           data_seed: int = 981648,
+                          pred_diff: bool = False,
 
                           # logging parameters
                           use_wandb: bool = False,
@@ -78,7 +79,8 @@ def regression_experiment(
     # provide data and sim
     x_train, y_train, x_test, y_test, sim_lf = provide_data_and_sim(
         data_source=data_source,
-        data_spec={'num_samples_train': num_samples_train},
+        data_spec={'num_samples_train': num_samples_train,
+                   'pred_diff': bool(pred_diff)},
         data_seed=data_seed)
 
     if model.endswith('_no_add_gp'):
@@ -114,6 +116,7 @@ def regression_experiment(
         'hidden_layer_sizes': [layer_size]*num_layers,
         'normalize_likelihood_std': normalize_likelihood_std,
         'learn_likelihood_std': bool(learn_likelihood_std),
+        'likelihood_exponent': likelihood_exponent,
     }
 
     if model == 'BNN_SVGD':
@@ -265,12 +268,13 @@ def main(args):
     parser.add_argument('--use_wandb', type=bool, default=False)
 
     # data parameters
-    parser.add_argument('--data_source', type=str, default='racecar_hf_no_angvel')
-    parser.add_argument('--num_samples_train', type=int, default=20000)
+    parser.add_argument('--data_source', type=str, default='pendulum_hf')
+    parser.add_argument('--num_samples_train', type=int, default=20)
     parser.add_argument('--data_seed', type=int, default=77698)
+    parser.add_argument('--pred_diff', type=int, default=0)
 
     # standard BNN parameters
-    parser.add_argument('--model', type=str, default='BNN_FSVGD_SimPrior_gp')
+    parser.add_argument('--model', type=str, default='BNN_SVGD')
     parser.add_argument('--model_seed', type=int, default=892616)
     parser.add_argument('--likelihood_std', type=float, default=None)
     parser.add_argument('--learn_likelihood_std', type=int, default=0)

experiments/lf_hf_transfer_exp/sweep_regression_exp.py

@@ -11,38 +11,38 @@
 MODEL_SPECIFIC_CONFIG = {
     'BNN_SVGD': {
         'bandwidth_svgd': {'distribution': 'log_uniform', 'min': -1., 'max': 4.},
-        'num_train_steps': {'values': [20000, 40000]}
+        'num_train_steps': {'values': [5000, 10000, 20000]}
     },
     'BNN_FSVGD': {
         'bandwidth_svgd': {'distribution': 'log_uniform_10', 'min': -1.0, 'max': 0.0},
         'bandwidth_gp_prior': {'distribution': 'log_uniform', 'min': -2., 'max': 0.},
-        'num_train_steps': {'values': [20000]},
+        'num_train_steps': {'values': [5000, 10000, 20000]},
         'num_measurement_points': {'values': [16, 32, 64, 128]},
     },
     'BNN_FSVGD_SimPrior_gp': {
         'bandwidth_svgd': {'distribution': 'log_uniform_10', 'min': -1.0, 'max': 0.0},
         'num_train_steps': {'values': [40000]},
-        'num_measurement_points': {'values': [16, 32]},
+        'num_measurement_points': {'values': [32]},
         'num_f_samples': {'values': [1024]},
-        'added_gp_lengthscale': {'distribution': 'uniform', 'min': 1., 'max': 20.0},   # racecar: 4 - 8
+        'added_gp_lengthscale': {'distribution': 'uniform', 'min': 2., 'max': 10.},   # racecar: 4 - 8  # pendulum (diff): > 8 - 15.
         #'added_gp_outputscale': {'distribution': 'uniform', 'min': 0.05, 'max': 0.5},  # racecar: 0.2 - 0.8
-        'added_gp_outputscale': {'distribution': 'uniform', 'min': 0.5, 'max': 2.0},   # racecar: 4 - 8
+        'added_gp_outputscale': {'distribution': 'uniform', 'min': 0.5, 'max': 4.},  # pendulum (diff): 1 - 3.0
     },
-    'BNN_FSVGD_SimPrior_ssge': {
+    'BNN_FSVGD_SimPrior_nu-method': {
         'bandwidth_svgd': {'distribution': 'log_uniform_10', 'min': -1.0, 'max': 0.0},
         'num_train_steps': {'values': [40000]},
-        'num_measurement_points': {'values': [8, 16, 32]},
+        'num_measurement_points': {'values': [32]},
         'num_f_samples': {'values': [512]},
-        'bandwidth_score_estim': {'distribution': 'log_uniform_10', 'min': -0.5, 'max': 1.},
+        'bandwidth_score_estim': {'distribution': 'uniform', 'min': 0.8, 'max': 2.0},
+        'added_gp_lengthscale': {'distribution': 'uniform', 'min': 2., 'max': 10.},
+        'added_gp_outputscale': {'distribution': 'uniform', 'min': 0.5, 'max': 4.},
     },
-    'BNN_FSVGD_SimPrior_nu-method': {
+    'BNN_FSVGD_SimPrior_ssge': {
         'bandwidth_svgd': {'distribution': 'log_uniform_10', 'min': -1.0, 'max': 0.0},
         'num_train_steps': {'values': [40000]},
-        'num_measurement_points': {'values': [16, 32]},
+        'num_measurement_points': {'values': [8, 16, 32]},
         'num_f_samples': {'values': [512]},
-        'bandwidth_score_estim': {'distribution': 'log_uniform_10', 'min': 0.0, 'max': 0.0},
-        'added_gp_lengthscale': {'distribution': 'uniform', 'min': 5., 'max': 10.0},
-        'added_gp_outputscale': {'distribution': 'uniform', 'min': 0.5, 'max': 2.0},   # racecar: 4 - 8
+        'bandwidth_score_estim': {'distribution': 'log_uniform_10', 'min': -0.5, 'max': 1.},
     },
     'BNN_FSVGD_SimPrior_gp+nu-method': {
         'bandwidth_svgd': {'distribution': 'log_uniform_10', 'min': -1.0, 'max': 0.0},
@@ -87,6 +87,7 @@ def main(args):
         #'likelihood_std': {'value': None},
         'num_particles': {'value': 20},
         'data_batch_size': {'value': 8},
+        'pred_diff': {'value': args.pred_diff},
     }
     # update with model specific sweep ranges
     model_name = args.model.replace('_no_add_gp', '')
@@ -134,7 +135,8 @@ def main(args):
     parser.add_argument('--yes', default=False, action='store_true')
 
     # data parameters
-    parser.add_argument('--data_source', type=str, default='pendulum')
+    parser.add_argument('--data_source', type=str, default='pendulum_hf')
+    parser.add_argument('--pred_diff', type=int, default=0)
 
     # # standard BNN parameters
     parser.add_argument('--model', type=str, default='BNN_SVGD')