Change config logic in:

- gym_manipulator
- find_joint_limits
- end_effector_utils

lerobot/common/envs/utils.py

@@ -44,13 +44,18 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
         if "images" not in key:
             continue
 
-        for imgkey, img in imgs.items():
-            # TODO(aliberts, rcadene): use transforms.ToTensor()?
+        # TODO(aliberts, rcadene): use transforms.ToTensor()?
+        if not torch.is_tensor(img):
             img = torch.from_numpy(img)
 
-            # sanity check that images are channel last
-            _, h, w, c = img.shape
-            assert c < h and c < w, f"expect channel last images, but instead got {img.shape=}"
+        if img.ndim == 3:
+            img = img.unsqueeze(0)
+
+        # sanity check that images are channel last
+        _, h, w, c = img.shape
+        assert c < h and c < w, (
+            f"expect channel last images, but instead got {img.shape=}"
+        )
 
         # sanity check that images are uint8
         assert img.dtype == torch.uint8, f"expect torch.uint8, but instead {img.dtype=}"