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

Update offline_trainer.py

Browse Source
This commit is contained in:
Vincent Moens
2024-10-26 00:33:27 +01:00
committed by GitHub
parent fad0d1be03
commit 3b5f67592c
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
tdmpc2/trainer/offline_trainer.py
View File

@@ -27,6 +27,7 @@ class OfflineTrainer(Trainer):
for _ in range(self.cfg.eval_episodes): for _ in range(self.cfg.eval_episodes):
obs, done, ep_reward, t = self.env.reset(task_idx), False, 0, 0 obs, done, ep_reward, t = self.env.reset(task_idx), False, 0, 0
while not done: while not done:
torch.compiler.cudagraph_mark_step_begin()
action = self.agent.act(obs, t0=t==0, eval_mode=True, task=task_idx) action = self.agent.act(obs, t0=t==0, eval_mode=True, task=task_idx)
obs, reward, done, info = self.env.step(action) obs, reward, done, info = self.env.step(action)
ep_reward += reward ep_reward += reward