set deterministic=True when sampling in diffusion evaluation

irom-princeton · Sep 26, 2024 · dd14c58 · dd14c58
1 parent 4962bbc
commit dd14c58
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/agent/eval/eval_diffusion_agent.py b/agent/eval/eval_diffusion_agent.py
@@ -50,7 +50,7 @@ def run(self):
                     .float()
                     .to(self.device)
                 }
-                samples = self.model(cond=cond)
+                samples = self.model(cond=cond, deterministic=True)
                 output_venv = (
                     samples.trajectories.cpu().numpy()
                 )  # n_env x horizon x act

diff --git a/agent/eval/eval_diffusion_img_agent.py b/agent/eval/eval_diffusion_img_agent.py
@@ -53,7 +53,7 @@ def run(self):
                     key: torch.from_numpy(prev_obs_venv[key]).float().to(self.device)
                     for key in self.obs_dims
                 }  # batch each type of obs and put into dict
-                samples = self.model(cond=cond)
+                samples = self.model(cond=cond, deterministic=True)
                 output_venv = (
                     samples.trajectories.cpu().numpy()
                 )  # n_env x horizon x act