fix(vlabench): align env obs/action with HF dataset conventions

The VLABench HF datasets (VLABench/vlabench_primitive_ft_lerobot_video
and _composite_) log 7D ee-frame tuples that don't match what the env
was producing/consuming, so a SmolVLA checkpoint trained on
lerobot/vlabench_unified saw garbage state and its actions drove the
arm with wrong IK targets — hence the "aggressive / IK-not-used"
behavior at eval time.

Three concrete mismatches fixed:

1. Position frame. Dataset positions (both `observation.state[:3]` and
   `actions[:3]`) are in robot-base frame — convert_to_lerobot.py does
   `ee_pos -= robot_frame_pos` for state, and the trajectory column is
   stored in the same frame. The env was:
   - exposing raw world-frame ee_pos as `agent_pos`
   - handing the policy's robot-frame action to `qpos_from_site_pose`
     which expects world-frame targets
   We now cache the robot base xyz on env build and add/subtract it at
   the obs/action boundary.

2. Orientation layout. Dataset euler comes from
   VLABench/utils/utils.py::quaternion_to_euler which is
   `scipy.Rotation.from_quat(...).as_euler('xyz')` (extrinsic xyz).
   The env was exposing the raw `[qw,qx,qy,qz]` quaternion as
   `agent_pos[3:7]` and filling `agent_pos[6]` with `qz`, so the
   policy saw a 7D state that doesn't exist anywhere in the training
   data. Now `_get_obs` emits `[pos_robot(3), euler_xyz(3), gripper(1)]`
   and `_build_ctrl_from_action` parses actions the same way.
   Replace the hand-rolled extrinsic-xyz->wxyz helper with
   `scipy.Rotation.from_euler('xyz', ...).as_quat()` (reorder to wxyz
   for dm_control) — exactly the formula VLABench uses upstream.

3. Gripper. VLABench's convert_to_lerobot.py encodes action gripper as
   `1 if finger_qpos > 0.03 else 0` — so dataset **1 = OPEN**, 0 =
   CLOSED. The env was using the opposite convention (1=closed) and
   mapping `finger_qpos = FINGER_OPEN + g*(CLOSED-OPEN)`, which flipped
   the gripper on every frame. Now `finger_qpos = g * FINGER_OPEN`
   (g=1 -> 0.04 open, g=0 -> 0.0 closed).

The `_get_obs` ee_state layout now mirrors the dataset exactly, so a
SmolVLA checkpoint trained on `lerobot/vlabench_unified` will consume
`observation.state` it actually saw during training and emit actions
in the frame the env expects.

Made-with: Cursor

src/lerobot/envs/vlabench.py

@@ -35,6 +35,7 @@ import cv2
 import gymnasium as gym
 import numpy as np
 from gymnasium import spaces
+from scipy.spatial.transform import Rotation as R
 
 from lerobot.types import RobotObservation
 
@@ -141,6 +142,12 @@ class VLABenchEnv(gym.Env):
         # refetch it via `self._env.physics` at the call site.
         self._env = None
         self.task_description = ""  # populated on first reset
+        # Cached world-frame XYZ of the robot base link. The VLABench datasets
+        # log both `observation.state` positions and `actions` positions in
+        # robot-base frame (see VLABench/scripts/convert_to_lerobot.py which
+        # subtracts `robot_frame_pos` from ee_pos). The robot is attached at a
+        # fixed offset per task so this is safe to cache once per env build.
+        self._robot_base_xyz: np.ndarray | None = None
 
         h, w = self.render_resolution
 
@@ -249,6 +256,17 @@ class VLABenchEnv(gym.Env):
         else:
             self.task_description = self.task
 
+        # Cache robot base world position so `_build_ctrl_from_action` and
+        # `_get_obs` can translate between robot-frame (dataset) and
+        # world-frame (dm_control) without hitting physics every call.
+        try:
+            self._robot_base_xyz = np.asarray(
+                self._env.get_robot_frame_position(), dtype=np.float64
+            ).reshape(3)
+        except Exception:
+            # Fallback to VLABench's default Franka base position.
+            self._robot_base_xyz = np.array([0.0, -0.4, 0.78], dtype=np.float64)
+
     def _get_obs(self) -> dict:
         """Get current observation from the environment."""
         assert self._env is not None
@@ -298,14 +316,29 @@ class VLABenchEnv(gym.Env):
             else:
                 images[key] = np.zeros((h, w, 3), dtype=np.uint8)
 
-        # Extract end-effector state — coerce to exactly (7,) so vector env concat
+        # Convert VLABench's raw ee_state `[pos_world(3), quat_wxyz(4), open(1)]`
-        # doesn't fail with shape-mismatch on buffer np.stack.
+        # to the dataset's observation.state layout `[pos_robot(3), euler_xyz(3),
-        ee_state = obs.get("ee_state", np.zeros(7, dtype=np.float64))
+        # gripper(1)]`. See VLABench/scripts/convert_to_lerobot.py — positions
-        ee_state = np.asarray(ee_state, dtype=np.float64).ravel()
+        # are stored in robot-base frame and orientations as scipy extrinsic
-        if ee_state.shape[0] != 7:
+        # 'xyz' euler angles.
-            fixed = np.zeros(7, dtype=np.float64)
+        raw = np.asarray(obs.get("ee_state", np.zeros(8)), dtype=np.float64).ravel()
-            fixed[: min(7, ee_state.shape[0])] = ee_state[:7]
+        pos_world = raw[:3] if raw.size >= 3 else np.zeros(3, dtype=np.float64)
-            ee_state = fixed
+        quat_wxyz = (
+            raw[3:7] if raw.size >= 7 else np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64)
+        )
+        gripper = float(raw[7]) if raw.size >= 8 else 0.0
+
+        base = (
+            self._robot_base_xyz
+            if self._robot_base_xyz is not None
+            else np.zeros(3, dtype=np.float64)
+        )
+        pos_robot = pos_world - base
+        euler_xyz = R.from_quat(
+            [quat_wxyz[1], quat_wxyz[2], quat_wxyz[3], quat_wxyz[0]]
+        ).as_euler("xyz", degrees=False)
+
+        ee_state = np.concatenate([pos_robot, euler_xyz, [gripper]]).astype(np.float64)
 
         if self.obs_type == "pixels":
             return {"pixels": images}
@@ -319,30 +352,19 @@ class VLABenchEnv(gym.Env):
 
     # ---- Action adaptation (EEF → joint ctrl) --------------------------------
     #
-    # The dataset logs 7D actions [x, y, z, rx, ry, rz, gripper] in end-effector
+    # The HF vlabench datasets log 7D actions
-    # space, but VLABench's dm_control task writes `data.ctrl[:] = action`
+    # `[x, y, z (robot frame), rx, ry, rz (scipy extrinsic xyz), gripper]`,
-    # directly — which for Franka expects 9 entries (7 arm joints + 2 gripper
+    # exactly matching VLABench's own eval pipeline (evaluator.base):
-    # fingers). Reproduce the same EEF-to-joint conversion that VLABench's data
+    #   pos, euler, g = policy(...)
-    # collector uses (`SingleArm.get_qpos_from_ee_pos`, a dm_control IK solver)
+    #   quat = euler_to_quaternion(*euler)      # extrinsic xyz -> wxyz
-    # so the policy's EEF commands actually drive the robot at eval time.
+    #   _, qpos = robot.get_qpos_from_ee_pos(physics, pos=pos + base, quat=quat)
+    #   env.step(np.concatenate([qpos, [g, g]]))
+    #
+    # VLABench's dm_control task writes `data.ctrl[:] = action` directly — for
+    # Franka that's 9 entries (7 arm joints + 2 gripper fingers). We mirror the
+    # above conversion so the policy's EEF commands actually drive the robot.
 
     _FRANKA_FINGER_OPEN = 0.04  # qpos when gripper fully open
-    _FRANKA_FINGER_CLOSED = 0.0
-
-    @staticmethod
-    def _euler_xyz_to_quat_wxyz(rx: float, ry: float, rz: float) -> np.ndarray:
-        """Euler (XYZ intrinsic) → quaternion (w, x, y, z) — dm_control convention."""
-        cx, cy, cz = np.cos(0.5 * np.array([rx, ry, rz]))
-        sx, sy, sz = np.sin(0.5 * np.array([rx, ry, rz]))
-        return np.array(
-            [
-                cx * cy * cz + sx * sy * sz,
-                sx * cy * cz - cx * sy * sz,
-                cx * sy * cz + sx * cy * sz,
-                cx * cy * sz - sx * sy * cz,
-            ],
-            dtype=np.float64,
-        )
 
     def _build_ctrl_from_action(self, action: np.ndarray, ctrl_dim: int) -> np.ndarray:
         """Convert a 7D EEF action into the `ctrl_dim`-sized joint command vector.
@@ -363,22 +385,34 @@ class VLABenchEnv(gym.Env):
 
         from dm_control.utils.inverse_kinematics import qpos_from_site_pose
 
-        pos = np.asarray(action[:3], dtype=np.float64)
+        # Action position is in robot-base frame (see convert_to_lerobot.py);
+        # dm_control's IK expects a world-frame target.
+        base = (
+            self._robot_base_xyz
+            if self._robot_base_xyz is not None
+            else np.zeros(3, dtype=np.float64)
+        )
+        pos_world = np.asarray(action[:3], dtype=np.float64) + base
         rx, ry, rz = float(action[3]), float(action[4]), float(action[5])
         gripper = float(np.clip(action[6], 0.0, 1.0))
-        quat = self._euler_xyz_to_quat_wxyz(rx, ry, rz)
+
+        # Dataset euler is scipy extrinsic 'xyz' (same as VLABench's
+        # `euler_to_quaternion`). scipy emits `[x, y, z, w]`; dm_control's IK
+        # and MuJoCo use `[w, x, y, z]`, so reorder.
+        qxyzw = R.from_euler("xyz", [rx, ry, rz], degrees=False).as_quat()
+        quat = np.array([qxyzw[3], qxyzw[0], qxyzw[1], qxyzw[2]], dtype=np.float64)
 
         assert self._env is not None
         robot = self._env.task.robot
         site_name = robot.end_effector_site.full_identifier
 
-        # Important: inplace=False so IK doesn't mutate physics state mid-step;
+        # inplace=False so IK doesn't mutate physics state mid-step — we only
-        # we only want the solved qpos. Fetch a fresh physics handle — caching
+        # want the solved qpos. Fetch a fresh physics handle — caching it can
-        # it can yield a stale weakref after a reset.
+        # yield a stale weakref after a reset.
         ik_result = qpos_from_site_pose(
             self._env.physics,
             site_name=site_name,
-            target_pos=pos,
+            target_pos=pos_world,
             target_quat=quat,
             inplace=False,
             max_steps=100,
@@ -386,11 +420,10 @@ class VLABenchEnv(gym.Env):
         n_dof = robot.n_dof  # 7 for Franka
         arm_qpos = ik_result.qpos[:n_dof]
 
-        # Gripper: action scalar in [0, 1] (0=open, 1=closed). Map linearly to
+        # Dataset gripper convention: 1 = open (finger qpos = 0.04),
-        # finger qpos in [CLOSED, OPEN]. Franka has 2 mirrored fingers.
+        # 0 = closed (finger qpos = 0.0). See VLABench/scripts/convert_to_lerobot.py
-        finger_qpos = self._FRANKA_FINGER_OPEN + gripper * (
+        # where `trajectory[i][-1] > 0.03` is encoded as `1`.
-            self._FRANKA_FINGER_CLOSED - self._FRANKA_FINGER_OPEN
+        finger_qpos = gripper * self._FRANKA_FINGER_OPEN
-        )
 
         ctrl = np.zeros(ctrl_dim, dtype=np.float64)
         ctrl[:n_dof] = arm_qpos