Co-authored-by: Simon Alibert <simon.alibert@huggingface.co>
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
Co-authored-by: Pablo <pablo.montalvo.leroux@gmail.com>
This commit is contained in:
Remi
2025-02-04 18:01:04 +01:00
committed by GitHub
parent dd974529cf
commit 638d411cd3
26 changed files with 2365 additions and 92 deletions

View File

@@ -319,7 +319,7 @@ class TDMPCPolicy(PreTrainedPolicy):
# (b, t) -> (t, b)
for key in batch:
if batch[key].ndim > 1:
if isinstance(batch[key], torch.Tensor) and batch[key].ndim > 1:
batch[key] = batch[key].transpose(1, 0)
action = batch["action"] # (t, b, action_dim)
@@ -502,7 +502,7 @@ class TDMPCPolicy(PreTrainedPolicy):
# Undo (b, t) -> (t, b).
for key in batch:
if batch[key].ndim > 1:
if isinstance(batch[key], torch.Tensor) and batch[key].ndim > 1:
batch[key] = batch[key].transpose(1, 0)
return info