hongshaorou/tdmpc2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add option to disable planning with mpc=false

Browse Source
This commit is contained in:
Nicklas Hansen
2024-01-04 19:17:47 -08:00
parent 13cac07759
commit a7ff00b0cd
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tdmpc2/tdmpc2.py
View File

@@ -85,7 +85,10 @@ class TDMPC2:
if task is not None: if task is not None:
task = torch.tensor([task], device=self.device) task = torch.tensor([task], device=self.device)
z = self.model.encode(obs, task) z = self.model.encode(obs, task)
if self.cfg.mpc:
a = self.plan(z, t0=t0, eval_mode=eval_mode, task=task) 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() return a.cpu()
@torch.no_grad() @torch.no_grad()