From b01887001dd63697216378fa3c03d56e0b88414d Mon Sep 17 00:00:00 2001 From: Ruihan Zhao Date: Sun, 25 Feb 2024 14:42:44 -0600 Subject: [PATCH 1/4] NutAssembly done. --- .../scripts/record_robomimic_data.py | 41 ++++++++++++++++--- 1 file changed, 36 insertions(+), 5 deletions(-) diff --git a/task_decomposition/scripts/record_robomimic_data.py b/task_decomposition/scripts/record_robomimic_data.py index cc83b6f..bbe6902 100644 --- a/task_decomposition/scripts/record_robomimic_data.py +++ b/task_decomposition/scripts/record_robomimic_data.py @@ -33,13 +33,15 @@ def get_data_to_record(env_name: str): # Unique to each environment if env_name == "NutAssemblySquare": - meta_data_to_record = ["nut0_pos", "nut1_pos"] + meta_data_to_record = ["nut0_pos", "peg1_pos"] elif env_name == "ToolHang": meta_data_to_record = ["tool_pos", "frame_pos", "stand_pos"] else: raise ValueError(f"Environment {env_name} has no defined data to record.") - return ["step"] + actions_to_record + obs_to_record + meta_data_to_record + return ( + ["step", "sub_task"] + actions_to_record + obs_to_record + meta_data_to_record + ) def query_sim_for_data(env, desired_obs): @@ -51,6 +53,10 @@ def query_sim_for_data(env, desired_obs): return env.env.sim.data.body_xpos[env.env.obj_body_id[env.env.nuts[0].name]] elif desired_obs == "nut1_pos": return env.env.sim.data.body_xpos[env.env.obj_body_id[env.env.nuts[1].name]] + elif desired_obs == "peg1_pos": + return env.env.sim.data.body_xpos[env.env.peg1_body_id] + elif desired_obs == "peg2_pos": + return env.env.sim.data.body_xpos[env.env.peg2_body_id] else: raise ValueError(f"Environment {env.name} has no defined data to record.") elif env.name == "ToolHang": @@ -92,8 +98,18 @@ def extract_trajectory(env, initial_state, states, actions, done_mode): traj_len = states.shape[0] frames = [] - for k in range(traj_len): + if env.name == "NutAssemblySquare": + sub_tasks = [ + "Reach for the Square Nut", + "Grasp the Square Nut", + "Align the Square Nut with the Squre Peg", + "Insert the Square Nut", + ] + + stage = 0 + + for k in range(traj_len): obs = env.reset_to({"states": states[k]}) frame = obs[IMAGE_VIEW_TO_RECORD] frames.append(frame) @@ -119,13 +135,28 @@ def extract_trajectory(env, initial_state, states, actions, done_mode): row_data[o] = np.around(actions[k], 2).tolist() elif o == "robot0_eef_pos": row_data[o] = np.around(obs[o], 2).tolist() + elif o == "sub_task": + pass else: row_data[o] = np.around( query_sim_for_data(env, desired_obs=o), 2 ).tolist() + + # Advange the stage + if env.name == "NutAssemblySquare": + if stage == 0 and actions[k][6] > 0: + stage = 1 + elif stage == 1 and row_data["nut0_pos"][2] > 0.83: + stage = 2 + elif stage == 2 and actions[k][6] < 0: + stage = 3 + + # Infer sub_task + row_data["sub_task"] = sub_tasks[stage] + df.loc[k] = row_data - print(" Done Running Simulation.") + print("Done Running Simulation.") return df, frames @@ -188,7 +219,7 @@ def record_dataset(args): # save data filename = env.name + f"_{idx}" save_video_fn(frames=frames, filename=filename) if save_video else None - save_df_to_txt(df=df, filename=filename) if save_txt else None + save_df_to_txt(df=df, filename=filename, kind="gt") if save_txt else None if __name__ == "__main__": From 0a303a2b75ad6fc2e68747869abf48dd6745f70c Mon Sep 17 00:00:00 2001 From: Ruihan Zhao Date: Sun, 25 Feb 2024 15:04:16 -0600 Subject: [PATCH 2/4] Added Tool Hang. --- .../scripts/record_robomimic_data.py | 28 ++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/task_decomposition/scripts/record_robomimic_data.py b/task_decomposition/scripts/record_robomimic_data.py index bbe6902..4bdea77 100644 --- a/task_decomposition/scripts/record_robomimic_data.py +++ b/task_decomposition/scripts/record_robomimic_data.py @@ -106,8 +106,19 @@ def extract_trajectory(env, initial_state, states, actions, done_mode): "Align the Square Nut with the Squre Peg", "Insert the Square Nut", ] + elif env.name == "ToolHang": + sub_tasks = [ + "Reach for the Frame", + "Grasp the Frame", + "Align the Frame with the Stand", + "Insert the Frame into the Stand", + "Reach for the Tool", + "Grasp the Tool", + "Align the Tool with the Frame", + "Hang the Tool", + ] - stage = 0 + stage = 0 for k in range(traj_len): obs = env.reset_to({"states": states[k]}) @@ -150,6 +161,21 @@ def extract_trajectory(env, initial_state, states, actions, done_mode): stage = 2 elif stage == 2 and actions[k][6] < 0: stage = 3 + elif env.name == "ToolHang": + if stage == 0 and actions[k][6] > 0: + stage = 1 + elif stage == 1 and row_data["frame_pos"][2] > 0.81: + stage = 2 + elif stage == 2 and actions[k][6] < 0: + stage = 3 + elif stage == 3 and row_data["frame_pos"][2] < 1: + stage = 4 + elif stage == 4 and actions[k][6] > 0: + stage = 5 + elif stage == 5 and row_data["tool_pos"][2] > 0.81: + stage = 6 + elif stage == 6 and actions[k][6] < 0: + stage = 7 # Infer sub_task row_data["sub_task"] = sub_tasks[stage] From c21ab162953f5f33dbe56600b25f5872c16854dc Mon Sep 17 00:00:00 2001 From: jsalfity Date: Sun, 25 Feb 2024 15:45:21 -0600 Subject: [PATCH 3/4] change to save two dfs --- .../scripts/record_robomimic_data.py | 44 ++++++++++++------- 1 file changed, 28 insertions(+), 16 deletions(-) diff --git a/task_decomposition/scripts/record_robomimic_data.py b/task_decomposition/scripts/record_robomimic_data.py index 4bdea77..16757b2 100644 --- a/task_decomposition/scripts/record_robomimic_data.py +++ b/task_decomposition/scripts/record_robomimic_data.py @@ -39,9 +39,7 @@ def get_data_to_record(env_name: str): else: raise ValueError(f"Environment {env_name} has no defined data to record.") - return ( - ["step", "sub_task"] + actions_to_record + obs_to_record + meta_data_to_record - ) + return ["step"] + actions_to_record + obs_to_record + meta_data_to_record def query_sim_for_data(env, desired_obs): @@ -95,19 +93,20 @@ def extract_trajectory(env, initial_state, states, actions, done_mode): data_to_record = get_data_to_record(env_name=env.name) df = pd.DataFrame(columns=get_data_to_record(env_name=env.name)) + gt_df = pd.DataFrame(columns=["step", "subtask", "stage"]) traj_len = states.shape[0] frames = [] if env.name == "NutAssemblySquare": - sub_tasks = [ + subtask_list = [ "Reach for the Square Nut", "Grasp the Square Nut", "Align the Square Nut with the Squre Peg", "Insert the Square Nut", ] elif env.name == "ToolHang": - sub_tasks = [ + subtask_list = [ "Reach for the Frame", "Grasp the Frame", "Align the Frame with the Stand", @@ -153,7 +152,9 @@ def extract_trajectory(env, initial_state, states, actions, done_mode): query_sim_for_data(env, desired_obs=o), 2 ).tolist() - # Advange the stage + df.loc[k] = row_data + + # Advange the stage for ground truth label if env.name == "NutAssemblySquare": if stage == 0 and actions[k][6] > 0: stage = 1 @@ -177,13 +178,10 @@ def extract_trajectory(env, initial_state, states, actions, done_mode): elif stage == 6 and actions[k][6] < 0: stage = 7 - # Infer sub_task - row_data["sub_task"] = sub_tasks[stage] + gt_row_data = {"step": k, "subtask": subtask_list[stage], "stage": stage} + gt_df.loc[k] = gt_row_data - df.loc[k] = row_data - - print("Done Running Simulation.") - return df, frames + return df, gt_df, frames def record_dataset(args): @@ -200,6 +198,7 @@ def record_dataset(args): ) save_txt = True if args.save_txt == 1 else False + save_gt = True if args.save_gt == 1 else False save_video = True if args.save_video == 1 else False print("==== Using environment with the following metadata ====") print(json.dumps(env.serialize(), indent=4)) @@ -234,7 +233,7 @@ def record_dataset(args): actions = f["data/{}/actions".format(ep)][()] timenow = datetime.now().strftime("%Y%m%d-%H%M%S") idx = timenow + f"_{idx}" - df, frames = extract_trajectory( + df, gt_df, frames = extract_trajectory( env=env, initial_state=initial_state, states=states, @@ -244,8 +243,14 @@ def record_dataset(args): # save data filename = env.name + f"_{idx}" + print(" Done Running Simulation.") save_video_fn(frames=frames, filename=filename) if save_video else None - save_df_to_txt(df=df, filename=filename, kind="gt") if save_txt else None + save_df_to_txt(df=df, filename=filename, kind="raw") if save_txt else None + ( + save_df_to_txt(df=gt_df, filename=filename + "_gt", kind="gt") + if save_gt + else None + ) if __name__ == "__main__": @@ -253,8 +258,8 @@ def record_dataset(args): parser.add_argument( "--dataset", type=str, - default=None, - required=True, + default="/Users/jonathansalfity/Documents/dev/task_decomposition/task_decomposition/data/robomimic/tool_hang/demo_v141.hdf5", + # required=True, help="path to input hdf5 dataset", ) @@ -272,6 +277,13 @@ def record_dataset(args): help="(Required) but default to True, save txt files", ) + parser.add_argument( + "--save_gt", + type=int, + default=1, + help="(Required) but default to True, save ground truth files", + ) + parser.add_argument( "--save_video", type=int, From 971ec3c5035b09832981dbfbd1ab8f77bb1368c0 Mon Sep 17 00:00:00 2001 From: jsalfity Date: Sun, 25 Feb 2024 16:33:24 -0600 Subject: [PATCH 4/4] robomimic data is saving --- task_decomposition/scripts/record_robomimic_data.py | 4 ++-- task_decomposition/utils/logging.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/task_decomposition/scripts/record_robomimic_data.py b/task_decomposition/scripts/record_robomimic_data.py index 16757b2..39989fb 100644 --- a/task_decomposition/scripts/record_robomimic_data.py +++ b/task_decomposition/scripts/record_robomimic_data.py @@ -258,8 +258,8 @@ def record_dataset(args): parser.add_argument( "--dataset", type=str, - default="/Users/jonathansalfity/Documents/dev/task_decomposition/task_decomposition/data/robomimic/tool_hang/demo_v141.hdf5", - # required=True, + default=None, # Example: "/Users/jonathansalfity/Documents/dev/task_decomposition/task_decomposition/data/robomimic/tool_hang/demo_v141.hdf5", + required=True, help="path to input hdf5 dataset", ) diff --git a/task_decomposition/utils/logging.py b/task_decomposition/utils/logging.py index 0015353..19d7af1 100644 --- a/task_decomposition/utils/logging.py +++ b/task_decomposition/utils/logging.py @@ -21,9 +21,10 @@ def make_env(config): return suite.make(**config) -def save_df_to_txt(df: pd.DataFrame, filename, kind): +def save_df_to_txt(df: pd.DataFrame, filename: str, kind: str): """ Dump pandas dataframe to file + kind: "raw" or "gt" """ if kind == "raw": savepath = DATA_RAW_TXT_PATH