fix(vlabench): resize images to declared observation_space shape

VLABench's load_env doesn't always honor render_resolution, so the
returned rgb frames may not match our declared (h, w, 3) observation
space. gymnasium's SyncVectorEnv.concatenate then raises
"Output array is the wrong shape". Resize to the declared (h, w).

Made-with: Cursor
This commit is contained in:
Pepijn
2026-04-16 21:13:18 +02:00
parent 5161d06f32
commit da7fb54c64
+6 -1
View File
@@ -237,11 +237,16 @@ class VLABenchEnv(gym.Env):
else: # HWC else: # HWC
images["image"] = rgb.astype(np.uint8) images["image"] = rgb.astype(np.uint8)
# Fill missing cameras with zeros # VLABench's native render size may not match our declared observation_space.
# Resize each image to (h, w) so gymnasium's vector env concatenate succeeds.
h, w = self.render_resolution h, w = self.render_resolution
for key in ["image", "second_image", "wrist_image"]: for key in ["image", "second_image", "wrist_image"]:
if key not in images: if key not in images:
images[key] = np.zeros((h, w, 3), dtype=np.uint8) images[key] = np.zeros((h, w, 3), dtype=np.uint8)
elif images[key].shape[:2] != (h, w):
import cv2
images[key] = cv2.resize(images[key], (w, h), interpolation=cv2.INTER_AREA).astype(np.uint8)
# Extract end-effector state # Extract end-effector state
ee_state = obs.get("ee_state", np.zeros(7, dtype=np.float64)) ee_state = obs.get("ee_state", np.zeros(7, dtype=np.float64))