From 3837407a445054f749e43421104d9b40e56dac4e Mon Sep 17 00:00:00 2001 From: ubless607 Date: Thu, 23 Jan 2025 22:18:02 +0900 Subject: [PATCH] Update rearrange_rlbench_demos.py Added AssertionError exception. If any of the pickle file is corrupted, it keeps generating recursive folder over and over. This fix will skip the episode with the problem. --- data_preprocessing/rearrange_rlbench_demos.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/data_preprocessing/rearrange_rlbench_demos.py b/data_preprocessing/rearrange_rlbench_demos.py index 8b64ae99..c1e38981 100644 --- a/data_preprocessing/rearrange_rlbench_demos.py +++ b/data_preprocessing/rearrange_rlbench_demos.py @@ -31,7 +31,14 @@ def main(root_dir, task): if os.path.isfile(f'{root_dir}/{task}/variation{variation}/variation_descriptions.pkl'): data1 = pickle.load(open(f'{root_dir}/{task}/all_variations/episodes/episode{num}/variation_descriptions.pkl', 'rb')) data2 = pickle.load(open(f'{root_dir}/{task}/variation{variation}/variation_descriptions.pkl', 'rb')) - assert data1 == data2 + try: + assert data1 == data2 + print(f"✓ Success: Variation {variation}, Episode {num} - descriptions match") + except AssertionError: + print(f"✗ Error: Mismatch in descriptions for Variation {variation}, Episode {num}") + print(f" Source data: {data1}") + print(f" Target data: {data2}") + continue else: call(['ln', '-s', f'{root_dir}/{task}/all_variations/episodes/episode{num}/variation_descriptions.pkl',