tangger/lerobot
1
0
Fork 1
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add Pi0 (#681)

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

4
lerobot/common/policies/tdmpc/modeling_tdmpc.py
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