feat: add Elite/Rebot robot support, extend XTac-UMI training, and clean up JAX attention by Hubo1231 · Pull Request #24 · XenseRobotics-AI/xense-openpi

Hubo1231 · 2026-10-29T05:22:24Z

本 PR 新增 Elite CS66 和 Rebot B601 双臂机器人的训练与推理支持，扩展 XTac-UMI 数据处理，并清理 JAX attention 实现及训练配置。

主要变更

Elite CS66 / Rebot B601：新增数据配置、策略适配和推理示例，支持 RTC、解耦执行、硬件 recipes 与 dry-run。
XTac-UMI：支持 G1 数据的头部位姿裁剪、头部相机选择、SE(3) 相对位姿动作及推理时绝对动作还原；新增将数据集重定位至 Flexiv 起始 TCP 的工具，移除评估步数限制。
数据加载与归一化：兼容 legacy 和独立 DatasetReader 两类 LeRobot API，仅解码所需视频流；支持显式数据集路径和视频后端，恢复本地 lerobot-xense 环境下的 parquet 归一化统计入口。
JAX attention 与训练启动：移除 cuDNN fused attention、FP16 custom VJP 及相关配置，统一使用显式 attention；撤下已发现会导致 loss 和梯度错误的 XLA 通信参数，并更新排障与迁移文档。
训练配置：显式声明学习率调度，补充和更新多机器人任务及 RTC 示例，清理过时配置。

兼容性与迁移

现有 checkpoint 权重保持兼容；训练和推理配置中的 use_cudnn_attention、cudnn_attention_dtype 必须删除，旧字段会导致配置加载失败。
启动脚本及环境变量需清理文档中已撤下的 XLA 通信参数。
attention 清理本身不代表训练 NaN 问题已解决。

测试

新增或扩展了机器人数据配置、策略变换、RTC、LeRobot 视频筛选及 UMI 数据重定位的回归测试。

… loading

…orm stats Include Elite CS66 runtime, bench recipes, policy adapters, training config registration and tests, plus Elite egg-carton and Flexiv strap-tape training configs. Revert UMI-specific normalization changes and import training loaders only for Torch/RLDS paths so parquet statistics work with the local lerobot-xense implementation.

Handle lerobot-xense builds without dataset_reader and filter video timestamps on the dataset for their legacy getitem path. Preserve reader filtering for implementations that delegate reads to DatasetReader. Exercise both installed API variants, including actual frame reads with default, selected, and empty video whitelists.

…ters

Retire unsafe XLA collective launch flags and document the findings. Include the whiteboard inspection config, remove the obsolete box assembly config, and preserve current prompt formatting updates.

Vertax42

Code review（Claude Code 自动审读，未跑测试验证）。共 10 条，建议合并前至少处理 🔴 两条和 🟠 两条（真机安全）。

🔴 norm stats 在错误空间计算（`scripts/compute_norm_stats.py` 的 `_STATE_ACTION_TRANSFORM_TYPES`，该行不在 diff 中故放这里）

默认 parquet 统计路径只应用 DeltaActions / SubsampleActions，跳过了本 PR 新增的 XtacUmiStateActions（去头部位姿 29D→20D）和 XtacUmiRelativeActions（相对 TCP 动作 + 仅夹爪 2D state）。

后果：compute_norm_stats.py --config-name pi05_base_bi_flexiv_umi_g1_clean_plates_1011_h100（drop_head_pose + use_relative_pose_actions）写出的是 29D 绝对位姿统计。训练时 Normalize 按 mean[..., :x.shape[-1]] 截取，于是 2D 夹爪 state 用了左臂 TCP x/y 的统计，相对动作用了绝对位姿统计——不报错，训练静默使用尺度错误的输入和目标。XtacUmiRelativeActions 的 docstring 声称支持 norm-stats batch，但这条路径从未调用它。

其余问题见行内评论。

Vertax42 · 2026-11-02T09:56:34Z

        return FakeDataset(model_config, num_samples=1024)

-    dataset_meta = lerobot_dataset.LeRobotDatasetMetadata(repo_id)
+    dataset_meta = lerobot_dataset.LeRobotDatasetMetadata(repo_id, root=data_config.dataset_root)


🔴 数据集路径在统计与训练之间不一致

训练读新增的 DataConfig.dataset_root，但 compute_norm_stats 只认 --dataset-root / OPENPI_LEROBOT_DATASET_ROOT；而训练又不认这个环境变量——1011 的 yaml 恰恰让用户设它。

场景：按 yaml 设 OPENPI_LEROBOT_DATASET_ROOT=/data/ds → 统计来自 /data/ds，训练 LeRobotDataset(root=None) 去读 HF cache 或从 Hub 拉。反过来只在 config 设 dataset_root 则颠倒。两边可能用的是不同的数据。建议两处共用同一个解析函数（config 优先、环境变量兜底）。

Vertax42 · 2026-11-02T09:56:34Z

+            time.sleep(0.001)
+        while self.robot.rt_moving:
+            if time.time() - t0 > _RESET_TIMEOUT_S:
+                logger.warn("Reset motion timeout, proceeding anyway")


🟠 reset 失败仍继续开跑（真机安全）

reset 动作 1s 内未开始或 15s 内未完成时只打 warning 继续；另外 go_to_start=False 时也总会调 reset_to_initial_position。

场景：手臂故障/卡在半路，episode 照常开始，策略从任意（可能仍在运动的）姿态输出笛卡尔 servo 目标，首个增量动作基于模型没见过的状态。同 PR 里 Rebot env 在 reset 失败时会明确拒绝继续，建议对齐为 raise。

Vertax42 · 2026-11-02T09:56:34Z

+
+    def disconnect(self) -> None:
+        """Disconnect both arms, grippers, and cameras."""
+        if self.robot.is_connected:


🟠 部分连接时资源不释放

仅当 robot.is_connected 才调 robot.disconnect()。

场景：左臂连上、右臂失败，或一只臂运行中掉线 → 双臂 facade 的 is_connected 为 False，disconnect 被跳过，已连接那只臂的 servo 循环、夹爪串口、相机都不释放/不回零。本 PR 中 Rebot env 的注释正好提到这个情况并总是调用 facade 的 disconnect，建议同样处理。

Vertax42 · 2026-11-02T09:56:34Z

+    out.reset_index(drop=True).to_parquet(dst / "data/chunk-000/file-000.parquet", index=False)
+
+    episodes = pd.concat([pd.read_parquet(p) for p in sorted((src / "meta/episodes").rglob("*.parquet"))])
+    episodes = episodes[episodes["episode_index"].isin(kept)].sort_values("episode_index").copy()


🟡 per-episode stats 未随 rebase 更新

episodes 元数据从源复制，只改了索引列；stats/observation.state/*、stats/action/* 仍是源数据 29D、Pico 坐标系的统计，而 info.json 已声明 20D 特征。

场景：任何聚合 per-episode stats 的 LeRobot 路径（merge/aggregate、从 episodes 重算 meta/stats）都会拿到维度和坐标系都不对的统计 → shape 不匹配或归一化错误。建议删除这些列或按新数据重算。

Vertax42 · 2026-11-02T09:56:35Z

+    # Sessions are detected over all episodes: a dropped one still sits in the recording order.
+    starts = np.stack([states[rows[e][0]] for e in episode_ids])
+    if args.session_starts is not None:
+        session_of = np.searchsorted(np.asarray(args.session_starts), episode_ids, side="right") - 1


🟡 --session-starts 未校验

直接传给 np.searchsorted，不排序，也不检查首个起点 ≤ 最小 episode 编号。

场景：--session-starts 15 0（无序）或 episode 从 0 开始却传 --session-starts 5 → session id 为 -1 或错误，episode 被静默分到错误的 anchor/yaw 下。建议排序 + 断言 starts[0] <= min(episode_index)。

Vertax42 · 2026-11-02T09:56:35Z

+                delta_timestamps=self.delta_timestamps,
+                image_transforms=self.image_transforms,
+            )
+            selective_reader.hf_dataset = reader.hf_dataset


🟢 _ensure_reader 重建 SelectiveVideoDatasetReader 时只复制了 hf_dataset 和 _absolute_to_relative_idx。如果某个 LeRobot 版本的 DatasetReader 在 __init__ 后还持有其它状态（lazy-load 标志、视频路径缓存等），重建后的行为会不同（例如每个 worker 重新加载 parquet）。单测只覆盖了这两个属性。可以考虑改 __class__ 或复制完整 __dict__ 再覆盖需要的字段。

Vertax42 · 2026-11-02T09:56:35Z

@@ -0,0 +1,52 @@
+# Launch template (8xH100, hb-openpi env). Do not set XLA_FLAGS: the retired collective
+# combining/pipelining flags corrupted FSDP gradients; see docs/training-optimization.md.
+#   python scripts/compute_norm_stats.py --config-name pi05_base_bi_elite_egg_carton_to_tray_1012_h100


🟢 注释里的启动命令写的是 pi05_base_bi_elite_egg_carton_to_tray_1012_h100，但文件/config 名是 ..._1016_h100，照抄会报 unknown config。

Vertax42 · 2026-11-02T09:56:35Z

@@ -151,7 +166,6 @@ class Args:
    # Runtime settings
    runtime_hz: float = _TRAINING_HZ
    num_episodes: int = 1


🟢 删掉了 max_episode_steps（Args 和 Runtime 调用里都删了），其它机器人示例都还保留。传了 max_episode_steps: 200 的 run 文件（各 dry-run.yaml 的写法）会解析失败；没有它，真机运行只能靠 Ctrl+C 停，少了按步数的安全上限。如果是有意为之，建议在 PR 描述里说明。

Vertax42 · 2026-11-02T09:56:35Z

            "Set `tactile: true` under the data config's `base_config` to decode them."
        )
-    return decode_keys
+    required_video_keys = _repack_source_keys(data_config) & set(video_keys)


🟢 _repack_source_keys(data_config) & set(video_keys) 算了两次（missing 检查一次、required_video_keys 一次），可以算一次复用，避免两处以后不一致。

…ot resolution - compute_norm_stats parquet path now applies XtacUmiStateActions and XtacUmiRelativeActions, so drop_head_pose / use_relative_pose_actions configs get 2D gripper-state and relative-action stats instead of 29D absolute-pose stats that Normalize silently truncated. - Add config.resolve_dataset_root (DataConfig.dataset_root first, OPENPI_LEROBOT_DATASET_ROOT fallback), used by both training and compute_norm_stats so they read the same dataset. --dataset-root stays a one-off override and warns when it differs from the training root. - Add scripts/compute_norm_stats_test.py covering both. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Hubo1231 · 2026-11-03T02:40:10Z

已修改，改法：在 src/openpi/training/config.py 新增一个共用函数 resolve_dataset_root(data_config)，按评审的建议：config 里的 dataset_root
优先，没有就用环境变量 OPENPI_LEROBOT_DATASET_ROOT。

训练（data_loader.py 的 create_torch_dataset）改用这个函数，所以按 1011 yaml
的说明设环境变量，现在对训练也生效。训练还会在日志里打出实际用的路径。
compute_norm_stats 用同一个函数。--dataset-root 还能作为单次覆盖使用（现在 torch loader
也能拿到这个值），但训练看不到这个参数，所以它和训练会用的路径不一致时会打 warning。
_FULL_REFERENCE.yaml 里 dataset_root 的注释和脚本 docstring 也同步改了。

Vertax42 · 2026-11-03T03:14:42Z

感谢修复！3ecae31a 已核对：

✅ norm stats 计算空间：_STATE_ACTION_TRANSFORM_TYPES 已加入 XtacUmiStateActions / XtacUmiRelativeActions，按 config 顺序执行（先去头部位姿，再算相对动作），顺序正确。
✅ 数据集路径一致性：resolve_dataset_root() 训练和统计共用，--dataset-root 不一致时 warning，符合建议。
✅ 新增的 scripts/compute_norm_stats_test.py 本地 3 个测试全部通过。

以下问题这次 commit 还没涉及，麻烦看一下（修复或回复说明即可）：

🟠 建议合并前处理（真机安全）

examples/bi_elite_cs66_rt/real_env.py:136 reset 未开始/未完成时只打 warning 继续，且 go_to_start=False 时也会调 reset_to_initial_position。建议对齐 Rebot env，reset 失败直接 raise。
examples/bi_elite_cs66_rt/real_env.py:151 仅在 is_connected 时才 disconnect()，单臂连接失败/掉线时另一臂的 servo 循环、夹爪串口、相机不会释放。建议总是调用 facade 的 disconnect。

🟡 数据工具

rebase_umi_dataset.py:280 episodes 元数据里的 per-episode stats 仍是源数据 29D、Pico 坐标系，与新的 20D info.json 不符。
rebase_umi_dataset.py:219 --session-starts 未排序、未校验首个起点 ≤ 最小 episode 编号。

🟢 次要（可后续 PR）

data_loader.py:245 _ensure_reader 只复制了部分 reader 状态。
pi05_base_bi_elite_egg_carton_to_tray_1016_h100.yaml:3 注释里的 config 名写成了 _1012_h100。
xtac_umi_bi_flexiv_rizon4_rt/main.py:168 删除了 max_episode_steps，若是有意为之请在 PR 描述中说明。
data_loader.py:278 同一个集合交集计算了两次。

细节见上一轮的行内评论。

Vertax42

norm stats 和数据集路径两处已修好，感谢。合并前还需要处理两条真机安全问题：

examples/bi_elite_cs66_rt/real_env.py:136：reset 失败（1s 未开始或 15s 未完成）时只打 warning 继续，应改为 raise；go_to_start=False 时不应调用 reset_to_initial_position。
examples/bi_elite_cs66_rt/real_env.py:151：只连上一部分时 disconnect() 被跳过，已连接那只臂的 servo 循环、夹爪串口和相机不会释放。应总是调用 facade 的 disconnect（参考同一 PR 里 Rebot env 的写法）。

另外请对 rebase_umi_dataset.py 的两条问题（per-episode stats 未更新、--session-starts 未校验）修复或说明。🟢 次要问题可以放到后续 PR。详见之前的行内评论和跟进评论。

Vertax42 · 2026-11-05T12:47:27Z

@Hubo1231 @Ding-YF 跟进一下这个 PR 的进度。

目前 PR 因 changes requested 被阻塞，合并前需要处理两条 Elite 真机安全问题（详见上面的 review）：

examples/bi_elite_cs66_rt/real_env.py：reset 失败（1s 未开始 / 15s 未完成）时只打 warning 继续，应改为 raise；go_to_start=False 时不应调用 reset_to_initial_position。
examples/bi_elite_cs66_rt/real_env.py：只连上一部分时 disconnect() 被跳过，已连接那只臂的 servo 循环、夹爪串口、相机不会释放。应总是调用 facade 的 disconnect（参考本 PR 中 Rebot env 的写法）。

@Ding-YF 注意到 feat/elite-rerun（基于本 PR 分支）也改了 real_env.py，但这两处还在。由于该分支依赖本 PR 先合并，麻烦两位确认一下由谁来修、修在哪个分支上，避免重复或冲突。

另外 rebase_umi_dataset.py 的两条（per-episode stats 未更新、--session-starts 未校验）请修复或回复说明即可。谢谢！

Hubo1231 · 2026-11-06T07:29:25Z

修改内容：

• Elite reset：1 秒未启动或 15 秒未完成时抛出 TimeoutError；go_to_start=False 时跳过复位。
• 资源释放：始终调用 facade 的 disconnect()；连接失败时也执行清理，覆盖部分连接场景。
• episode 统计：根据重定位后的 20D 数据重算 state/action 的统计及分位数。
• session 起点：排序、去重，并校验能否覆盖首个 episode；非法参数不会删除已有目标数据集。

Vertax42

LGTM，已核对 d0990368 和 bb538359：

✅ Elite reset 失败抛 TimeoutError，go_to_start=False 跳过复位
✅ disconnect 总是调用 facade，connect 失败时也清理
✅ rebase 后 per-episode stats 按 20D 输出重算；--session-starts 排序去重并校验
✅ 本地 Elite env / rebase / norm-stats 测试 16 passed, 1 skipped；CI 通过

剩余 🟢 次要问题（_ensure_reader 状态复制、1016 yaml 注释中的 config 名、max_episode_steps、重复的集合交集）留给后续 PR。

@Ding-YF 本 PR 合并后请把 feat/elite-rerun rebase 到 main（real_env.py 两边都改过，试合并无冲突）。

Hubo1231 added 12 commits October 11, 2026 16:10

chore(training): remove JAX cuDNN attention and standardize YAML configs

c506757

feat(xtac-umi): support G1 relative pose training and selective video…

7ea4179

… loading

fix(bi-elite): align inference recipe with diagonal-08 assembly

bbaabdd

feat(xtac-umi): rebase datasets onto Flexiv start TCPs

98034c6

fix(xtac-umi): remove evaluation step limit and document launch

ded96ae

chore(config): remove XTac-UMI cube pickup example

3b62593

chore(config): add launch template for Pi0Config with training parame…

c55e6b5

…ters

feat(config): add Rebot B601 support and RTC training examples

1afdcba

feat(bi-rebot): add B601 RTC inference example and regression tests

cbf4519

chore(training): make LR schedules explicit and refresh training configs

c663f61

Retire unsafe XLA collective launch flags and document the findings. Include the whiteboard inspection config, remove the obsolete box assembly config, and preserve current prompt formatting updates.

Hubo1231 requested a review from Vertax42 October 29, 2026 05:22

Vertax42 reviewed Nov 2, 2026

View reviewed changes

Hubo1231 requested a review from Vertax42 November 3, 2026 02:40

Vertax42 requested changes Nov 3, 2026

View reviewed changes

Hubo1231 added 2 commits November 6, 2026 15:15

fix(bi-elite): abort failed resets and clean up partial connections

d099036

fix(xtac-umi): rebuild episode stats and validate session starts

bb53835

Hubo1231 requested a review from Vertax42 November 6, 2026 07:29

Vertax42 approved these changes Nov 6, 2026

View reviewed changes

Vertax42 merged commit c3831d9 into main Nov 6, 2026
1 check passed

Vertax42 deleted the chore/remove-jax-cudnn-attention branch November 6, 2026 08:55

Vertax42 mentioned this pull request Nov 12, 2026

feat(rlt): sync with main and align online RL with the RLT paper #26

Merged

Conversation

Hubo1231 commented Oct 29, 2026

主要变更

兼容性与迁移

测试

Uh oh!

Vertax42 left a comment

Choose a reason for hiding this comment

🔴 norm stats 在错误空间计算（scripts/compute_norm_stats.py 的 _STATE_ACTION_TRANSFORM_TYPES，该行不在 diff 中故放这里）

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Hubo1231 commented Nov 3, 2026

Uh oh!

Vertax42 commented Nov 3, 2026

Uh oh!

Vertax42 left a comment

Choose a reason for hiding this comment

Uh oh!

Vertax42 commented Nov 5, 2026

Uh oh!

Hubo1231 commented Nov 6, 2026

Uh oh!

Vertax42 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

🔴 norm stats 在错误空间计算（`scripts/compute_norm_stats.py` 的 `_STATE_ACTION_TRANSFORM_TYPES`，该行不在 diff 中故放这里）