fix(tasks): index Franka gripper/arm joints by sorted-name order in ManiSkill rewards (#774)

MetaSim reports RobotState.joint_pos/joint_vel in alphabetically-sorted
joint-name order. For the Franka Panda that is [panda_finger_joint1,
panda_finger_joint2, panda_joint1..7] -> fingers at columns 0-1, arm at 2-8
(franka_cfg.py ee_joint_indices=[0,1]). The dense rewards sliced joint_pos[:,7:9]
for the gripper and joint_vel[:,:7] for the arm (URDF order), so is_grasped,
gripper_width and the static/velocity terms read the wrong joints on every
backend. Correct to [:,0:2] (gripper) and [:,2:9] (arm) in pick_cube_v1,
lift_peg_upright, place_sphere, poke_cube, pick_single_ycb, stack_cube.

The _native/ bitwise-parity rewards are fed raw native-order qpos and are
intentionally left untouched.

Adds a backend-free regression test (constructs RobotState/ObjectState and
calls the pure _reward): grasp reward must respond to the finger columns, plus
a source guard against the URDF-order slice across all six files. Verified the
test is red on the pre-fix slices and green after.

Author: geng-haoran

roboverse_pack/tasks/maniskill/lift_peg_upright.py

@@ -119,7 +119,7 @@ def _reward(self, env_states):
 
         to_grip = torch.linalg.norm(peg - tcp, dim=-1)
         reaching = 1.0 - torch.tanh(5.0 * to_grip)
-        gripper_width = rs.joint_pos[:, 7:9].sum(dim=-1)
+        gripper_width = rs.joint_pos[:, 0:2].sum(dim=-1)
         is_grasping = (to_grip < 0.04) & (gripper_width < 0.07)
         reaching = torch.where(is_grasping, torch.ones_like(reaching), reaching)
         reaching = reaching / 5.0

roboverse_pack/tasks/maniskill/pick_cube_v1.py

@@ -153,7 +153,7 @@ def _reward(self, env_states):
         reaching = 1.0 - torch.tanh(5.0 * tcp_to_obj)
         reward = reaching
 
-        finger_q = rs.joint_pos[:, 7:9]
+        finger_q = rs.joint_pos[:, 0:2]
         gripper_width = finger_q.sum(dim=-1)
         is_grasped = ((tcp_to_obj < 0.03) & (gripper_width < 0.07)).float()
         reward = reward + is_grasped
@@ -162,7 +162,7 @@ def _reward(self, env_states):
         place = 1.0 - torch.tanh(5.0 * obj_to_goal)
         reward = reward + place * is_grasped
 
-        arm_qvel = rs.joint_vel[:, :7]
+        arm_qvel = rs.joint_vel[:, 2:9]
         static = 1.0 - torch.tanh(5.0 * torch.linalg.norm(arm_qvel, dim=-1))
         is_obj_placed = (obj_to_goal <= _GOAL_THRESH).float()
         reward = reward + static * is_obj_placed

roboverse_pack/tasks/maniskill/pick_single_ycb.py

@@ -98,7 +98,7 @@ def _reward(self, env_states):
 
         tcp_to_obj = torch.linalg.norm(obj - tcp, dim=-1)
         reaching = 1.0 - torch.tanh(5.0 * tcp_to_obj)
-        gripper_w = rs.joint_pos[:, 7:9].sum(dim=-1)
+        gripper_w = rs.joint_pos[:, 0:2].sum(dim=-1)
         is_grasped = ((tcp_to_obj < 0.03) & (gripper_w < 0.07)).float()
         reward = reaching + is_grasped
 
@@ -109,7 +109,7 @@ def _reward(self, env_states):
         is_obj_placed = (obj_to_goal <= _YCB_GOAL_THRESH).float()
         reward = reward + is_obj_placed * is_grasped
 
-        arm_qvel = rs.joint_vel[:, :7]
+        arm_qvel = rs.joint_vel[:, 2:9]
         static = 1.0 - torch.tanh(5.0 * torch.linalg.norm(arm_qvel, dim=-1))
         reward = reward + static * is_obj_placed * is_grasped
 

roboverse_pack/tasks/maniskill/place_sphere.py

@@ -157,14 +157,14 @@ def _reward(self, env_states):
 
         o_to_bin = torch.linalg.norm(bin_top - obj, dim=-1)
         place = 1.0 - torch.tanh(5.0 * o_to_bin)
-        gripper_w = rs.joint_pos[:, 7:9].sum(dim=-1)
+        gripper_w = rs.joint_pos[:, 0:2].sum(dim=-1)
         is_grasped = (o_to_tcp < 0.03) & (gripper_w < 0.07)
         reward = torch.where(is_grasped, 4.0 + place, reward)
 
-        ungrasp = rs.joint_pos[:, 7:9].sum(dim=-1) / _PS_GRIPPER_WIDTH
+        ungrasp = rs.joint_pos[:, 0:2].sum(dim=-1) / _PS_GRIPPER_WIDTH
         ungrasp = torch.where(is_grasped, ungrasp, torch.full_like(ungrasp, 16.0))
         static = 1.0 - torch.tanh(torch.linalg.norm(obj_v, dim=-1) * 10.0 + torch.linalg.norm(obj_w, dim=-1))
-        arm_qvel = rs.joint_vel[:, :7]
+        arm_qvel = rs.joint_vel[:, 2:9]
         robot_static = (arm_qvel.abs().amax(dim=-1) < 0.2).float()
 
         off = obj - objs["bin"].root_state[:, :3]

roboverse_pack/tasks/maniskill/poke_cube.py

@@ -195,15 +195,15 @@ def _reward(self, env_states):
         h2c = torch.linalg.norm(peg_head[:, :2] - cube[:, :2], dim=-1)
         close = 1.0 - torch.tanh(5.0 * h2c)
 
-        gripper_w = rs.joint_pos[:, 7:9].sum(dim=-1)
+        gripper_w = rs.joint_pos[:, 0:2].sum(dim=-1)
         grasped = (tcp_to_peg < 0.04) & (gripper_w < 0.07) & reached
         reward = torch.where(grasped, 4.0 + close + align, reward)
 
         place = 1.0 - torch.tanh(5.0 * torch.linalg.norm(goal - cube, dim=-1))
         fit = (angle < 0.05) & (h2c <= _PK_CUBE_HALF + 0.005) & grasped
         reward = torch.where(fit, 7.0 + place, reward)
 
-        arm_qvel = rs.joint_vel[:, :7]
+        arm_qvel = rs.joint_vel[:, 2:9]
         static = 1.0 - torch.tanh(5.0 * torch.linalg.norm(arm_qvel, dim=-1))
         cube_placed = torch.linalg.norm(cube[:, :2] - goal[:, :2], dim=-1) < _PK_GOAL_RADIUS
         reward = torch.where(cube_placed, reward + static, reward)

roboverse_pack/tasks/maniskill/stack_cube.py

@@ -120,11 +120,11 @@ def _reward(self, env_states):
         a_to_goal = torch.linalg.norm(goal_xyz - cubeA, dim=-1)
         place = 1.0 - torch.tanh(5.0 * a_to_goal)
 
-        gripper_w = rs.joint_pos[:, 7:9].sum(dim=-1)
+        gripper_w = rs.joint_pos[:, 0:2].sum(dim=-1)
         is_grasped = (a_to_tcp < 0.03) & (gripper_w < 0.07)
         reward = torch.where(is_grasped, 4.0 + place, reward)
 
-        ungrasp = rs.joint_pos[:, 7:9].sum(dim=-1) / _SC_GRIPPER_WIDTH
+        ungrasp = rs.joint_pos[:, 0:2].sum(dim=-1) / _SC_GRIPPER_WIDTH
         ungrasp = torch.where(is_grasped, ungrasp, torch.ones_like(ungrasp))
         static = 1.0 - torch.tanh(torch.linalg.norm(cubeA_vel, dim=-1) * 10.0 + torch.linalg.norm(cubeA_ang, dim=-1))
 

tests/test_maniskill_reward_joint_order.py

@@ -0,0 +1,103 @@
+"""Regression: ManiSkill dense rewards index Franka joints by the sorted-name contract.
+
+MetaSim reports ``RobotState.joint_pos`` / ``joint_vel`` in alphabetically-sorted
+joint-name order. For the Franka Panda that order is
+``[panda_finger_joint1, panda_finger_joint2, panda_joint1..panda_joint7]`` — i.e.
+the two gripper fingers are columns 0-1 and the seven arm joints are columns 2-8
+(see ``roboverse_pack/robots/franka_cfg.py`` ``ee_joint_indices = [0, 1]``).
+
+The dense reward functions used to slice ``joint_pos[:, 7:9]`` for the gripper and
+``joint_vel[:, :7]`` for the arm — the *URDF* declaration order — so ``is_grasped``,
+``gripper_width`` and the static/velocity terms were computed from the wrong joints
+on every backend. These tests pin the corrected indexing without needing a
+simulator backend or downloaded assets (``_reward`` is a pure function of the
+passed ``env_states``).
+
+Backend-free by design: no existing ``tests/test_maniskill_*`` file is usable here
+because they all ``pytest.importorskip("mani_skill")`` at module load.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from metasim.types import ObjectState, RobotState
+
+
+def _pick_cube_states(finger_val: float, arm_tail_val: float = 1.5):
+    """Build a minimal env_states for PickCube where the cube sits at the TCP.
+
+    Fingers occupy sorted columns 0-1; ``arm_tail_val`` is written to columns 7-8
+    (the columns the buggy code mistook for the gripper) so a regression cannot
+    pass by reading the arm. ``finger_val`` per finger: 0.01 -> closed/grasped,
+    0.04 -> open.
+    """
+    import roboverse_pack.tasks.maniskill.pick_cube_v1 as M
+
+    jp = torch.zeros(1, 9)
+    jp[0, 0] = finger_val
+    jp[0, 1] = finger_val
+    jp[0, 7] = arm_tail_val
+    jp[0, 8] = arm_tail_val
+    body_state = torch.zeros(1, 2, 13)
+    body_state[0, 1, 3] = 1.0  # panda_hand quaternion = wxyz identity, at origin
+    rs = RobotState(
+        root_state=torch.zeros(1, 13),
+        body_names=["panda_link0", "panda_hand"],
+        body_state=body_state,
+        joint_pos=jp,
+        joint_vel=torch.zeros(1, 9),
+    )
+    cube = torch.zeros(1, 13)
+    cube[0, :3] = torch.tensor(M._TCP_OFFSET)  # cube exactly at the TCP -> reaching satisfied
+    goal = torch.zeros(1, 13)
+    goal[0, 2] = 0.5  # goal far away -> not placed -> no success override
+
+    class _S:
+        pass
+
+    s = _S()
+    s.robots = {"franka": rs}
+    s.objects = {"cube": ObjectState(root_state=cube), "goal_site": ObjectState(root_state=goal)}
+    return s
+
+
+def test_pick_cube_grasp_reward_reads_finger_columns():
+    """Closing the fingers (sorted cols 0-1) must register as a grasp.
+
+    On the pre-fix code (``joint_pos[:, 7:9]``) the reward read arm columns 7-8
+    (here a large constant), so it never saw the grasp and both branches scored
+    identically — this asserts the gap that only the corrected indexing produces.
+    """
+    from roboverse_pack.tasks.maniskill.pick_cube_v1 import PickCubeV1DenseTask
+
+    grasped = float(PickCubeV1DenseTask._reward(None, _pick_cube_states(0.01))[0])
+    ungrasped = float(PickCubeV1DenseTask._reward(None, _pick_cube_states(0.04))[0])
+
+    assert grasped > ungrasped + 0.9, (
+        f"grasp (fingers closed) should raise the reward via the is_grasped term; "
+        f"got grasped={grasped:.4f} ungrasped={ungrasped:.4f}. A near-zero gap means "
+        f"the reward is reading arm columns, not the gripper fingers (cols 0-1)."
+    )
+
+
+def test_no_maniskill_reward_uses_urdf_order_finger_slice():
+    """Guard the whole dense-reward family against the URDF-order regression.
+
+    The five sibling tasks (stack_cube, poke_cube, place_sphere, pick_single_ycb,
+    lift_peg_upright) share the same Franka layout but build different object
+    states, so this asserts the corrected slice convention at the source level
+    rather than re-deriving each task's full env_states.
+    """
+    import pathlib
+
+    pack = pathlib.Path(__file__).resolve().parents[1] / "roboverse_pack" / "tasks" / "maniskill"
+    offenders = []
+    for py in pack.glob("*.py"):
+        text = py.read_text()
+        if "rs.joint_pos[:, 7:9]" in text or "rs.joint_vel[:, :7]" in text:
+            offenders.append(py.name)
+    assert not offenders, (
+        f"these files slice Franka joints in URDF order (gripper 7:9 / arm :7) "
+        f"instead of the sorted-name order (gripper 0:2 / arm 2:9): {offenders}"
+    )