Remove batch length restrictions in select_action (#123)

lerobot/common/policies/act/modeling_act.py

@@ -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
         queue is empty.
         """
+        assert "observation.images.top" in batch
+        assert "observation.state" in batch
+
         self.eval()
 
         batch = self.normalize_inputs(batch)

lerobot/common/policies/diffusion/modeling_diffusion.py

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

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)