add separate eval model class that also initializes the pre-trained policy for early denoising steps

2025-02-04 11:51:47 -05:00
parent e7f73dffc1
commit ace2bbdab9
4 changed files with 138 additions and 2 deletions
@@ -128,7 +128,7 @@ class TrainAgent:
        data = {
            "itr": self.itr,
            "model": self.model.state_dict(),
-        }
+        }  # right now `model` includes weights for `network`, `actor`, `actor_ft`. Weights for `network` is redundant, and we can use `actor` weights as the base policy (earlier denoising steps) and `actor_ft` weights as the fine-tuned policy (later denoising steps) during evaluation.
        savepath = os.path.join(self.checkpoint_dir, f"state_{self.itr}.pt")
        torch.save(data, savepath)
        log.info(f"Saved model to {savepath}")