# Created by BVH, Mar 2026. # Action parsing utilities for various dataset formats. # # Each dataset stores different action/state structures in its lowdim npz files. # After BaseWebbed unpacks the npz, sample['action'][time_key] is a Python dict # whose keys vary by dataset. This module normalizes them into a common format: # { (t, c): {'action': np.ndarray (20,), 'desired': np.ndarray (20,) or None, 'actual': np.ndarray (20,) or None} } # # The 3 output keys follow the LBM convention: # action[t] = the commanded next EE state (= desired[t+1]) # desired[t] = the desired EE state at time t (= action[t-1]) # actual[t] = the actual measured EE state at time t # see also https://tri-internal.slack.com/archives/C09BQDCDYF2/p1771519244381309?thread_ts=1771487466.877189&cid=C09BQDCDYF2 # # All datasets output a unified 20-D representation per timestep: # [0:3] right arm xyz # [3:9] right arm rot6d # [9:12] left arm xyz # [12:18] left arm rot6d # [18] right gripper # [19] left gripper # # Single-arm datasets (DROID, RH20T) populate right arm slots and zero the left arm. # Bimanual datasets (LBM, HumanoidEveryday) populate both arms. # # Per-dataset conversion details: # # Dataset Native format Conversion # ----------------- -------------------------------- ------------------------------------------- # LBM xyz(3)+rot6d(6)+grip(1) x2 = 20 Already in target format (no conversion) # DROID xyz(3)+euler_xyz(3)+grip(1) = 7 euler -> rot6d, single arm # RH20T xyz(3)+quat_xyzw(4)+grip(1) = 8 quaternion -> rot6d, single arm # HumanoidEveryday right/left_pose (4x4) + grip 4x4 homogeneous -> xyz + rot6d, bimanual # AgiBotWorld end pos(2,3)+quat(2,4)+grip(2) quaternion -> rot6d, bimanual # YAMxdof joint_raw: arm joints(6)+eef(1)x2 JOINT-space passthrough (no EE pose / FK), bimanual # (generic) varies best-effort, zero-padded to 20 # # TODO: gripper normalization across datasets (different scales/conventions). import numpy as np import torch from rich import print def _to_np(val): ''' Convert a value to a 1-D numpy array (handles scalars, lists, tensors, and existing arrays). ''' if isinstance(val, torch.Tensor): val = val.detach().cpu().numpy() arr = np.asarray(val, dtype=np.float32).ravel() return arr def _to_np_2d(val): ''' Convert a value to a 2-D numpy array (for matrices like 4x4 poses). ''' if isinstance(val, torch.Tensor): val = val.detach().cpu().numpy() return np.asarray(val, dtype=np.float32) def _current_action(action): ''' Extract the current timestep action from a potentially multi-step array/tensor. ''' if isinstance(action, torch.Tensor): action = action.detach().cpu().numpy() if action.ndim == 2: action = action[:1] # take the first timestep if multiple are provided, e.g. in legacy LBM return action # --------------------------------------------------------------------------- # Rotation conversion helpers (pure numpy, no external dependencies) # --------------------------------------------------------------------------- def _matrix_to_rot6d(R): ''' Convert a 3x3 rotation matrix to 6D representation (first 2 rows flattened). ''' return R[:2, :].ravel().astype(np.float32) # (6,) def _euler_xyz_to_rot6d(euler): ''' Convert XYZ extrinsic euler angles to rot6d. euler: (3,) array [roll_x, pitch_y, yaw_z]. Convention: R = Rz @ Ry @ Rx (extrinsic XYZ = intrinsic ZYX). ''' rx, ry, rz = float(euler[0]), float(euler[1]), float(euler[2]) cx, sx = np.cos(rx), np.sin(rx) cy, sy = np.cos(ry), np.sin(ry) cz, sz = np.cos(rz), np.sin(rz) Rx = np.array([[1, 0, 0], [0, cx, -sx], [0, sx, cx]], dtype=np.float32) Ry = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]], dtype=np.float32) Rz = np.array([[cz, -sz, 0], [sz, cz, 0], [0, 0, 1]], dtype=np.float32) R = Rz @ Ry @ Rx return _matrix_to_rot6d(R) def _quat_xyzw_to_rot6d(quat): ''' Convert quaternion [x, y, z, w] to rot6d via rotation matrix. Uses the standard quaternion-to-matrix formula (Hamilton convention). ''' x, y, z, w = float(quat[0]), float(quat[1]), float(quat[2]), float(quat[3]) # Normalize n = np.sqrt(x*x + y*y + z*z + w*w) if n < 1e-10: return np.zeros(6, dtype=np.float32) x, y, z, w = x/n, y/n, z/n, w/n R = np.array([ [1 - 2*(y*y + z*z), 2*(x*y - z*w), 2*(x*z + y*w)], [ 2*(x*y + z*w), 1 - 2*(x*x + z*z), 2*(y*z - x*w)], [ 2*(x*z - y*w), 2*(y*z + x*w), 1 - 2*(x*x + y*y)], ], dtype=np.float32) return _matrix_to_rot6d(R) def _homogeneous_to_xyz_rot6d(T): ''' Extract xyz and rot6d from a 4x4 homogeneous transform matrix. Returns (xyz (3,), rot6d (6,)). ''' T = np.asarray(T, dtype=np.float32) xyz = T[:3, 3].copy() rot6d = _matrix_to_rot6d(T[:3, :3]) return xyz, rot6d def _pack_single_arm(xyz, rot6d, grip): ''' Pack a single-arm (right arm) action into the unified 20D format. Left arm slots are zeroed. Gripper goes to index 18 (right). ''' out = np.zeros(20, dtype=np.float32) out[0:3] = xyz out[3:9] = rot6d out[18] = float(grip) if np.isscalar(grip) else float(grip.ravel()[0]) return out def _pack_bimanual(right_xyz, right_rot6d, left_xyz, left_rot6d, right_grip, left_grip): ''' Pack bimanual action into the unified 20D format. ''' out = np.zeros(20, dtype=np.float32) out[0:3] = right_xyz out[3:9] = right_rot6d out[9:12] = left_xyz out[12:18] = left_rot6d out[18] = float(right_grip) if np.isscalar(right_grip) else float(np.asarray(right_grip).ravel()[0]) out[19] = float(left_grip) if np.isscalar(left_grip) else float(np.asarray(left_grip).ravel()[0]) return out # --------------------------------------------------------------------------- # Per-dataset parsers # --------------------------------------------------------------------------- # Each returns {'action': np.ndarray (20,), 'desired': np.ndarray (20,) or None, 'actual': np.ndarray (20,) or None} # for a single (time, cam) entry from sample['action'][t]. _LBM_POSE_KEYS = [ 'poses__right::panda__xyz', # (3,) 'poses__right::panda__rot_6d', # (6,) 'poses__left::panda__xyz', # (3,) 'poses__left::panda__rot_6d', # (6,) 'grippers__right::panda_hand', # (1,) 'grippers__left::panda_hand', # (1,) ] def _parse_lbm(entry): ''' LBM (bimanual panda): flat dict with 'action' (20,) plus robot__actual__* and robot__desired__* keys at top level. All 3 keys are in the same 20-D EE-pose space: xyz(3)+rot6d(6)+grip(1) per arm. action[t] == desired[t+1] always holds. ''' # NOTE(dc): for now we take only the action for the current timestep action = _current_action(entry['action']) result = {'action': action, 'desired': None, 'actual': None} # Desired state: robot__desired__poses + robot__desired__grippers -> (20,) desired_keys = [f'robot__desired__{k}' for k in _LBM_POSE_KEYS] if all(k in entry for k in desired_keys): parts = [_current_action(entry[k]).ravel() for k in desired_keys] result['desired'] = np.concatenate(parts) # (20,) # Actual state: robot__actual__poses + robot__actual__grippers -> (20,) actual_keys = [f'robot__actual__{k}' for k in _LBM_POSE_KEYS] if all(k in entry for k in actual_keys): parts = [_current_action(entry[k]).ravel() for k in actual_keys] result['actual'] = np.concatenate(parts) # (20,) return result def _droid_7d_to_20d(cart_pos, grip_pos): ''' Convert DROID 7D (xyz(3) + euler_xyz(3) + grip(1)) to unified 20D. ''' xyz = cart_pos[:3] euler = cart_pos[3:6] rot6d = _euler_xyz_to_rot6d(euler) grip = grip_pos[0] if len(grip_pos) > 0 else 0.0 return _pack_single_arm(xyz, rot6d, grip) def _parse_droid(entry): ''' DROID: dict with cartesian_position (6,) = xyz(3) + XYZ euler(3), gripper_position (scalar). Converts to unified 20D: euler -> rot6d. desired = not directly available, recommend to use action[t-1] ''' cart_pos = _to_np(entry['cartesian_position']) # (6,) grip_pos = _to_np(entry['gripper_position']) # (1,) action = _droid_7d_to_20d(cart_pos, grip_pos) # Actual: observed state from robot_state, same conversion actual = None robot_state = entry.get('robot_state', {}) if isinstance(robot_state, dict) and 'cartesian_position' in robot_state: actual_cart = _to_np(robot_state['cartesian_position']) # (6,) actual_grip = _to_np(robot_state.get('gripper_position', 0.0)) # (1,) actual = _droid_7d_to_20d(actual_cart, actual_grip) return {'action': action, 'desired': None, 'actual': actual} def _rh20t_8d_to_20d(tcp, grip): ''' Convert RH20T 8D (xyz(3) + quat_xyzw(4) + grip(1)) to unified 20D. ''' xyz = tcp[:3] quat = tcp[3:7] rot6d = _quat_xyzw_to_rot6d(quat) grip_val = grip[0] if len(grip) > 0 else 0.0 return _pack_single_arm(xyz, rot6d, grip_val) def _parse_rh20t(entry): ''' RH20T: action_tcp(7) = xyz(3) + quaternion_xyzw(4), action_gripper(1). Converts to unified 20D: quaternion -> rot6d. desired = not directly available, use action[t-1] via time-shift. ''' # RH20T_d12 webbing wraps the flat fields in an extra {'action': {...}} layer. if 'action' in entry and isinstance(entry['action'], dict) and 'action_tcp' not in entry: entry = entry['action'] if 'action_tcp' in entry: tcp = _to_np(entry['action_tcp']) # (7,) grip = _to_np(entry.get('action_gripper', 0.0)) # (1,) action = _rh20t_8d_to_20d(tcp, grip) else: print(f'[yellow] RH20T: no action_tcp in entry, keys={list(entry.keys())}[/yellow]') action = np.zeros(20, dtype=np.float32) # Actual: state_tcp(7) + state_gripper(1), same conversion actual = None if 'state_tcp' in entry: state_tcp = _to_np(entry['state_tcp']) # (7,) state_grip = _to_np(entry.get('state_gripper', 0.0)) # (1,) actual = _rh20t_8d_to_20d(state_tcp, state_grip) return {'action': action, 'desired': None, 'actual': actual} def _parse_humanoid(entry): ''' HumanoidEveryday (G1): per-frame dict with 'actions' and 'states' sub-dicts. What this dataset provides for the unified output keys -- HAVE vs MISSING: action[t] (commanded EE pose): HAVE. actions['right_pose'/'left_pose'] are 4x4 homogeneous EE transforms (IK-solved), used directly (xyz + rot6d). NO forward kinematics needed -- this is the trajectory we condition on / visualize. desired[t]: MISSING. HumEv has no separate desired stream (use action[t-1] downstream if ever needed). Returned as None. actual[t] (measured EE pose): MISSING / not implemented. states has only joint angles (arm_state/hand_state/leg_state), NO right_pose/left_pose, so a measured EE pose would require forward kinematics from arm_state -- which we do NOT do here. Returned as None. (The states-pose branch below is inert and only fires if a future webbing variant ever exposes state EE poses.) gripper [18/19]: IMPLEMENTED but crude PROXY. G1 has a 14-dim dexterous hand (no single gripper width), so we map mean states['hand_state'] per hand to [0=closed .. 0.1=open]. Raw values are signed multi-DOF joint angles (not pure flexion), so it is approximate/uncalibrated. NOTE: this feeds BOTH the marker color AND the model's action conditioning -- set r_grip=l_grip=0.0 to disable. Only activates for 14-dim hands (G1); other robots (e.g. 12-dim) -> 0.0. If actions lacks pose keys entirely -> all-zero action. Do NOT fall back to joint angles: they have a different meaning and would pollute the shared 20D vector. ''' actions_dict = entry.get('actions', {}) states_dict = entry.get('states', {}) # Gripper openness PROXY (uncalibrated) in [0(closed) .. 0.1(open)]; assumes # hand_state layout [right 7, left 7]. Only 14-dim hands (G1); else stays 0.0. r_grip = l_grip = 0.0 hs = _to_np(states_dict['hand_state']) if 'hand_state' in states_dict else np.zeros(0, np.float32) if hs.shape[0] >= 14: FLEX_FULL = 1.6 # approx G1 finger flexion (rad) near full close; uncalibrated r_grip = 0.1 * (1.0 - float(np.clip(np.mean(hs[0:7]) / FLEX_FULL, 0.0, 1.0))) l_grip = 0.1 * (1.0 - float(np.clip(np.mean(hs[7:14]) / FLEX_FULL, 0.0, 1.0))) # HAVE: commanded EE pose (right_pose/left_pose 4x4) -> unified 20D, no FK needed. if 'right_pose' in actions_dict and 'left_pose' in actions_dict: right_T = _to_np_2d(actions_dict['right_pose']) # (4, 4) left_T = _to_np_2d(actions_dict['left_pose']) # (4, 4) r_xyz, r_rot6d = _homogeneous_to_xyz_rot6d(right_T) l_xyz, l_rot6d = _homogeneous_to_xyz_rot6d(left_T) action = _pack_bimanual(r_xyz, r_rot6d, l_xyz, l_rot6d, r_grip, l_grip) else: # No EE-pose keys: zeros. Do NOT fall back to joint angles (different meaning). print('[yellow] HumanoidEveryday: no right_pose/left_pose in actions; ' 'returning zero action[/yellow]') action = np.zeros(20, dtype=np.float32) # MISSING: measured EE pose. states has only joint angles (no right_pose/left_pose), # so 'actual' would need forward kinematics from arm_state (not done) -> stays None. # This branch is inert; it only fires if a future variant exposes state EE poses. actual = None if 'right_pose' in states_dict and 'left_pose' in states_dict: right_T = _to_np_2d(states_dict['right_pose']) left_T = _to_np_2d(states_dict['left_pose']) r_xyz, r_rot6d = _homogeneous_to_xyz_rot6d(right_T) l_xyz, l_rot6d = _homogeneous_to_xyz_rot6d(left_T) actual = _pack_bimanual(r_xyz, r_rot6d, l_xyz, l_rot6d, r_grip, l_grip) return {'action': action, 'desired': None, 'actual': actual} def _parse_agibot(entry): ''' AgiBotWorld: bimanual humanoid with nested action/state dicts. Structure: action/state -> end -> position (2, 3), orientation (2, 4) [xyzw quat] -> effector -> position (2,) [right_grip, left_grip] -> joint -> position (14,) -> head/waist/robot -> misc We use end-effector pose (xyz + quat -> rot6d) for both arms, plus effector position for grippers. Row 0 = right arm, row 1 = left arm. ''' act_dict = entry.get('action', entry) end = act_dict.get('end', {}) effector = act_dict.get('effector', {}) pos = _to_np_2d(end.get('position', np.zeros((2, 3)))) # (2, 3) ori = _to_np_2d(end.get('orientation', np.zeros((2, 4)))) # (2, 4) xyzw grip = _to_np(effector.get('position', np.zeros(2))) # (2,) r_xyz = pos[0] # (3,) l_xyz = pos[1] # (3,) r_rot6d = _quat_xyzw_to_rot6d(ori[0]) # (6,) l_rot6d = _quat_xyzw_to_rot6d(ori[1]) # (6,) r_grip = grip[0] if len(grip) > 0 else 0.0 l_grip = grip[1] if len(grip) > 1 else 0.0 action = _pack_bimanual(r_xyz, r_rot6d, l_xyz, l_rot6d, r_grip, l_grip) # Actual state: same structure under 'state' actual = None state_dict = entry.get('state', {}) if 'end' in state_dict: s_end = state_dict['end'] s_eff = state_dict.get('effector', {}) s_pos = _to_np_2d(s_end.get('position', np.zeros((2, 3)))) s_ori = _to_np_2d(s_end.get('orientation', np.zeros((2, 4)))) s_grip = _to_np(s_eff.get('position', np.zeros(2))) actual = _pack_bimanual( s_pos[0], _quat_xyzw_to_rot6d(s_ori[0]), s_pos[1], _quat_xyzw_to_rot6d(s_ori[1]), s_grip[0] if len(s_grip) > 0 else 0.0, s_grip[1] if len(s_grip) > 1 else 0.0, ) return {'action': action, 'desired': None, 'actual': actual} def _parse_yam(entry): ''' YAMxdof (XDOF teleop, bimanual YAM arms): joint_raw format, no EE poses (would need FK). actions: left/right_arm_leader (6,) joint angles + left/right_eef_leader (1,) gripper. states: left/right_arm_proprio (18,) = joints(6)+more, left/right_eef_proprio (3,) = grip+more. JOINT-SPACE EXCEPTION: dims [0:6]/[9:15] hold joint angles instead of xyz+rot6d (mirrors the per-arm block layout; rot6d tails stay zero). OK for YAM-only training. TODO(bvh): before mixing YAM with EE-pose datasets (X-embodiment AnyAct runs), convert to EE pose via FK: raw data has NO pose anywhere (verified at mcap level), but the official 6-DOF URDF exists (github.com/i2rt-robotics/i2rt yam.urdf) -> pytorch-kinematics over these joint angles gives flange pose per arm (in PER-ARM base frame; rig calib unavailable). ''' def _pack_joints(d, right_key, left_key, right_grip, left_grip): out = np.zeros(20, dtype=np.float32) out[0:6] = _to_np(d[right_key])[:6] out[9:15] = _to_np(d[left_key])[:6] out[18] = float(_to_np(right_grip)[0]) if right_grip is not None else 0.0 out[19] = float(_to_np(left_grip)[0]) if left_grip is not None else 0.0 return out actions_dict = entry.get('actions', {}) states_dict = entry.get('states', {}) if 'right_arm_leader' in actions_dict and 'left_arm_leader' in actions_dict: action = _pack_joints(actions_dict, 'right_arm_leader', 'left_arm_leader', actions_dict.get('right_eef_leader'), actions_dict.get('left_eef_leader')) else: print(f'[yellow] YAMxdof: no arm_leader keys in actions, keys={list(actions_dict.keys())}; ' f'zero action[/yellow]') action = np.zeros(20, dtype=np.float32) actual = None if 'right_arm_proprio' in states_dict and 'left_arm_proprio' in states_dict: actual = _pack_joints(states_dict, 'right_arm_proprio', 'left_arm_proprio', states_dict.get('right_eef_proprio'), states_dict.get('left_eef_proprio')) return {'action': action, 'desired': None, 'actual': actual} def _parse_generic(entry): ''' Fallback parser for unknown datasets. If an 'action' key exists, use it directly. Otherwise, concatenate all numeric values in the dict. ''' if 'action' in entry: action = entry['action'] if isinstance(action, (np.ndarray, torch.Tensor)): return {'action': _current_action(action), 'desired': None, 'actual': None} return {'action': _to_np(action), 'desired': None, 'actual': None} # Try to concatenate all numeric arrays parts = [] for k, v in sorted(entry.items()): try: parts.append(_to_np(v)) except (TypeError, ValueError): pass if parts: return {'action': np.concatenate(parts), 'desired': None, 'actual': None} print(f'[yellow] Unknown dataset: empty action dict, keys={list(entry.keys())}[/yellow]') return {'action': np.zeros(1, dtype=np.float32), 'desired': None, 'actual': None} # --------------------------------------------------------------------------- # Dataset name -> parser dispatch # --------------------------------------------------------------------------- _DATASET_PARSERS = { 'lbm': _parse_lbm, 'droid': _parse_droid, 'rh20t': _parse_rh20t, 'humanoideveryday': _parse_humanoid, 'agibot': _parse_agibot, 'yam': _parse_yam, } def _get_parser(dset_name, parsers=None): ''' Find the right parser for a dataset name (case-insensitive substring match). `parsers` defaults to _DATASET_PARSERS; pass a custom table to override (e.g. action_format_v2). ''' name = dset_name.lower() for key, parser in (parsers or _DATASET_PARSERS).items(): if key in name: return parser print(f'[yellow] action_utils: no parser for dataset "{dset_name}", using generic fallback[/yellow]') return _parse_generic # --------------------------------------------------------------------------- # Main entry point # --------------------------------------------------------------------------- ACTION_DIM = 20 # Unified action dimension: xyz(3)+rot6d(6)+grip(1) per arm. def _pad_to_action_dim(arr, target_dim=ACTION_DIM): ''' Coerce any array to exactly target_dim float32: flatten, truncate if longer, zero-pad if shorter. Returns None only for None input (desired/actual may be None). ''' if arr is None: return None arr = np.asarray(arr, dtype=np.float32).reshape(-1) padded = np.zeros(target_dim, dtype=np.float32) n = min(arr.shape[0], target_dim) padded[:n] = arr[:n] return padded def parse_anydata_raw_robot_action(sample, dset_name='', parsers=None): ''' Parse raw action dicts from an anydata sample into a normalized format. Preserves the key shape of `sample['action']`: if input keys are scalar time (after Base.py post_process_sample strips action keys), output keys are also scalar time. If input keys are (time, cam) tuples, output keys match. :param sample: anydata sample dict with an 'action' key. :param dset_name: dataset identifier string (e.g. 'lbmv12', 'DROID15'). :param parsers: optional {substr: parser} table overriding _DATASET_PARSERS (used by action_format_v2 to swap in its own _parse_droid). None = default. :return actions: dict with the same keys as sample['action'], mapping to {'action': array, 'desired': array|None, 'actual': array|None}. ''' parser = _get_parser(dset_name, parsers) actions = {} for key, entry in sample['action'].items(): if not entry: continue # CONTRACT: this wrapper ALWAYS returns a fixed ACTION_DIM structure per key, # regardless of per-dataset parser bugs or missing/broken/corrupt/variable-length # raw data (e.g. HumEv right_angles vary 7..12). A parser that raises or returns # a malformed dict must NOT propagate variable shapes into collate downstream. try: parsed = parser(entry) if not isinstance(parsed, dict): raise TypeError(f'parser returned {type(parsed).__name__}, expected dict') except Exception as e: pname = getattr(parser, '__name__', str(parser)) print(f'[yellow] parse_anydata_raw_robot_action: {pname} failed for "{dset_name}" ' f'key {key} ({type(e).__name__}: {e}); zero action[/yellow]') parsed = {'action': None, 'desired': None, 'actual': None} # Pad all keys to ACTION_DIM so the rest of the pipeline sees uniform shape. # 'action' is NEVER None (zeros fallback); 'desired'/'actual' are ACTION_DIM or None. out = {} for sub in ('action', 'desired', 'actual'): arr = _pad_to_action_dim(parsed.get(sub, None)) if arr is not None: # Defensive: raw fields can carry NaN/inf (e.g. unavailable telemetry); zero them # so they never reach the model. nan_to_num copies, so raw entry is untouched. arr = np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0) out[sub] = arr if out['action'] is None: out['action'] = np.zeros(ACTION_DIM, dtype=np.float32) actions[key] = out return actions