from argparse import ArgumentParser
import numpy as np
import pickle

parser = ArgumentParser()
parser.add_argument("pkl_path")
# parser.add_argument("out")
args, _ = parser.parse_known_args()
data = pickle.load(open(args.pkl_path, 'rb'))

for traj in data:
    traj['observations'] = traj['actions'].astype(np.float32)

out = args.pkl_path[:-4] + "_actions.pkl"
print(f"Success! Dumping to {out}")
pickle.dump(data, open(out, 'wb'))
