From a7ff00b0cd7b87cc97d574f6ee6f2570cb335765 Mon Sep 17 00:00:00 2001 From: Nicklas Hansen Date: Thu, 4 Jan 2024 19:17:47 -0800 Subject: [PATCH] add option to disable planning with mpc=false --- tdmpc2/tdmpc2.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tdmpc2/tdmpc2.py b/tdmpc2/tdmpc2.py index 3925359..9d49cf8 100755 --- a/tdmpc2/tdmpc2.py +++ b/tdmpc2/tdmpc2.py @@ -85,7 +85,10 @@ class TDMPC2: 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()