Remove batch length restrictions in select_action (#123)

2024-05-04 15:33:55 +02:00
parent bccee745c3
commit c015252e20
3 changed files with 3 additions and 2 deletions
--- a/lerobot/common/policies/tdmpc/modeling_tdmpc.py
+++ b/lerobot/common/policies/tdmpc/modeling_tdmpc.py
@@ -122,7 +122,6 @@ class TDMPCPolicy(nn.Module, PyTorchModelHubMixin):
        """Select a single action given environment observations."""
        assert "observation.image" in batch
        assert "observation.state" in batch
-        assert len(batch) == 2

        batch = self.normalize_inputs(batch)