diff --git a/CMakeLists.txt b/CMakeLists.txt index 01e34be..c6354ca 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -156,8 +156,13 @@ if(GGML_CUDA) src/kernels/bitvla/bitvla_vit_cuda.cu src/kernels/bitvla/bitvla_fp32head_cuda.cu ) + # Device-link each archive on its own. Without this, vla_core nvlinks the + # relocatable device code of both CUDA libraries together, and CUDA 12.x + # rejects the mix because only bitvla compiles with --use_fast_math: + # "nvlink fatal: Cicc option values for '-ftz' do not match". set_target_properties(bitvla_cuda_kernels PROPERTIES CUDA_SEPARABLE_COMPILATION ON + CUDA_RESOLVE_DEVICE_SYMBOLS ON POSITION_INDEPENDENT_CODE ON ) target_include_directories(bitvla_cuda_kernels PRIVATE @@ -180,6 +185,7 @@ if(GGML_CUDA) ) set_target_properties(vla_cuda_ops PROPERTIES CUDA_SEPARABLE_COMPILATION ON + CUDA_RESOLVE_DEVICE_SYMBOLS ON POSITION_INDEPENDENT_CODE ON ) target_compile_features(vla_cuda_ops PRIVATE cxx_std_17) diff --git a/eval/client/run_ALOHA_client_direct.py b/eval/client/run_ALOHA_client_direct.py index c6c64a2..cb416f7 100644 --- a/eval/client/run_ALOHA_client_direct.py +++ b/eval/client/run_ALOHA_client_direct.py @@ -287,8 +287,12 @@ def _s(v): return np.array([[[float(v)]]], dtype=np.float32) return { "video.image": _hwc_u8(front_rgb, image_size)[None, None], "video.wrist_image": _hwc_u8(wrist_rgb, image_size)[None, None], + # Both namings of the same numbers: EEF-style checkpoints read the + # six scalars, joint-space ones read the grouped "single_arm" vector. + # The client picks whichever its statistics file declares. "state.x": _s(j[0]), "state.y": _s(j[1]), "state.z": _s(j[2]), "state.roll": _s(j[3]), "state.pitch": _s(j[4]), "state.yaw": _s(j[5]), + "state.single_arm": j[:ARM_DOF].reshape(1, 1, -1), "state.gripper": gripper.reshape(1, 1, -1), "task": task, } @@ -322,19 +326,31 @@ def __init__(self, args: argparse.Namespace): ) self.bridge = CvBridge() + # Which follower a single-arm policy drives. Dual-arm always uses both, + # left first, so --arm-side only matters for the single-arm case. + side = args.arm_side + front_topic = args.front_topic + wrist_topic = args.wrist_topic or f"/camera_wrist_{side}/camera/color/image_rect_raw" + state_topic = args.state_topic or f"/follower_{side}/joint_states" + # Publishers self.left_arm_pub = self.create_publisher(JointGroupCommand, "/follower_left/commands/joint_group", 10) self.left_hand_pub = self.create_publisher(JointSingleCommand, "/follower_left/commands/joint_single", 10) - if self.dual_arm: + if self.dual_arm or side == "right": self.right_arm_pub = self.create_publisher(JointGroupCommand, "/follower_right/commands/joint_group", 10) self.right_hand_pub = self.create_publisher(JointSingleCommand, "/follower_right/commands/joint_single", 10) + # Single-arm chunks go out on the chosen follower; the observation + # buffer stays the "left" one either way so build_obs is unchanged. + self._pub_single = self._pub_right if (side == "right" and not self.dual_arm) else self._pub_left # Subscribers - self.create_subscription(Image, "/camera_high/camera/color/image_raw", self._cb_front, qos) - self.create_subscription(Image, "/camera_wrist_left/camera/color/image_raw", self._cb_wrist_left, qos) - self.create_subscription(JointState, "/follower_left/joint_states", self._cb_left_state, 10) + self.create_subscription(Image, front_topic, self._cb_front, qos) + self.create_subscription(Image, wrist_topic, self._cb_wrist_left, qos) + self.create_subscription(JointState, state_topic, self._cb_left_state, 10) if self.dual_arm: self.create_subscription(JointState, "/follower_right/joint_states", self._cb_right_state, 10) + self.log.info(f"topics front={front_topic} wrist={wrist_topic} " + f"state={state_topic} commands=/follower_{side if not self.dual_arm else 'left+right'}") # Sensor buffers self.lock = Lock() @@ -386,6 +402,7 @@ def __init__(self, args: argparse.Namespace): n_action_steps = args.n_action_steps, stats_json = args.stats_json, bitvla_unnorm_key = args.bitvla_unnorm_key, + rel_stats_json = args.rel_stats_json, ) self.client.reset() self.log.info(f"client ready arch={self.arch} addr={args.vla_addr}") @@ -578,11 +595,7 @@ def _run_inference_sync(self): prev_right = self.prev_right_state if self.dual_arm else None # noqa: F841 snapshot for downstream consumers if front is None or wrist is None or left_state is None: - self.log.info( - f"Waiting for data front={front is not None}" - f" wrist={wrist is not None}" - f" left_state={left_state is not None}" - ) + self._log_waiting(front is not None, wrist is not None, left_state is not None) self._end_action = True return @@ -624,11 +637,9 @@ def _run_inference_async(self): # First call or after a timeout fallback: trigger and wait. with self.lock: if self.front_rgb is None or self.wrist_left_rgb is None or self.left_state is None: - self.log.info( - f"Waiting for data front={self.front_rgb is not None}" - f" wrist={self.wrist_left_rgb is not None}" - f" left_state={self.left_state is not None}" - ) + self._log_waiting(self.front_rgb is not None, + self.wrist_left_rgb is not None, + self.left_state is not None) self._end_action = True return if self.dual_arm and self.right_state is None: @@ -692,7 +703,7 @@ def _execute_chunk(self, chunk: np.ndarray): for i in range(n_steps): t_step = time.time() row = chunk[i] - self._pub_left(row[:ARM_DOF], row[ARM_DOF]) + self._pub_single(row[:ARM_DOF], row[ARM_DOF]) if self.dual_arm and row.size >= JOINT_DOF * 2: self._pub_right(row[JOINT_DOF:JOINT_DOF + ARM_DOF], row[JOINT_DOF + ARM_DOF]) time.sleep(max(0.0, 1.0 / 200.0 - (time.time() - t_step))) @@ -726,6 +737,19 @@ def _get_chunk(self, front, wrist, left_state, right_state) -> np.ndarray: # Publishers # ------------------------------------------------------------------ + def _log_waiting(self, front_ok: bool, wrist_ok: bool, state_ok: bool) -> None: + """One line per second while a sensor stream is still missing. + + The control loop retries at its full rate, so an unthrottled log here + buries every other message under thousands of identical lines. + """ + now = time.time() + if now - getattr(self, "_last_wait_log", 0.0) < 1.0: + return + self._last_wait_log = now + self.log.info( + f"Waiting for data front={front_ok} wrist={wrist_ok} state={state_ok}") + def _pub_left(self, arm: np.ndarray, gripper: float): arm_msg = JointGroupCommand() arm_msg.name = "arm" @@ -794,7 +818,20 @@ def main(): parser.add_argument("--smooth-step", type=int, default=20, help="Interpolation sub-steps between actions (1 = no smoothing).") parser.add_argument("--dual-arm", action="store_true") + parser.add_argument("--arm-side", choices=("left", "right"), default="left", + help="follower a single-arm policy reads and drives (default: left)") + parser.add_argument("--front-topic", type=str, + default="/camera_high/camera/color/image_rect_raw", + help="overhead camera topic") + parser.add_argument("--wrist-topic", type=str, default=None, + help="wrist camera topic (default: /camera_wrist_/camera/color/image_rect_raw)") + parser.add_argument("--state-topic", type=str, default=None, + help="joint state topic (default: /follower_/joint_states)") parser.add_argument("--stats-json", type=str, default=None) + parser.add_argument("--rel-stats-json", type=str, default=None, + help="GR00T meta/relative_stats.json. Required for checkpoints whose " + "action modalities are RELATIVE: the chunk is then an offset from " + "the observed state and is added back onto it.") parser.add_argument("--bitvla-unnorm-key", type=str, default=None) parser.add_argument("--embodiment", type=str, default="new_embodiment", help="Embodiment key in statistics.json (default: new_embodiment).") diff --git a/eval/client/vla_cpp_client.py b/eval/client/vla_cpp_client.py index 70f18da..6ffec82 100644 --- a/eval/client/vla_cpp_client.py +++ b/eval/client/vla_cpp_client.py @@ -142,6 +142,7 @@ def __init__( stats_json: str | Path | None = None, bitvla_unnorm_key: str | None = None, + rel_stats_json: str | Path | None = None, ): if arch not in ARCH_PRESETS: raise ValueError(f"unknown arch {arch!r}; expected one of {sorted(ARCH_PRESETS)}") @@ -284,6 +285,9 @@ def _oft_norm(x, q01=q01, q99=q99, mask=mask): self._gr00t_action_unnorm = None self._gr00t_state_norm = None + self._gr00t_state_keys = None + self._gr00t_state_dims = None + self._gr00t_last_state_raw = None if arch == "gr00t_n1_7" and stats_json is not None: stats_path = Path(stats_json) if not stats_path.exists(): @@ -303,38 +307,92 @@ def _oft_norm(x, q01=q01, q99=q99, mask=mask): f"--bitvla-unnorm-key explicitly for arch=gr00t_n1_7.") if key not in blob: raise KeyError(f"embodiment {key!r} not in {stats_path}; have {list(blob)}") - modalities = ["x", "y", "z", "roll", "pitch", "yaw", "gripper"] action_stats = blob[key]["action"] - q01 = np.array([action_stats[m]["q01"][0] for m in modalities], dtype=np.float32) - q99 = np.array([action_stats[m]["q99"][0] for m in modalities], dtype=np.float32) - rng = (q99 - q01).astype(np.float32) - def _unnorm(chunk_132: np.ndarray, q01=q01, q99=q99, rng=rng) -> np.ndarray: + modalities, mod_dims = self._gr00t_modality_layout(action_stats) + q01 = self._gr00t_quantile(action_stats, modalities, "q01") + q99 = self._gr00t_quantile(action_stats, modalities, "q99") + act_dim = int(q01.size) + + # Checkpoints trained with use_relative_action predict, for the + # modalities listed in meta/relative_stats.json, the offset from the + # observed state rather than an absolute target - and those offsets + # carry their own per-chunk-step statistics. GR00T builds them in + # data/stats.py::load_relative_actions and undoes them with + # JointActionChunk.to_absolute_chunking, i.e. reference + offset, + # every step of the chunk sharing the one reference state. + rel_stats = {} + if rel_stats_json is not None: + rel_path = Path(rel_stats_json) + if not rel_path.exists(): + raise FileNotFoundError(f"relative stats JSON not found at {rel_path}") + rel_stats = json.loads(rel_path.read_text()) + rel_names = [m for m in modalities if m in rel_stats] + + if rel_names: + horizon = min(len(rel_stats[m]["min"]) for m in rel_names) + q01_t = np.tile(q01, (horizon, 1)).astype(np.float32) + q99_t = np.tile(q99, (horizon, 1)).astype(np.float32) + is_rel = np.zeros(act_dim, dtype=bool) + off = 0 + for m, dim in zip(modalities, mod_dims): + if m in rel_stats: + r = rel_stats[m] + # min/max, never q01/q99, even when the checkpoint sets + # use_percentiles. GR00T swaps the whole norm_params entry + # for a relative key with the raw relative_stats dict + # (state_action_processor.py), which bypasses the branch + # that would otherwise substitute the percentiles. Using + # q01/q99 here scales every offset down by ~2x and the arm + # never reaches the pose the policy is steering to. + a = np.asarray(r["min"], dtype=np.float32)[:horizon] + b = np.asarray(r["max"], dtype=np.float32)[:horizon] + if a.shape[1] != dim: + raise ValueError( + f"relative stats for {m!r} are {a.shape[1]}-wide, " + f"statistics say {dim}") + q01_t[:, off:off + dim] = a + q99_t[:, off:off + dim] = b + is_rel[off:off + dim] = True + off += dim + rng_t = (q99_t - q01_t).astype(np.float32) + + def _unnorm(chunk_132, q01_t=q01_t, rng_t=rng_t, is_rel=is_rel, + act_dim=act_dim, horizon=horizon): + n = min(len(chunk_132), horizon) + norm = np.clip(chunk_132[:n, :act_dim].astype(np.float32), -1.0, 1.0) + raw = (norm + 1.0) * 0.5 * rng_t[:n] + q01_t[:n] + ref = self._gr00t_last_state_raw + if ref is None: + raise RuntimeError("relative actions need the observation state; " + "none was recorded for this request") + raw[:, is_rel] += np.asarray(ref, dtype=np.float32)[:act_dim][is_rel] + return raw.astype(np.float32) + else: + rng = (q99 - q01).astype(np.float32) - norm = chunk_132[..., :7].astype(np.float32) + def _unnorm(chunk_132: np.ndarray, q01=q01, q99=q99, rng=rng, + act_dim=act_dim) -> np.ndarray: - norm = np.clip(norm, -1.0, 1.0) + norm = chunk_132[..., :act_dim].astype(np.float32) + + norm = np.clip(norm, -1.0, 1.0) - raw = (norm + 1.0) * 0.5 * rng[None, :] + q01[None, :] + raw = (norm + 1.0) * 0.5 * rng[None, :] + q01[None, :] - return raw.astype(np.float32) + return raw.astype(np.float32) self._gr00t_action_unnorm = _unnorm print(f"vla-cpp-direct[arch=gr00t_n1_7]: action unnormalizer " - f"(q01/q99 + clip + gripper flip) via {stats_path}::{key}.action " - f"[modalities={modalities}, q01={q01.tolist()}, q99={q99.tolist()}]", + f"(q01/q99 + clip) via {stats_path}::{key}.action " + f"[modalities={list(zip(modalities, mod_dims))} -> {act_dim}-D, " + f"relative={rel_names or 'none'}]", flush=True) state_stats = blob[key]["state"] - s_q01_parts, s_q99_parts = [], [] - for m, dim in zip(self._GR00T_STATE_KEYS, self._GR00T_STATE_DIMS): - if m not in state_stats: - raise KeyError(f"state modality {m!r} not in {stats_path}::{key}.state") - q01_m = np.asarray(state_stats[m]["q01"], dtype=np.float32) - q99_m = np.asarray(state_stats[m]["q99"], dtype=np.float32) - if q01_m.size != dim or q99_m.size != dim: - raise ValueError(f"state.{m}: stats dim {q01_m.size}/{q99_m.size} != expected {dim}") - s_q01_parts.append(q01_m); s_q99_parts.append(q99_m) - s_q01 = np.concatenate(s_q01_parts) - s_q99 = np.concatenate(s_q99_parts) + state_keys, state_dims = self._gr00t_modality_layout(state_stats) + self._gr00t_state_keys = tuple(state_keys) + self._gr00t_state_dims = tuple(state_dims) + s_q01 = self._gr00t_quantile(state_stats, state_keys, "q01") + s_q99 = self._gr00t_quantile(state_stats, state_keys, "q99") s_rng = (s_q99 - s_q01).astype(np.float32) def _state_norm(state_8d: np.ndarray, q01=s_q01, q99=s_q99, rng=s_rng) -> np.ndarray: @@ -343,7 +401,8 @@ def _state_norm(state_8d: np.ndarray, q01=s_q01, q99=s_q99, rng=s_rng) -> np.nda self._gr00t_state_norm = _state_norm print(f"vla-cpp-direct[arch=gr00t_n1_7]: state normalizer " f"(q01/q99 + clip) via {stats_path}::{key}.state " - f"[q01={s_q01.tolist()}, q99={s_q99.tolist()}]", flush=True) + f"[modalities={list(zip(state_keys, state_dims))} -> {s_q01.size}-D, " + f"q01={s_q01.tolist()}, q99={s_q99.tolist()}]", flush=True) if arch == "gr00t_n1_6" and stats_json is not None: stats_path = Path(stats_json) @@ -1084,9 +1143,42 @@ def _gr00t_eval_image_transform(img_u8_hwc: np.ndarray, target_size: int, img = cv2.resize(img, (target_size, target_size), interpolation=cv2.INTER_AREA) return np.ascontiguousarray(img, dtype=np.uint8) + # Fallback layout: the EEF-style embodiments (x/y/z/rpy + 2-finger gripper) + # the GR00T paths were first written against. Checkpoints trained on a joint + # space name their modalities differently ("single_arm", "gripper", ...), so + # the real layout is read out of the checkpoint's statistics JSON instead - + # see _gr00t_modality_layout. _GR00T_STATE_KEYS = ("x", "y", "z", "roll", "pitch", "yaw", "gripper") _GR00T_STATE_DIMS = (1, 1, 1, 1, 1, 1, 2) + @classmethod + def _gr00t_modality_layout(cls, group_stats: dict) -> tuple[list[str], list[int]]: + """Ordered modality names and their widths for one statistics group. + + Keeps the historical EEF order when the checkpoint has every one of + those modalities; otherwise follows the order the checkpoint itself + declares, which is the order GR00T's processor concatenates them in. + """ + if all(k in group_stats for k in cls._GR00T_STATE_KEYS): + keys = list(cls._GR00T_STATE_KEYS) + else: + keys = list(group_stats.keys()) + dims = [int(np.asarray(group_stats[k]["q01"], dtype=np.float32).reshape(-1).size) + for k in keys] + return keys, dims + + @staticmethod + def _gr00t_quantile(group_stats: dict, keys: list[str], field: str) -> np.ndarray: + """Concatenate one quantile field across modalities, in the given order.""" + parts = [] + for k in keys: + if k not in group_stats: + raise KeyError(f"modality {k!r} not in statistics group; have {list(group_stats)}") + if field not in group_stats[k]: + raise KeyError(f"modality {k!r} lacks {field!r}; GR00T N1.7 needs q01/q99") + parts.append(np.asarray(group_stats[k][field], dtype=np.float32).reshape(-1)) + return np.concatenate(parts).astype(np.float32) + def _predict_chunk_gr00t_n1_7(self, observations: dict[str, Any]) -> np.ndarray: import re @@ -1112,10 +1204,14 @@ def _predict_chunk_gr00t_n1_7(self, observations: dict[str, Any]) -> np.ndarray: images_f32.append(img_f32) state_chunks = [] - for key, dim in zip(self._GR00T_STATE_KEYS, self._GR00T_STATE_DIMS): + state_keys = getattr(self, "_gr00t_state_keys", None) or self._GR00T_STATE_KEYS + state_dims = getattr(self, "_gr00t_state_dims", None) or self._GR00T_STATE_DIMS + for key, dim in zip(state_keys, state_dims): mk = f"state.{key}" if mk not in observations: - raise KeyError(f"gr00t_n1_7 state key '{mk}' missing; got {list(observations.keys())}") + raise KeyError( + f"gr00t_n1_7 state key '{mk}' missing; this checkpoint's statistics " + f"declare {list(state_keys)}; got {list(observations.keys())}") v = observations[mk] if isinstance(v, torch.Tensor): v = v.numpy() @@ -1125,6 +1221,9 @@ def _predict_chunk_gr00t_n1_7(self, observations: dict[str, Any]) -> np.ndarray: raise ValueError(f"gr00t_n1_7 state '{mk}': expected {dim}-d, got {v.size}-d") state_chunks.append(v) state_raw = np.concatenate(state_chunks, axis=0).astype(np.float32) + # Reference frame for relative actions: the raw, un-normalised state of + # this observation. Every step of the chunk is an offset from it. + self._gr00t_last_state_raw = state_raw.copy() if self._gr00t_state_norm is not None: state_raw = self._gr00t_state_norm(state_raw) state_padded = np.zeros(self.max_state_dim, dtype=np.float32) diff --git a/scripts/convert_gr00t_n1_7_to_gguf.py b/scripts/convert_gr00t_n1_7_to_gguf.py index c4c0755..7009f9f 100644 --- a/scripts/convert_gr00t_n1_7_to_gguf.py +++ b/scripts/convert_gr00t_n1_7_to_gguf.py @@ -43,6 +43,49 @@ resolve_out ) +def _uses_relative_actions(ckpt, processor_json: str) -> bool: + """Whether any action modality of a shipped embodiment is RELATIVE. + + The per-modality `rep` in processor_config.json is the signal that matters: + the global `use_relative_action` in experiment_cfg/ is a training switch that + is set on checkpoints whose modalities are all ABSOLUTE too, so reading it + alone mislabels them. Falls back to that flag only when the processor config + carries no action_configs at all. + """ + import re + try: + mods = json.loads(processor_json).get("processor_kwargs", {}).get("modality_configs", {}) + except (ValueError, TypeError): + mods = {} + # Only the embodiment this checkpoint was finetuned for counts. statistics.json + # carries every stock embodiment, plenty of which are relative, so keying off it + # marks every checkpoint relative. experiment_cfg/dataset_statistics.json holds + # just the finetuned one. + finetuned = ckpt / "experiment_cfg" / "dataset_statistics.json" + try: + shipped = set(json.loads(finetuned.read_text())) if finetuned.exists() else set() + except ValueError: + shipped = set() + saw_action_configs = False + for emb, cfg in mods.items(): + cfgs = (cfg.get("action") or {}).get("action_configs") or [] + if cfgs: + saw_action_configs = True + if emb not in shipped: + continue + if any(str(c.get("rep", "")).upper() == "RELATIVE" for c in cfgs): + return True + if saw_action_configs: + return False + for name in ("config.yaml", "conf.yaml"): + path = ckpt / "experiment_cfg" / name + if path.exists(): + m = re.search(r"^\s*use_relative_action:\s*(true|false)\s*$", + path.read_text(errors="replace"), re.IGNORECASE | re.MULTILINE) + if m: + return m.group(1).lower() == "true" + return False + ARCH = "gr00t_n1_7" KV = kv_prefix(ARCH) @@ -137,7 +180,7 @@ def main() -> int: CROP_FRACTION = float(cfg_json.get("crop_fraction", 0.95) or 0.95) ICS = cfg_json.get("image_crop_size", [230, 230]) or [230, 230] ITS = cfg_json.get("image_target_size", [256, 256]) or [256, 256] - USE_RELATIVE_ACTION = bool(cfg_json.get("use_relative_action", False)) + USE_RELATIVE_ACTION = False # resolved from the sidecars once they are read APPLY_SINCOS_STATE = bool(cfg_json.get("apply_sincos_state_encoding", False)) print(f"loading sharded safetensors from {ckpt} ...") @@ -173,6 +216,13 @@ def main() -> int: proc_kwargs = json.loads(processor_json).get("processor_kwargs", {}) if processor_json != "{}" else {} USE_PERCENTILES = bool(proc_kwargs.get("use_percentiles", True)) CLIP_OUTLIERS = bool(proc_kwargs.get("clip_outliers", True)) + USE_RELATIVE_ACTION = bool(cfg_json.get( + "use_relative_action", + _uses_relative_actions(ckpt, processor_json))) + if USE_RELATIVE_ACTION: + print(" NOTE: this checkpoint predicts RELATIVE actions. The engine returns them " + "as-is; the caller must add the observation state back (eval/client does this " + "given --rel-stats-json).") print(f"resolved cfg: vit=Qwen3-VL {VIT['vit_hidden']}d×{VIT['vit_layers']}L×{VIT['vit_heads']}h (Conv3d patch {VIT['patch_size']}², temporal {VIT['temporal_patch_size']}, " f"learned pos {VIT['vit_num_position_embeddings']}=48² + 2D rope; deepstack@{DEEPSTACK_IDXS}; merger LN={VIT['vit_hidden']} pre-merge / deepstack LN={c_merged} post-merge ⇒ " diff --git a/scripts/gen_gr00t_relative_stats.py b/scripts/gen_gr00t_relative_stats.py new file mode 100755 index 0000000..53e30c7 --- /dev/null +++ b/scripts/gen_gr00t_relative_stats.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generate GR00T relative-action statistics for the ALOHA right-arm datasets. + +Mirrors gr00t/data/stats.py::load_relative_actions for ActionType.NON_EEF: +for every frame i, the reference is state[i] and the chunk is action[i .. i+H-1]; +the relative action is a plain element-wise subtraction (JointPose.__sub__, see +JointActionChunk.relative_chunking). Stats are per chunk-step and per joint, +matching the [H, D] layout of meta/relative_stats.json. +""" +import json, sys, glob +from pathlib import Path +import numpy as np, pandas as pd + +H = 16 # len(modality action delta_indices) +SLICE = slice(0, 6) # modality.json: action/state single_arm = [0:6] +KEY = "single_arm" # only RELATIVE modalities get relative stats + +def deltas_for_dataset(root: Path) -> np.ndarray: + out = [] + for f in sorted(glob.glob(str(root / "data" / "**" / "*.parquet"), recursive=True)): + df = pd.read_parquet(f, columns=["action", "observation.state"]) + act = np.stack(df["action"].values).astype(np.float32)[:, SLICE] + st = np.stack(df["observation.state"].values).astype(np.float32)[:, SLICE] + usable = len(df) - (H - 1) + if usable <= 0: + continue + idx = np.arange(usable)[:, None] + np.arange(H)[None, :] # (usable, H) + out.append(act[idx] - st[:usable][:, None, :]) # (usable, H, 6) + return np.concatenate(out, axis=0) + +def stats_of(d: np.ndarray) -> dict: + return {"max": d.max(axis=0).tolist(), + "min": d.min(axis=0).tolist(), + "q01": np.quantile(d, 0.01, axis=0).tolist(), + "q99": np.quantile(d, 0.99, axis=0).tolist(), + "mean": d.mean(axis=0).tolist(), + "std": d.std(axis=0).tolist()} + +def main(roots, agg_out): + pooled = [] + for r in roots: + r = Path(r) + d = deltas_for_dataset(r) + pooled.append(d) + p = r / "meta" / "relative_stats.json" + p.write_text(json.dumps({KEY: stats_of(d)}, indent=4)) + print(f" {r.name}: {d.shape[0]} chunks -> {p}") + allp = np.concatenate(pooled, axis=0) + Path(agg_out).write_text(json.dumps({KEY: stats_of(allp)}, indent=4)) + print(f" aggregate: {allp.shape[0]} chunks, shape {allp.shape[1:]} -> {agg_out}") + print(f" delta range over all steps: min {allp.min():+.4f} max {allp.max():+.4f} mean |d| {np.abs(allp).mean():.4f}") + +if __name__ == "__main__": + main(sys.argv[1:-1], sys.argv[-1])