add option to disable planning with mpc=false

YeranLee · Jan 5, 2024 · a7ff00b · a7ff00b
1 parent 13cac07
commit a7ff00b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tdmpc2/tdmpc2.py b/tdmpc2/tdmpc2.py
@@ -85,7 +85,10 @@ def act(self, obs, t0=False, eval_mode=False, task=None):
 		if task is not None:
 			task = torch.tensor([task], device=self.device)
 		z = self.model.encode(obs, task)
-		a = self.plan(z, t0=t0, eval_mode=eval_mode, task=task)
+		if self.cfg.mpc:
+			a = self.plan(z, t0=t0, eval_mode=eval_mode, task=task)
+		else:
+			a = self.model.pi(z, task)[int(not eval_mode)][0]
 		return a.cpu()
 
 	@torch.no_grad()