cleanup.py

import os
import glob
if __name__ == '__main__':
    task = 'TD'
    log_dirs = glob.glob(os.path.join(os.path.dirname(__file__), 'logs', f'final_{task}_*'))
    for log_dir in log_dirs:
        for seed in range(10):
            ckpt_path = os.path.join(log_dir, str(seed), 'checkpoints')
            for ckpt in range(1000, 50_001, 500):
                if ckpt % 2000 == 0:
                    continue
                file_path = os.path.join(ckpt_path, f'ckpt_{ckpt}.pt')
                if os.path.exists(file_path):
                    os.remove(file_path)
                    print(f"{file_path} has been deleted.")
                else:
                    print(f"{file_path} does not exist.")