wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Remove batch length restrictions in select_action (#123)

This commit is contained in:
Simon Alibert 2024-05-04 15:33:55 +02:00 committed by GitHub
parent bccee745c3
commit c015252e20
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
lerobot/common/policies/act/modeling_act.py
View File

@ -71,6 +71,9 @@ class ACTPolicy(nn.Module, PyTorchModelHubMixin):
environment. It works by managing the actions in a queue and only calling `select_actions` when the environment. It works by managing the actions in a queue and only calling `select_actions` when the
queue is empty. queue is empty.
""" """
assert "observation.images.top" in batch
assert "observation.state" in batch
self.eval() self.eval()
batch = self.normalize_inputs(batch) batch = self.normalize_inputs(batch)

1
lerobot/common/policies/diffusion/modeling_diffusion.py
View File

@ -115,7 +115,6 @@ class DiffusionPolicy(nn.Module, PyTorchModelHubMixin):
""" """
assert "observation.image" in batch assert "observation.image" in batch
assert "observation.state" in batch assert "observation.state" in batch
assert len(batch) == 2
batch = self.normalize_inputs(batch) batch = self.normalize_inputs(batch)

1
lerobot/common/policies/tdmpc/modeling_tdmpc.py
View File

@ -122,7 +122,6 @@ class TDMPCPolicy(nn.Module, PyTorchModelHubMixin):
"""Select a single action given environment observations.""" """Select a single action given environment observations."""
assert "observation.image" in batch assert "observation.image" in batch
assert "observation.state" in batch assert "observation.state" in batch
assert len(batch) == 2
batch = self.normalize_inputs(batch) batch = self.normalize_inputs(batch)