import os
import torch

_ = "Exp-20220112-145213-OLKnapsackDecision"

for sub in ["warmup"]:
    ckpt_path = os.path.join(_, "checkpoint", sub)
    ckpts = os.listdir(ckpt_path)
    for ckpt in ckpts:
        ckpt_package = torch.load(os.path.join(ckpt_path, ckpt), map_location="cpu")
        for env in ckpt_package["envs"]:
            try:
                env.clean()
            except:
                print("1")
        torch.save(ckpt_package, os.path.join(ckpt_path, ckpt))
