diff --git a/experiments/data_provider.py b/experiments/data_provider.py index 762e22e..0eee028 100644 --- a/experiments/data_provider.py +++ b/experiments/data_provider.py @@ -96,7 +96,7 @@ 'likelihood_std': {'value': _RACECAR_NOISE_STD_ENCODED.tolist()}, 'num_samples_train': {'value': 200}, } for name in ['real_racecar_new', 'real_racecar_new_only_pose', 'real_racecar_new_no_angvel', - 'real_racecar_new_actionstack'] + 'real_racecar_new_actionstack', 'real_racecar_v2'] }) @@ -187,18 +187,25 @@ def _rccar_transitions_to_dataset(transitions: Transition, encode_angles: bool = def get_rccar_recorded_data_new(encode_angle: bool = True, skip_first_n_points: int = 10, + dataset: str = 'all', 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 + assert dataset in ['all', 'v1'] recordings_dir = [os.path.join(DATA_DIR, 'recordings_rc_car_v1')] elif car_id == 2: - num_train_traj = 12 - recordings_dir = [os.path.join(DATA_DIR, 'recordings_rc_car_v2'), - os.path.join(DATA_DIR, 'recordings_rc_car_v3'), - os.path.join(DATA_DIR, 'recordings_rc_car_v4')] + if dataset == 'all': + recordings_dir = [os.path.join(DATA_DIR, 'recordings_rc_car_v2'), + os.path.join(DATA_DIR, 'recordings_rc_car_v3'), + os.path.join(DATA_DIR, 'recordings_rc_car_v4')] + num_test_points = 20_000 + elif dataset in ['v2', 'v3', 'v4']: + recordings_dir = [os.path.join(DATA_DIR, f'recordings_rc_car_{dataset}')] + num_test_points = 10_000 + else: + raise ValueError(f"Unknown dataset {dataset} for car_id {car_id}") else: raise ValueError(f"Unknown car id {car_id}") files = [sorted(glob.glob(rd + '/*.pickle')) for rd in recordings_dir] @@ -208,18 +215,16 @@ def get_rccar_recorded_data_new(encode_angle: bool = True, skip_first_n_points: # 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] # 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, y = map(lambda x: jnp.concatenate(x, axis=0), zip(*map(prep_fn, transitions))) - # x_test, y_test = map(lambda x: jnp.concatenate(x, axis=0), zip(*map(prep_fn, transitions[num_train_traj:]))) indices = jnp.arange(start=0, stop=x.shape[0], step=1) indices = jax.random.shuffle(key=jax.random.PRNGKey(9345), x=indices) x, y = x[indices], y[indices] - num_test_points = 20_000 + + # split into train and test x_train, y_train, x_test, y_test = x[:-num_test_points], y[:-num_test_points], \ x[-num_test_points:], y[-num_test_points:] return x_train, y_train, x_test, y_test @@ -380,6 +385,10 @@ def provide_data_and_sim(data_source: str, data_spec: Dict[str, Any], data_seed: elif data_source.startswith('real_racecar_new'): x_train, y_train, x_test, y_test = get_rccar_recorded_data_new(encode_angle=True, action_stacking=False, action_delay=3, car_id=car_id) + elif data_source.startswith('real_racecar_v2'): + x_train, y_train, x_test, y_test = get_rccar_recorded_data_new(encode_angle=True, action_stacking=False, + action_delay=3, car_id=car_id, + dataset='v2') else: x_train, y_train, x_test, y_test = get_rccar_recorded_data(encode_angle=True) diff --git a/experiments/util.py b/experiments/util.py index 1ddf645..cc366e5 100644 --- a/experiments/util.py +++ b/experiments/util.py @@ -321,20 +321,14 @@ def collect_exp_results(exp_name: str, dir_tree_depth: int = 3, verbose: bool = def ucb(row): - if row.shape[0] > 1: - return np.nan return np.quantile(row, q=0.95, axis=0) def lcb(row): - if row.shape[0] > 1: - return np.nan return np.quantile(row, q=0.05, axis=0) def median(row): - if row.shape[0] > 1: - return np.nan return np.quantile(row, q=0.5, axis=0)