import sklearn
import os
import pickle
import numpy as np

#an example folder path for cheetah data
dir_path_1 = "enter_path_here/cheetah-data/"
def get_folder_data(dir_path):

    f_paths = []

    for file_path in os.listdir(dir_path):
            f_paths.append(file_path)

    np_all_states = np.array([])
    count = 0
    for f_name in f_paths:
        if int(f_name.split("_")[0]) < 200:
            with open(dir_path + f_name, "rb") as f_in:
                state_arrs = pickle.load(f_in)
                for state_arr in state_arrs:
                    if count == 0:
                        np_all_states = state_arr["observations"]
                        count += 1
                    else:
                        np_all_states = np.concatenate((np_all_states, state_arr["observations"]))
    return np_all_states

np_all_states = get_folder_data(dir_path_1)
print(np_all_states.shape)
with open("all_states_cheetah.pkl", "wb") as f_out:
    pickle.dump(np_all_states, f_out)




