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

separate episodes with nans

Browse Source
This commit is contained in:
Nicklas Hansen
2024-01-07 14:21:38 -08:00
parent 33876d124f
commit 31249a8961
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tdmpc2/trainer/online_trainer.py
View File

@@ -47,8 +47,10 @@ class OnlineTrainer(Trainer):
episode_success=np.nanmean(ep_successes), episode_success=np.nanmean(ep_successes),
) )
def to_td(self, obs, action=None, reward=None): def to_td(self, obs=None, action=None, reward=None):
"""Creates a TensorDict for a new episode.""" """Creates a TensorDict for a new episode."""
if obs is None:
obs = torch.full((*self.cfg.obs_shape[self.cfg.obs],), float('nan'))
if isinstance(obs, dict): if isinstance(obs, dict):
obs = TensorDict(obs, batch_size=(), device='cpu') obs = TensorDict(obs, batch_size=(), device='cpu')
else: else:
@@ -88,6 +90,7 @@ class OnlineTrainer(Trainer):
) )
train_metrics.update(self.common_metrics()) train_metrics.update(self.common_metrics())
self.logger.log(train_metrics, 'train') self.logger.log(train_metrics, 'train')
self._tds.append(self.to_td()) # Separate episodes with NaNs
self._ep_idx = self.buffer.add(torch.cat(self._tds)) self._ep_idx = self.buffer.add(torch.cat(self._tds))
obs = self.env.reset() obs = self.env.reset()