From d1cc7d96ad21d074f2683250b02dc270a1966cbe Mon Sep 17 00:00:00 2001 From: huangyanpei <18200966375@163.com> Date: Wed, 23 Sep 2026 20:58:39 +0800 Subject: [PATCH] Add Wall-X serving and Turtle2 TCP WebSocket bridge --- .gitignore | 5 + TURTLE2_RETRAIN_DEPLOYMENT.md | 84 + TURTLE2_STARTUP.md | 159 + deploy_configs/7dh10_0/serving_config.yml | 260 ++ docs/WALLX_RTC_USAGE.md | 200 ++ docs/libero.md | 103 + scripts/README.md | 41 + scripts/run_serving.sh | 36 + scripts/run_serving_rtc.sh | 342 ++ scripts/tcp_ws_bridge.py | 833 +++++ scripts/tcp_ws_bridge_rtc.py | 872 +++++ setup.py | 2 + tests/test_lora_scale.py | 45 + tests/test_robot_preprocessor.py | 29 + tests/test_rtc_wallx.py | 149 + tests/test_run_serving_contract.py | 40 + tests/test_tcp_ws_bridge.py | 137 + tests/test_turtle2_inference_input.py | 47 + wall_x/_vendor/harrix/adapters/qwen_vlact.py | 6 +- .../serving/_wallx_infer/model_wrapper.py | 23 +- .../serving/_wallx_infer/model_wrapper_rtc.py | 802 +++++ .../harrix/serving/_wallx_infer/robot.py | 9 +- .../harrix/serving/launch_serving_rtc.py | 319 ++ .../harrix/serving/policy/wall_x_policy.py | 12 + .../serving/policy/wall_x_policy_rtc.py | 472 +++ wall_x/_vendor/harrix/serving/rtc_wallx.py | 119 + wall_x/_vendor/harrix/utils/ckpt_load.py | 102 +- wall_x/_vendor/harrix/utils/normalizer.py | 35 +- wall_x/config/model_config.py | 1 + wall_x/data/backends/lerobot/loader.py | 22 + wall_x/data/backends/lerobot/utils.py | 27 + wall_x/model/core/ops/moe.py | 17 +- wall_x/model/core/vla_mixin.py | 19 +- .../model/qact/qwen2_5/modeling_qwen2_5_vl.py | 16 +- .../qact/qwen2_5/modeling_qwen2_5_vl_act.py | 36 + .../qwen2_5/modeling_qwen2_5_vl_act_rtc.py | 3139 +++++++++++++++++ wall_x/trainer/adapters/base_adapter.py | 4 + wall_x/trainer/adapters/vla_model_adapter.py | 47 +- wall_x/trainer/fsdp_trainer/fsdp_trainer.py | 18 +- wall_x/trainer/trainer_utils.py | 2 + 40 files changed, 8591 insertions(+), 40 deletions(-) create mode 100644 TURTLE2_RETRAIN_DEPLOYMENT.md create mode 100644 TURTLE2_STARTUP.md create mode 100644 deploy_configs/7dh10_0/serving_config.yml create mode 100644 docs/WALLX_RTC_USAGE.md create mode 100644 docs/libero.md create mode 100755 scripts/run_serving_rtc.sh create mode 100644 scripts/tcp_ws_bridge.py create mode 100644 scripts/tcp_ws_bridge_rtc.py create mode 100644 tests/test_lora_scale.py create mode 100644 tests/test_robot_preprocessor.py create mode 100644 tests/test_rtc_wallx.py create mode 100644 tests/test_run_serving_contract.py create mode 100644 tests/test_tcp_ws_bridge.py create mode 100644 tests/test_turtle2_inference_input.py create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper_rtc.py create mode 100644 wall_x/_vendor/harrix/serving/launch_serving_rtc.py create mode 100644 wall_x/_vendor/harrix/serving/policy/wall_x_policy_rtc.py create mode 100644 wall_x/_vendor/harrix/serving/rtc_wallx.py create mode 100644 wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act_rtc.py diff --git a/.gitignore b/.gitignore index b243f82..0e37ac0 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,8 @@ # Runtime artifacts __pycache__/ *.py[cod] +*.bak +*.bak2 .ruff_cache/ .pytest_cache/ .mypy_cache/ @@ -29,6 +31,9 @@ ckpt/ *.safetensors *.mp4 +# Accidental files created from command-line arguments +--* + # Editor / OS noise .DS_Store .vscode/ diff --git a/TURTLE2_RETRAIN_DEPLOYMENT.md b/TURTLE2_RETRAIN_DEPLOYMENT.md new file mode 100644 index 0000000..4bede76 --- /dev/null +++ b/TURTLE2_RETRAIN_DEPLOYMENT.md @@ -0,0 +1,84 @@ +# Wall-X 重新训练后部署到 Turtle2：完整迁移流程 + +本流程对应 `/home/xiehaolv/huanghuagui/zhanyifeng/wall-x` 的当前 serving 与 TCP/WebSocket 桥接代码。当前已验收的目标 checkpoint 是 `/home/xiehaolv/huanghuagui/wall-x-model/finetuned/1`。实际启动命令见 [TURTLE2_STARTUP.md](TURTLE2_STARTUP.md)。训练侧的 `config.yml` 是训练快照，部署时保留原件；推理参数改在启动命令或部署代码中。 + +## 一、先判断这次究竟改变了什么 + +| 训练变化 | 部署时必须核对或修改 | +| --- | --- | +| 只改学习率、batch size、epoch、随机种子，数据格式和动作定义不变 | 更换 checkpoint 路径及其同目录 `config.yml`；核对 LoRA 参数、归一化文件和实际推理结果。桥接接口通常不变。 | +| 换数据集，但仍为三路相机、右臂 7D 原始状态/动作 | 更新训练配置中的数据根目录、`repo_id`、`norm_stats_path`、任务文本及相机 key mapping；部署使用新 checkpoint 自带的 normalizer，并更新 serving 的 `norm-key` 与任务文本。 | +| 相机数量、名称、顺序或实际安装位置变化 | 更新训练数据的相机映射与 serving `cam-names`；核对机器人发图顺序及桥接 `camera_left`/`camera_front`/`camera_right` 的物理对应。先做逐路图像检查，不要直接执行动作。 | +| `action_horizon_flow` 变化 | serving 和桥接两处 `--action-horizon` 都改为新值；重新计算动作裁剪与插值后的包长，并运行协议检查。启动脚本现会拒绝与训练配置不一致的 serving 长度。 | +| 绝对/相对动作、欧拉角/6D 旋转、左右臂、关节/末端、夹爪范围变化 | 训练配置、归一化、serving 的动作重建、桥接的 7D `follow{1,2}_pos` 协议都要重新验收。不能只替换权重路径。 | +| LoRA 的 rank、alpha、`use_rslora` 或目标模块变化 | 从训练服务器导出实际使用的 LoRA JSON 为 checkpoint 内的 `lora_config.json`；加载器依据 rank 和 alpha 计算合并系数。若使用逐层不同缩放，当前加载器会拒绝，需实现并测试后才能部署。 | + +## 二、训练前检查数据和配置 + +1. 保留一份原始 LeRobot v3 数据。当前训练配置指向 LeRobot v2.1 数据目录；继续使用同一训练环境时，先生成其可读取的 v2.1 训练副本，不要把 v3 目录直接写成 v2.1 路径。 +2. 核对每条 episode 的三路视频能解码、帧数与 `observation.state`/`action` 对齐，并剔除明显不合格的示教。当前物理映射是 camera1 左臂固定视角、camera2 面部视角、camera3 右臂视角。 +3. 当前训练 key mapping 为：`observation.images.faceImg → face_view`，`observation.images.leftImg → left_wrist_view`，`observation.images.rightImg → right_wrist_view`，`observation.state → state`，`action → action`。换数据集后按**新数据的真实字段**更新；字段名相同并不能证明镜头位置相同。 +4. 确认机器人状态和示教动作的坐标系、单位、夹爪开合方向相同。当前模型训练右臂，训练数据原始状态/动作为 7D；当前配置在训练加载时把欧拉角转换为 6D，并把动作转为相对当前末端的表示，最终形成“位置 3＋旋转 6＋夹爪 1＋填充 16”的 26D 布局。 +5. 在训练服务器上更新 `data.lerobot_config.root`、`repo_id`、`data.norm_stats_path`、`data.key_mappings.camera`、图像分辨率、`task.action_horizon` 和 `task.action_horizon_flow`。如果只改优化超参数，仍核对保存出的 `config.yml` 这些字段未意外变化。 +6. `norm_stats.json` 与 normalizer 必须由**这次训练所用的数据和配置**生成。当前 6D 配置部署时优先读取 checkpoint 内的 `normalizer_action.pth` 与 `normalizer_propri.pth`；不要沿用上一模型的两个 `.pth`。 +7. 保存训练实际使用的 LoRA JSON，记录 `lora_r`、`lora_alpha`、`use_rslora` 和目标模块。不要只保存配置文件的绝对路径，训练服务器路径到部署服务器通常不可达。 + +## 三、把 checkpoint 导出到部署服务器 + +为每次训练创建独立目录，例如 `/home/xiehaolv/huanghuagui/wall-x-model/finetuned/<新版本>`。不要覆盖正在部署的 checkpoint。至少保留： + +```text +model.safetensors # 或受加载器支持的 pytorch_model_fsdp.bin +config.yml # 对应这次训练的原始配置 +config.json # 模型结构 +preprocessor_config.json +tokenizer.json +tokenizer_config.json +vocab.json # 若本次 tokenizer 导出包含该文件 +normalizer_action.pth +normalizer_propri.pth +lora_config.json # 从这次训练实际使用的 LoRA JSON 原样复制 +``` + +`norm_stats.json`、`global_step.pth` 可一并保留，便于审计。`optimizer.pt`、`scheduler.pt`、`rng_state.pt` 不参与推理加载，可以不传到部署服务器。若训练保存的是仅含 LoRA adapter 的轻量文件，而不是该项目预期的完整 checkpoint，当前 `run_serving.sh` 不能直接当完整模型加载；先按训练代码的导出流程补齐基础权重。 + +当前 checkpoint `1` 具有完整 `model.safetensors`、配置、processor 和 normalizer，但**缺少训练原始 `lora_config.json`**。部署代码会兼容旧模型，警告后暂按 `alpha/r=2` 合并。其 LoRA 权重 rank 为 16；只有核对训练原始 JSON 的 alpha 为 32 且 `use_rslora=false`，才能确认这个系数。若原始 JSON 不在手，不要把示例配置当作训练凭据。 + +## 四、部署前静态核对 + +在部署服务器进入仓库目录： + +```bash +cd /home/xiehaolv/huanghuagui/zhanyifeng/wall-x +export WALLX_TEST_CHECKPOINT=/home/xiehaolv/huanghuagui/wall-x-model/finetuned/1 +export WALLX_TEST_PYTHON=/home/xiehaolv/huanghuagui/conda-envs/wallox_0.5/bin/python +PYTHONDONTWRITEBYTECODE=1 "$WALLX_TEST_PYTHON" -m pytest -q \ + tests/test_lora_scale.py \ + tests/test_tcp_ws_bridge.py \ + tests/test_run_serving_contract.py \ + tests/test_turtle2_inference_input.py +``` + +这些测试当前以 checkpoint `1` 的 32 步、三相机布局为基准；换成动作长度、相机数或特征布局不同的 checkpoint 时，应先更新相应测试的期望值，再运行。即使测试通过，也要继续做一次不发送动作的真实模型推理，因为测试不会证明新权重本身能完成抓取任务。 + +核对 `config.yml` 的 `task.dof_config`、`task.agent_pos_config`、`task.action_horizon_flow`、`data.key_mappings.camera` 和 `data.lerobot_config.repo_id`。将 serving 的 `--model-config.norm-key` 设置为 checkpoint 内 normalizer 的数据集键。当前键为 `pick_paper_lerobot_v21_3cam`；如果新 checkpoint 的键变了，要同步修改启动命令。 + +## 五、更新启动命令与不发送动作的推理 + +在 [TURTLE2_STARTUP.md](TURTLE2_STARTUP.md) 中将 `--checkpoint-path` 与 `--train-config-path` 都指向新目录；将服务端、桥接端的 `--action-horizon` 设为训练的 `task.action_horizon_flow`；相机名称、`norm-key`、任务指令同步为新训练实际值。保持 `--serialize-actions`。桥接端只接受序列化的 `follow2_pos`；它会拒绝缺少该字段的原始 `predict_action`，避免按旧布局误读填充维。 + +先启动 serving，观察 LoRA 合并日志和 `load_state_dict result: `。当前模型可能打印“词嵌入表被裁剪”和“checkpoint 内的 normalizer 参数未直接用作模型权重”：它们分别来自当前流式动作 tokenizer 与单独加载的 normalizer；需要关注是否还有其他非预期未匹配权重。若 LoRA scale 出现缺少 metadata 的警告，先核对上一步的训练 JSON。 + +然后启动**不带** `--allow-send` 的桥接端，在 `arm-pc` 以单包模式发送真实相机和状态。核对三路画面朝向、日志中的 `recv state`、serving 返回行数、桥接的 `predicted follow2 T=<插值后长度>`、右臂首尾目标及夹爪范围。对于当前 32 步、`action-end-ratio=0.2`、插值倍数 32，桥接预期返回 192 点。模型推理可用不等于动作方向正确，现场仍须验证坐标与夹爪符号。 + +## 六、单包真机与连续执行 + +单包时停止 dry-run 桥接进程，按启动页相同命令增加 `--allow-send`，保留 `--max-action-cycles 1`、`--require-right-feedback`、`--clip-action-delta` 和当前小幅位置/旋转限制。现场人员确认机器人周围无障碍、右臂控制器已使能、升降柱目标高度合适，能即时急停。观察 `/follow2_pos_back` 与 `/joint_information2`，并核对夹爪开合方向与物体接触情况。 + +单包结果符合预期后，才把 `--max-action-cycles` 改为 `0` 进入连续推理；`0` 意味着不限动作包数。当前桥接的反馈检查只验证右臂末端是否发生运动，**不提供可靠的物品滑脱检测或有限次数自动重抓**。如果任务需要这些能力，应单独设计感知判定和执行状态机，而不能把 `max-action-cycles` 当作重抓次数。 + +停止顺序：机器人端先退出，再停止桥接端和 serving。保留上一版 checkpoint 与其匹配的启动命令作为回退入口；回退必须同时恢复权重、训练配置、动作长度、相机和 normalizer 参数，不能只切换 `model.safetensors`。 + +## 七、当前验证范围 + +在部署服务器 GPU 1 上，checkpoint `1` 已完成载入，日志显示 72 组 LoRA 合并和 ``。三路 448×448 合成图像已通过完整回环：TCP 状态/图像输入 → WebSocket serving → Turtle 动作序列化 → TCP 动作回传；返回包含 `follow1_pos`、`follow2_pos`、`head_pos`、`lift`、`car_pose`，右臂轨迹为 192 点，左臂 192 点均保持输入位姿。另有 12 项回归测试通过。三路 448×448 图像此前触发的 1024 token 截断已修复为推理时保留完整图像 token。回环只把动作发给本机合成客户端，没有连接或驱动机器人。真实相机内容、抓取效果、右臂执行方向、夹爪范围和训练原始 LoRA alpha 仍需要现场与训练产物确认。 diff --git a/TURTLE2_STARTUP.md b/TURTLE2_STARTUP.md new file mode 100644 index 0000000..9541540 --- /dev/null +++ b/TURTLE2_STARTUP.md @@ -0,0 +1,159 @@ +# Wall-X + Turtle2 启动命令 + +本页对应 checkpoint `/home/xiehaolv/huanghuagui/wall-x-model/finetuned/1`。推流端和桥接端运行在推理服务器上，机器人端运行在 `arm-pc` 上。更换数据集或训练配置时，先按 [重训与部署迁移流程](TURTLE2_RETRAIN_DEPLOYMENT.md) 核对配置，再更新本页命令。 + +相机对应关系：机器人 camera1/`camera_left` = 左臂固定视角 → `left_wrist_view`；camera2/`camera_front` = 面部视角 → `face_view`；camera3/`camera_right` = 右臂视角 → `right_wrist_view`。当前模型只输出右臂动作，桥接端保持左臂不动，并对底盘发送相对零位移。 + +## 1. Wall-X 推流服务端 + +如果 `32196` 端口上还没有 serving 服务，在推理服务器终端执行： + +```bash +cd /home/xiehaolv/huanghuagui/zhanyifeng/wall-x + +PYTHON_BIN=/home/xiehaolv/huanghuagui/conda-envs/wallox_0.5/bin/python \ +bash scripts/run_serving.sh \ + --checkpoint-path /home/xiehaolv/huanghuagui/wall-x-model/finetuned/19 \ + --train-config-path /home/xiehaolv/huanghuagui/wall-x-model/finetuned/19/config.yml \ + --port 32196 \ + --cuda-id 1 \ + --env X2ROBOT \ + --robot-type turtle \ + --serialize-actions \ + --action-horizon 32 \ + --robot-action-interpolate-multiplier 1 \ + --robot-action-end-ratio 1.0 \ + --model-device cuda \ + -- \ + --model-config.cam-names face_view left_wrist_view right_wrist_view \ + --model-config.norm-key pick_paper_lerobot_v21_3cam +``` + +该命令使用 checkpoint `/home/xiehaolv/huanghuagui/wall-x-model/finetuned/1` 的训练配置：三路相机最长边均为 448，动作长度 32，右臂相对位置 + 相对 6D 旋转 + 夹爪，后接 16 维虚拟填充。服务端会使用 checkpoint 内的两份 26 维 `.pth` 归一化文件。`run_serving.sh` 会在载入模型前核对命令里的动作长度与 `config.yml`，不一致会直接退出。 + +新 checkpoint 尚未携带训练时的 `lora_config.json`，因此加载日志会提示暂用旧系数 `alpha/r=2`。这不妨碍服务启动，但在真机执行前，应核对训练时使用的 LoRA JSON 是否为 `lora_r=16`、`lora_alpha=32` 且未启用 `use_rslora`；如不同，复制原始 JSON 到 checkpoint 目录并命名为 `lora_config.json`，加载器会自动按该文件计算系数。不要用部署仓库里的同名示例文件冒充训练原件。 + +## 2. TCP/WebSocket 桥接端 + +第一次先进行不发送动作的推理与协议检查： + +```bash +cd /home/xiehaolv/huanghuagui/zhanyifeng/wall-x + +/home/xiehaolv/huanghuagui/conda-envs/wallox_0.5/bin/python \ + scripts/tcp_ws_bridge.py \ + --ws-url ws://127.0.0.1:32196 \ + --tcp-port 30124 \ + --instruction "pick up the paper towel" \ + --action-horizon 32 \ + --fixed-lift 0.1 \ + --action-end-ratio 0.2 \ + --action-interpolate-multiplier 32 \ + --max-position-delta 0.005 \ + --max-rotation-delta 0.0125 \ + --clip-action-delta \ + --max-action-cycles 1 \ + --require-right-feedback +``` +/home/xiehaolv/huanghuagui/conda-envs/wallox_0.5/bin/python \ + scripts/tcp_ws_bridge.py \ + --ws-url ws://127.0.0.1:32196 \ + --tcp-port 30124 \ + --instruction "pick up the paper" \ + --action-horizon 32 \ + --fixed-lift 0.3 \ + --action-end-ratio 0.7 \ + --action-interpolate-multiplier 6 \ + --max-position-delta 2 \ + --max-rotation-delta 2 \ + --clip-action-delta \ + --max-action-cycles 0 \ + --require-right-feedback \ + --allow-send + + + + + + + + --action-horizon 32 \ + --fixed-lift 0.1 \ + --require-right-feedback \ + --max-position-delta 0.010 \ + --max-rotation-delta 0.0250 \ + --clip-action-delta \ + --max-action-cycles 0 \ + --allow-send + + + +`--clip-action-delta` 会保持模型轨迹形状，并把整个推理动作包相对当前机械臂位姿的最大变化限制在上述范围。它不会再把每个采样点分别累加成固定步长斜坡。 + +上述命令不带 `--allow-send`，只检查三路图像、推理返回、32 步序列和 `follow2_pos` 解析，不发送动作包。`--fixed-lift 0.1` 会把升降柱目标设为 0.1 米；若当前高度不是 0.1 米，在真机试验前确认这项高度变化可执行。底盘默认发送相对零位移 `[0, 0, 0]`。 + +先确认 serving 日志出现 LoRA 合并、`load_state_dict result: `；桥接日志出现 `predicted follow2 T=192`，并检查右臂第一点、末点及夹爪数值。若报 `Image features and image tokens do not match`，不要发送动作，检查三路图像尺寸与当前推理预处理代码。若只有 `predict_action` 而没有 `follow2_pos`，检查 serving 是否带 `--serialize-actions`，桥接端会拒绝未序列化动作。 + +确认以上结果和现场状态后，**停止 dry-run 桥接端**，用完全相同的命令重新启动，并在最后一行改为： + +```bash + --require-right-feedback \ + --allow-send +``` + +保持 `--max-action-cycles 1`，仅发送一个动作包。机器人端需由现场人员确认空域、夹爪和升降柱状态后启动。执行完检查 `/follow2_pos_back` 与 `/joint_information2` 的右臂反馈，再考虑连续运行。 + +确认一次动作后 `/follow2_pos_back` 与 `/joint_information2` 都产生了相应的小变化，且右臂末端和夹爪动作方向符合任务，再把 `--max-action-cycles 1` 改为 `--max-action-cycles 0` 并重启桥接端，进入连续推理。`0` 表示连接内动作包数量不限，**并不是重抓次数上限**。 + +不要添加 `--allow-base-motion`。默认情况下桥接端发送相对零位移，禁止模型动作驱动底盘。 + +连续运行时桥接端使用以下参数（单包命令中的其余参数保持不变）： + +```bash + --action-horizon 32 \ + --fixed-lift 0.1 \ --max-position-delta 0.005 \ + --max-rotation-delta 0.0125 \ + --clip-action-delta \ + --max-action-cycles 0 \ + --allow-send +``` + +## 3. Turtle2 机器人端 + +在 `arm-pc` 上执行： + +```bash +source /opt/ros/noetic/setup.bash +source /home/arm/prj/turtle2/modules/devel/setup.bash + +export PYTHONPATH=/home/arm/prj/turtle2/modules/src/turtlesys/turtle_monitor/scripts:$PYTHONPATH + +infer 172.25.5.237 30124 +``` + +## 模式选择 + +- 单包 dry-run 时，在推理模式提示输入 `2`；桥接端不带 `--allow-send`，因此不执行机械臂动作。完成解析检查后，停止旧桥接端并带 `--allow-send` 重复单包测试。 +- 当前 checkpoint 只训练右臂。桥接默认保持左臂的当前位姿；不要添加 `--allow-left-arm-motion`。 +- 单包实机验证建议增加 `--require-right-feedback`。若下一个机器人状态中右臂末端反馈没有变化，桥接会在发送下一包前停止，并提示检查 `/follow_pos_cmd_2`、`/follow2_pos_back`、`/joint_information2` 与右臂控制器。 +- 连续推理时输入 `1`，同时桥接端必须使用 `--max-action-cycles 0`。机器人端仍会调用底盘 pose 分支，但桥接端发送的是相对零位移，因此会看到 `exec car pose [0.0, 0.0, 0.0]` 日志。 +- `mode1: 0`、`mode2: 0` 是现有机器人脚本固定填写的字段；R5 末端控制回调不使用它们，不是机械臂停止的原因。 +- 桥接端每次打印 `recv state`、`predicted ... T=<插值后长度>`、`sent action back`，表示完成了一次动作包；32 步动作按默认旧版节奏在 `--action-end-ratio 0.2`、`--action-interpolate-multiplier 32` 下通常会扩展为约192个点，具体以日志为准。`Robot closed the connection` 表示机器人端或桥接端已退出，不是推理失败。 +- `--clip-action-delta` 目前只限制右臂位置和姿态，不限制夹爪。日志中若出现 `gripper` 大幅跳变（例如 `0.8` 或 `-0.8`），先确认夹爪动作范围和现场安全，再继续执行抓取任务。 +- `--require-right-feedback` 只检查右臂末端是否按上一包发生运动，不检测物品滑脱或抓取成功。 +- `/head/control unavailable` 是头部服务缺失产生的独立错误；它不会改变右臂末端控制，但会造成高频日志。 + +## 验证要求（checkpoint `/home/xiehaolv/huanghuagui/wall-x-model/finetuned/1`） + +- 先用不带 `--allow-send` 的命令确认 serving 日志出现 LoRA 合并、`All keys matched successfully`，桥接日志出现 `predicted follow2 T=192`。现在的 checkpoint 如出现 `LoRA scale metadata is missing` 警告，按上文核对训练时的原始 LoRA JSON 后再做真机试验。 +- 单包实机验证只在机器人端预览确认后增加 `--allow-send`，并保持 `--max-action-cycles 1`。 +- 确认 `/follow2_pos_back` 与 `/joint_information2` 的右臂反馈发生变化后，再改为 `--max-action-cycles 0` 连续执行。 + +本次服务器侧已验证：checkpoint 能在 GPU 1 上加载；三路 640×480 合成图像的推理返回 33 行 `follow2_pos`，桥接端解析为 192 点；三路 448×448 合成图像的完整图像 token 也通过回归测试。以上输入均为合成观测，不能证明真实相机映射、模型抓取效果或真机动作方向。 + +## 端口和停止 + +- Wall-X serving：`32196` +- TCP 桥接端：`30124` +- 机器人连接地址：`172.25.5.237:30124` +- 停止服务：在对应终端按 `Ctrl+C`，建议先停止机器人端，再停止桥接端和推流端。 diff --git a/deploy_configs/7dh10_0/serving_config.yml b/deploy_configs/7dh10_0/serving_config.yml new file mode 100644 index 0000000..97e5423 --- /dev/null +++ b/deploy_configs/7dh10_0/serving_config.yml @@ -0,0 +1,260 @@ +_raw_data: + dataset_type: lerobot + key_mappings: + action: action + camera: + observation.images.faceImg: face_view + observation.images.leftImg: left_wrist_view + observation.images.rightImg: right_wrist_view + state: observation.state + lerobot_config: + repo_id: pick_paper_lerobot_mixed_v21_3cam + root: /data1/students/huangyanpei/walloss_bag_to_lerobot_v3_3cam/pick_paper_lerobot_v3_3cam_mixed_v21 + max_length: 1024 + norm_stats_path: /data1/students/huangyanpei/walloss_bag_to_lerobot_v3_3cam/pick_paper_lerobot_v3_3cam_mixed_v21/norm_stats.json + num_workers: 4 + resolution: + face_view: 448 + left_wrist_view: 448 + right_wrist_view: 448 + train_test_split: 0.95 +_raw_yaml: + _attn_implementation: sdpa + checkpoint: + resume_from: /data5/students/huanghuagui/walloss_Tweak/migration/models/wall-oss-0.5/model.safetensors + save_path: /data1/students/huangyanpei/walloss_bag_to_lerobot_v3_3cam/pick_paper_lerobot_v3_3cam_mixed_v21/output_weights_7d_H10 + data: + dataset_type: lerobot + key_mappings: + action: action + camera: + observation.images.faceImg: face_view + observation.images.leftImg: left_wrist_view + observation.images.rightImg: right_wrist_view + state: observation.state + lerobot_config: + repo_id: pick_paper_lerobot_mixed_v21_3cam + root: /data1/students/huangyanpei/walloss_bag_to_lerobot_v3_3cam/pick_paper_lerobot_v3_3cam_mixed_v21 + max_length: 1024 + norm_stats_path: /data1/students/huangyanpei/walloss_bag_to_lerobot_v3_3cam/pick_paper_lerobot_v3_3cam_mixed_v21/norm_stats.json + num_workers: 4 + resolution: + face_view: 448 + left_wrist_view: 448 + right_wrist_view: 448 + train_test_split: 0.95 + debug: + nvtx: false + profile: false + distributed: + bf16: true + use_fsdp: true + use_mixed_precision: true + hyperparams: + batch_size_per_gpu: 8 + gradient_accumulation_steps: 8 + num_epoch: 6 + optimizer: + betas: + - 0.9 + - 0.95 + enable_grad_clip: true + eps: 1.0e-08 + learning_rate: 5.0e-05 + max_grad_norm: 1.0 + optimizer_type: adamw + weight_decay: 1.0e-08 + scheduler: + min_lr: 1.0e-06 + num_training_steps: 3200 + num_warmup_steps: 200 + scheduler_type: cosine + seed: 10222 + logging: + epoch_save_interval: 1 + log_entity: null + log_interval: 1 + log_name: pick_paper_3cam_mixed_lora + log_project: null + save_interval: 700 + use_wandb: false + val_interval: 700 + model: + ar_loss_weight: 0.01 + attn_deterministic: true + backbone: qwen2_5 + config_path: /data5/students/huanghuagui/walloss_Tweak/migration/models/wall-oss-0.5/config.json + flow_loss_weight: 1.0 + lora_config_path: /data5/students/huanghuagui/huangyanpei/wall-x/workspace/example/wall-oss-0.5-lora.json + pretrained_path: /data5/students/huanghuagui/walloss_Tweak/migration/models/Qwen2.5-VL-3B-Instruct + processor_path: /data5/students/huanghuagui/walloss_Tweak/migration/models/Qwen2.5-VL-3B-Instruct + use_ema: false + model_type: qwen2_5 + task: + action_horizon: 10 + action_horizon_flow: 10 + agent_pos_config: + action_padding: 19 + follow_right_ee_cartesian_pos: 3 + follow_right_ee_rotation: 3 + follow_right_gripper: 1 + ar_dof_config: + action_padding: 19 + follow_right_ee_cartesian_pos_relative: 3 + follow_right_ee_rotation_relative: 3 + follow_right_gripper: 1 + dof_config: + action_padding: 19 + follow_right_ee_cartesian_pos_relative: 3 + follow_right_ee_rotation_relative: 3 + follow_right_gripper: 1 + use_state_string_representation: false +checkpoint: + resume_data: null + resume_ema: null + resume_from: /data5/students/huanghuagui/walloss_Tweak/migration/models/wall-oss-0.5/model.safetensors + resume_model: null + resume_optimizer: null + resume_rng: null + resume_scheduler: null + resume_step: null + save_path: /data1/students/huangyanpei/walloss_bag_to_lerobot_v3_3cam/pick_paper_lerobot_v3_3cam_mixed_v21/output_weights_7d_H10 + validate_first: false +data: + action_tokenizer_path: null + camera_name_mapping: null + dataset_type: lerobot + key_mappings: + action: action + camera: + observation.images.faceImg: face_view + observation.images.leftImg: left_wrist_view + observation.images.rightImg: right_wrist_view + state: observation.state + lerobot_config: + repo_id: pick_paper_lerobot_mixed_v21_3cam + root: /data1/students/huangyanpei/walloss_bag_to_lerobot_v3_3cam/pick_paper_lerobot_v3_3cam_mixed_v21 + noise_scheduler: null + norm_stats_path: /data1/students/huangyanpei/walloss_bag_to_lerobot_v3_3cam/pick_paper_lerobot_v3_3cam_mixed_v21/norm_stats.json + normalizer_config: null + num_workers: 4 + padding_side: left + priority_order: null + resolution: + face_view: 448 + left_wrist_view: 448 + right_wrist_view: 448 + train_test_split: 0.95 + use_fast_tokenizer: false +dataset_path: null +debug: + enable_mfu: false + enable_mfu_profile: false + nvtx: false + profile: false + profile_active_iters: 3 + profile_save_path: ./profile + profile_wait_iters: 1 + profile_warmup_iters: 1 + save_debug_batch_path: null + show_time_details: false + visualize_sample: false +distributed: + bf16: true + broadcast_buffers: true + bucket_cap_mb: 25 + find_unused_parameters: false + fsdp_backward_prefetch: backward_pre + fsdp_cpu_offload: false + fsdp_forward_prefetch: false + fsdp_hsdp_replicate_size: null + fsdp_limit_all_gathers: true + fsdp_reduce_dtype: bf16 + fsdp_save_policy: full + fsdp_sharding_strategy: full_shard + fsdp_sync_module_states: true + fsdp_use_orig_params: true + use_amp: false + use_fsdp: true + use_gradient_checkpointing: false + use_gradient_checkpointing_offload: false + use_mixed_precision: true + use_selective_recompute: false +hyperparams: + batch_size_per_gpu: 8 + gradient_accumulation_steps: 8 + num_epoch: 6 + optimizer: + action_expert_learning_rate: null + action_lr_keywords: null + betas: !!python/tuple + - 0.9 + - 0.95 + enable_grad_clip: true + eps: 1.0e-08 + foreach: null + fused: true + learning_rate: 5.0e-05 + lr_groups: null + max_grad_norm: 1.0 + optimizer_type: adamw + train_action_expert_only: false + weight_decay: 1.0e-08 + scheduler: + min_lr: 1.0e-06 + num_training_steps: 3200 + num_warmup_steps: 200 + scheduler_type: cosine + seed: 10222 +logging: + epoch_save_interval: 1 + gc_interval_steps: 1000 + ignore_until_interval: 0 + log_entity: null + log_interval: 1 + log_name: pick_paper_3cam_mixed_lora + log_project: null + loss_log_smooth_window: 1 + save_interval: 700 + use_wandb: false + val_interval: 700 + wandb_offline: false +model: + action_tokenizer_checkpoint_path: null + action_tokenizer_config_dir: null + action_tokenizer_path: null + action_tokenizer_type: null + ar_loss_weight: 0.01 + attn_deterministic: true + attn_implementation: null + backbone: qwen2_5 + config_path: /data5/students/huanghuagui/walloss_Tweak/migration/models/wall-oss-0.5/config.json + customized_robot_config: null + enable_customized_robot_config: false + flow_loss_weight: 1.0 + lora_config_path: /data5/students/huanghuagui/huangyanpei/wall-x/workspace/example/wall-oss-0.5-lora.json + new_special_tokens: null + pretrained_path: /data5/students/huanghuagui/walloss_Tweak/migration/models/Qwen2.5-VL-3B-Instruct + processor_path: /data5/students/huanghuagui/walloss_Tweak/migration/models/Qwen2.5-VL-3B-Instruct + use_ema: false +model_type: qwen2_5 +task: + action_horizon: 10 + action_horizon_flow: 10 + agent_pos_config: + action_padding: 19 + follow_right_ee_cartesian_pos: 3 + follow_right_ee_rotation: 3 + follow_right_gripper: 1 + ar_dof_config: + action_padding: 19 + follow_right_ee_cartesian_pos_relative: 3 + follow_right_ee_rotation_relative: 3 + follow_right_gripper: 1 + dof_config: + action_padding: 19 + follow_right_ee_cartesian_pos_relative: 3 + follow_right_ee_rotation_relative: 3 + follow_right_gripper: 1 + noise_scheduler: null + use_state_string_representation: false diff --git a/docs/WALLX_RTC_USAGE.md b/docs/WALLX_RTC_USAGE.md new file mode 100644 index 0000000..8b2a360 --- /dev/null +++ b/docs/WALLX_RTC_USAGE.md @@ -0,0 +1,200 @@ +# Wall-X Turtle2 RTC 使用说明 + +本文描述独立 RTC 副本链路。原 Wall-X serving、bridge 和 Turtle2 客户端文件未被覆盖。 + +## 1. 实现结构 + +RTC 链路由三部分组成： + +1. `launch_serving_rtc.py` 加载 RTC policy/model 副本； +2. `tcp_ws_bridge_rtc.py` 透传 session、request、已消费步数和推理延迟； +3. `SlaveRos2SocketPort_TURTLE2_RTC.py` 用独立控制线程在推理期间继续执行旧动作队列。 + +RTC guidance 在 Wall-X 的归一化相对动作空间中执行。上一段剩余绝对末端轨迹会基于最新 `follow2_pos` 重新转换为相对平移和相对 6D 旋转，然后才送入下一轮 flow inference。 + +当前仅支持： + +- `model_type=qwen2_5`； +- `infer_mode=flow`； +- 单 GPU 串行推理； +- Turtle2 WebSocket/TCP 部署链路。 + +RTC serving 不支持 dynamic batching。RTC VJP 会增加推理时间和显存占用。 + +## 2. 参数关系 + +推荐初始参数： + +```text +action_horizon = 32 +rtc_execution_horizon = 6 # 模型步 +action_interpolate_multiplier = 6 +action_end_ratio = 1.0 +``` + +含义： + +- 每次模型生成完整 32 步，作为推理超时时的备用队列； +- 真机执行约 6 个模型步后申请下一次推理； +- 每个模型步在 bridge 端插值为 6 个控制点； +- 因此正常情况下执行约 `6 × 6 = 36` 个控制点后重规划； +- 完整备用队列是 `32 × 6 = 192` 个控制点。 + +RTC 模式不要把 `action_end_ratio` 设为 `0.2`。这会删除备用队列后半段，推理稍慢时容易出现队列下溢。 + +## 3. 启动 RTC serving + +建议先使用不同端口，不影响原 serving： + +```bash +cd /home/xiehaolv/huanghuagui/zhanyifeng/wall-x + +bash scripts/run_serving_rtc.sh \ + --checkpoint-path /home/xiehaolv/huanghuagui/wall-x-model/finetuned/7 \ + --train-config-path /home/xiehaolv/huanghuagui/wall-x-model/finetuned/7/config.yml \ + --port 32197 \ + --env X2ROBOT \ + --cuda-id 0 \ + --serialize-actions \ + --robot-type turtle \ + --action-horizon 32 \ + --robot-action-interpolate-multiplier 1 \ + --robot-action-end-ratio 1.0 \ + --rtc-execution-horizon 6 \ + --rtc-max-guidance-weight 10 \ + --rtc-prefix-schedule linear \ + -- \ + --model-config.cam-names face_view left_wrist_view right_wrist_view \ + --model-config.norm-key pick_paper_lerobot_v21_3cam +``` + +不要添加 `--max-batch-size`。RTC policy 会保存每个机器人 session 的上一段动作。 + +## 4. 首次启动 RTC bridge：只做 dry-run + +首次运行不要添加 `--allow-send`： + +```bash +cd /home/xiehaolv/huanghuagui/zhanyifeng/wall-x + +/home/xiehaolv/huanghuagui/conda-envs/wallox_0.5/bin/python \ + scripts/tcp_ws_bridge_rtc.py \ + --ws-url ws://127.0.0.1:32197 \ + --tcp-port 30126 \ + --instruction "pick up the paper" \ + --action-horizon 32 \ + --action-end-ratio 1.0 \ + --action-interpolate-multiplier 6 \ + --rtc-execution-horizon 6 \ + --fixed-lift 0.3 \ + --max-position-delta 0.05 \ + --max-rotation-delta 0.25 \ + --clip-action-delta \ + --max-action-cycles 0 +``` + +dry-run 会在第一轮推理后关闭连接，这是预期行为。确认图像、状态和动作维度正常后，再进行真机发送测试。 + +## 5. 部署真机 RTC 副本 + +本机生成的真机文件： + +```text +/home/xiehaolv/huanghuagui/robot_self/turtle2/modules/src/ui/model2arm/communicationPort/SlaveRos2SocketPort_TURTLE2_RTC.py +/home/xiehaolv/huanghuagui/robot_self/turtle2/modules/src/ui/model2arm/communicationPort/infer_rtc.py +``` + +复制到真机对应目录，不覆盖原文件： + +```bash +scp \ + /home/xiehaolv/huanghuagui/robot_self/turtle2/modules/src/ui/model2arm/communicationPort/SlaveRos2SocketPort_TURTLE2_RTC.py \ + arm@172.28.57.93:/home/arm/prj/turtle2/modules/src/ui/model2arm/communicationPort/ + +scp \ + /home/xiehaolv/huanghuagui/robot_self/turtle2/modules/src/ui/model2arm/communicationPort/infer_rtc.py \ + arm@172.28.57.93:/home/arm/prj/turtle2/modules/src/ui/model2arm/communicationPort/ +``` + +## 6. 真机发送测试 + +重新启动 RTC bridge，并增加 `--allow-send`： + +```bash +/home/xiehaolv/huanghuagui/conda-envs/wallox_0.5/bin/python \ + scripts/tcp_ws_bridge_rtc.py \ + --ws-url ws://127.0.0.1:32197 \ + --tcp-port 30126 \ + --instruction "pick up the paper" \ + --action-horizon 32 \ + --action-end-ratio 1.0 \ + --action-interpolate-multiplier 6 \ + --rtc-execution-horizon 6 \ + --fixed-lift 0.3 \ + --max-position-delta 0.05 \ + --max-rotation-delta 0.25 \ + --clip-action-delta \ + --max-action-cycles 0 \ + --allow-send +``` + +初次 RTC 测试不要添加 `--require-right-feedback`。如果右臂反馈 topic 没有变化，该选项会让 bridge 主动关闭 TCP，真机随后会出现 `BrokenPipeError`。确认 `/follow2_pos_back` 正常后再启用。 + +在真机终端运行： + +```bash +source /opt/ros/noetic/setup.bash +source /home/arm/prj/turtle2/modules/devel/setup.bash +cd /home/arm/prj/turtle2/modules/src/ui/model2arm/communicationPort +python3 infer_rtc.py <推理服务器IP> 30126 +``` + +## 7. 正常日志判据 + +第一轮没有旧 chunk，应看到： + +```text +guided=False +``` + +第二轮开始应看到： + +```text +guided=True +RTC merged request_id=... delay_points=... delay_model_steps=... remaining=... +``` + +bridge 应持续输出不同的 request ID，不应出现： + +```text +stale/mismatched request_id +``` + +如果出现： + +```text +RTC response ... is stale +``` + +说明推理时间已经超过完整 32 步备用队列。应先降低 flow inference steps、提高控制点周期或增加 action horizon，而不是继续提高 guidance weight。 + +## 8. 回退 + +RTC 使用独立文件和端口。停止 RTC 三个进程后，重新运行原始命令即可回退： + +```text +launch_serving.py / run_serving.sh +tcp_ws_bridge.py +infer.py +``` + +不需要恢复或覆盖任何原始文件。 + +## 9. 离线验证 + +```bash +cd /home/xiehaolv/huanghuagui/zhanyifeng/wall-x +/home/xiehaolv/huanghuagui/conda-envs/wallox_0.5/bin/python \ + -m pytest -q tests/test_rtc_wallx.py tests/test_tcp_ws_bridge.py +``` + diff --git a/docs/libero.md b/docs/libero.md new file mode 100644 index 0000000..11a86d7 --- /dev/null +++ b/docs/libero.md @@ -0,0 +1,103 @@ +三、仿真任务全流程（LIBERO 数据集微调+评估） +适用于单臂机器人仿真场景，基于 LeRobot 格式 LIBERO 数据集完成模型微调与仿真测试 +3.1 下载 LIBERO 数据集 +libero的huggingface路径：https://huggingface.co/datasets/lerobot/libero +huggingface-cli download lerobot/libero \ + --repo-type dataset \ + --local-dir /path/to/libero_all +3.2 配置训练 YAML 文件 +复制官方示例配置，替换所有本地路径参数，适配本地环境 +cp workspace/example/libero.yml /path/to/my_libero_config.yml +核心必填配置（需完整替换）： +参考：https://github.com/X-Square-Robot/wall-x/blob/main/workspace/example/libero.yml +model: + config_path: /path/to/wall-oss-0.5/config.json # 模型配置文件路径 + processor_path: /path/to/Qwen2.5-VL-3B-Instruct # 处理器路径 + pretrained_path: /path/to/Qwen2.5-VL-3B-Instruct # 预训练基座路径 + +data: + lerobot_config: + repo_id: /path/to/libero_all # 本地LIBERO数据集根目录 + norm_stats_path: /path/to/libero_all_norm_stats.json # 归一化统计文件 + key_mappings: # 数据集字段映射（固定适配LIBERO） + camera: + observation.images.faceImg: face_view + observation.images.rightImg: right_wrist_view + state: observation.state + action: action + +checkpoint: + save_path: /path/to/libero_training_output # 训练结果保存路径 + resume_from: /path/to/wall-oss-0.5/model.safetensors # 预训练权重路径 + +dof_config配置 + # Libero delta action: pos3 + rot3 + gripper1 = 7, plus action_padding(19) = 26. + dof_config: + master_right_ee_cartesian_pos: 3 # delta position + master_right_ee_rotation: 3 # delta rotation (ZYX euler) + master_right_gripper: 1 + action_padding: 19 + ar_dof_config: + master_right_ee_cartesian_pos: 3 + master_right_ee_rotation: 3 + master_right_gripper: 1 + action_padding: 19 + # State: pos3 + rot3 + gripper2 = 8, plus action_padding(18) = 26. + agent_pos_config: + follow_right_ee_cartesian_pos: 3 + follow_right_ee_rotation: 3 + follow_right_gripper: 2 + action_padding: 18 + action_horizon: 10 + action_horizon_flow: 10 + + +关键适配说明：LIBERO 为7维单臂动作，需通过配置中 action_padding 补全至26维，匹配模型预训练维度（参考示例配置注释） +3.3 生成数据集归一化统计文件 +训练前必须执行，生成数据均值、方差统计文件，保证训练稳定性 +python scripts/compute_norm_stats.py \ + --train_config /path/to/my_libero_config.yml \ + --data_root /path/to/libero_all \ + --output_path /path/to/libero_all_norm_stats.json +执行完成后，需确认配置文件中 norm_stats_path 与输出路径一致 +3.4 启动模型微调训练 +硬件要求：单卡训练最低需要 48G 显存，多卡训练推荐开启 FSDP 分布式训练 +单GPU训练命令 +CUDA_VISIBLE_DEVICES=0 \ + python wall_x/trainer/fsdp_trainer/train_fsdp.py \ + --config /path/to/my_libero_config.yml +多GPU训练（推荐） +CUDA_VISIBLE_DEVICES=0,1,2,3 \ + torchrun --nproc_per_node=4 \ + wall_x/trainer/fsdp_trainer/train_fsdp.py \ + --config /path/to/my_libero_config.yml +断点合并说明 +多卡FSDP训练会生成分片权重文件，推理前需合并为完整权重： +python scripts/merge_sharded_weights.py \ + /path/to/sharded_checkpoint \ + /path/to/merged_checkpoint +3.5 LIBERO 仿真推理评估 +批量测试模型在仿真场景的任务完成效果，支持指定任务套件、测试次数 +前置依赖校验（未安装需重新执行仿真依赖安装命令）：需提前安装 robosuite、MuJoCo、PyOpenGL 等仿真组件 +常规批量评估 +参考脚本：https://github.com/X-Square-Robot/wall-x/blob/main/scripts/run_libero.sh +CHECKPOINT_PATH=/path/to/checkpoint \ +TRAIN_CONFIG_PATH=/path/to/my_libero_config.yml \ +TASK_SUITE_NAME=libero_spatial \ +NUM_TRIALS_PER_TASK=50 \ +bash scripts/run_libero.sh +快速冒烟测试（调试用，单任务1次测试） +SMOKE=1 CHECKPOINT_PATH=/path/to/checkpoint bash scripts/run_libero.sh +核心环境变量说明 +环境变量 +参数说明 +CHECKPOINT_PATH +训练完成的模型断点目录 +TRAIN_CONFIG_PATH +训练使用的YAML配置文件路径 +TASK_SUITE_NAME +仿真任务套件：libero_spatial / libero_object / libero_goal / libero_10 +ALL_SUITES=1 +开启后批量运行全部4类仿真任务套件 +TASK_INDICES +指定测试任务序号，多任务用逗号分隔（如0,1,2） \ No newline at end of file diff --git a/scripts/README.md b/scripts/README.md index 7a35df2..4cefbc2 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -22,6 +22,47 @@ python scripts/fake_inference.py \ --train-config-path /path/to/config.yml ``` +## Converting LeRobot v3 data for the pinned v2.1 loader + +`convert_lerobot_v3_to_v21.py` creates a separate, episode-sharded v2.0 +dataset compatible with Wall-X's pinned LeRobot v2.1 loader. It converts both +the Parquet data and video streams; the source directory is never changed. + +```bash +python scripts/convert_lerobot_v3_to_v21.py \ + --source /path/to/libero_v3 \ + --output /path/to/libero_v21 \ + --dry-run + +python scripts/convert_lerobot_v3_to_v21.py \ + --source /path/to/libero_v3 \ + --output /path/to/libero_v21 +``` + +Use `--max-episodes 1` only to test the conversion pipeline. It creates a +subset and must not be used for the full training run. + +## Merging compatible LeRobot v2.1 datasets + +`merge_lerobot_v21_datasets.py` combines compatible episode-sharded v2.1 +datasets into one root for a single training run. It renumbers global task, +episode, and frame indices. On one filesystem videos are hard-linked, so the +merge does not duplicate MP4 contents. The sources are never modified. + +```bash +python scripts/merge_lerobot_v21_datasets.py \ + --sources /path/to/libero_spatial /path/to/libero_object /path/to/libero_goal /path/to/libero_10 \ + --output /path/to/libero_all_v21 \ + --dry-run + +python scripts/merge_lerobot_v21_datasets.py \ + --sources /path/to/libero_spatial /path/to/libero_object /path/to/libero_goal /path/to/libero_10 \ + --output /path/to/libero_all_v21 +``` + +Pass `--copy-videos` only when the output is on another filesystem and cannot +use hard links; it duplicates the video storage. + ## LIBERO evaluation `run_libero.sh` is a small shell wrapper around `infer_libero.py`. It requires diff --git a/scripts/run_serving.sh b/scripts/run_serving.sh index 14bfeeb..028ee0e 100755 --- a/scripts/run_serving.sh +++ b/scripts/run_serving.sh @@ -226,6 +226,42 @@ if [[ -z "${CHECKPOINT_PATH}" ]]; then exit 2 fi +# The CLI action horizon overrides the checkpoint config. Reject mismatches +# before loading a large model or handing an incompatible packet to Turtle2. +if [[ -z "${TRAIN_CONFIG_PATH}" && -f "${CHECKPOINT_PATH%/}/config.yml" ]]; then + TRAIN_CONFIG_PATH="${CHECKPOINT_PATH%/}/config.yml" +fi +if [[ -n "${TRAIN_CONFIG_PATH}" ]]; then + if ! "${PYTHON_BIN}" - "${TRAIN_CONFIG_PATH}" "${ACTION_HORIZON}" <<'PY' +import sys +import yaml + +path, requested = sys.argv[1], int(sys.argv[2]) +try: + with open(path, encoding="utf-8") as stream: + config = yaml.load(stream, Loader=yaml.FullLoader) or {} +except (OSError, yaml.YAMLError) as exc: + raise SystemExit(f"error: cannot read training config {path}: {exc}") + +task = config.get("task") or {} +data = config.get("data") or {} +configured = ( + task.get("action_horizon_flow") + or task.get("action_horizon") + or data.get("action_horizon_flow") + or data.get("action_horizon") +) +if configured is not None and requested != int(configured): + raise SystemExit( + f"error: action horizon {requested} disagrees with training config " + f"{path}: {configured}; use --action-horizon {configured}" + ) +PY + then + exit 2 + fi +fi + export CUDA_VISIBLE_DEVICES="${CUDA_ID}" export ENABLE_FAST_PREPROCESS="${ENABLE_FAST_PREPROCESS:-true}" diff --git a/scripts/run_serving_rtc.sh b/scripts/run_serving_rtc.sh new file mode 100755 index 0000000..b0dec9d --- /dev/null +++ b/scripts/run_serving_rtc.sh @@ -0,0 +1,342 @@ +#!/usr/bin/env bash +set -euo pipefail + +usage() { + cat <<'EOF' +Usage: + bash scripts/run_serving_rtc.sh --checkpoint-path /path/to/checkpoint [options] + +Required: + --checkpoint-path PATH Checkpoint directory or checkpoint file. + +Common options: + --train-config-path PATH Training config used by the checkpoint. + --port PORT WebSocket port. Default: 32195. + --host HOST Bind host. Default: 0.0.0.0. + --env X2ROBOT|LIBERO Serving environment. Default: X2ROBOT. + --cuda-id ID Sets CUDA_VISIBLE_DEVICES. Default: 0. + --image-passing-mode MODE base64 or numpy. Default: base64. + --action-horizon N Model action horizon. Default: 32. + --robot-type TYPE desktop, turtle, or ex001. Default: desktop. + --raw-actions Alias for --no-serialize-actions. + --serialize-actions Return robot-serialized actions. + --no-serialize-actions Return raw model action chunks. Default. + --max-batch-size N Enable dynamic batching. + --enable-cuda-graph Enable CUDA graph in the serving runtime. + --enable-experimental-engine Enable the experimental inference engine. + --rtc-execution-horizon N Model steps between RTC replans. Default: 6. + --rtc-max-guidance-weight F RTC guidance clamp. Default: 10.0. + --rtc-prefix-schedule NAME zeros, ones, linear, or exp. Default: linear. + --debug Enable debug logging. + --dry-run Print the command without running it. + +Additional arguments after "--" are forwarded to launch_serving_rtc.py, for example: + bash scripts/run_serving_rtc.sh --checkpoint-path /ckpt -- \ + --model-config.norm-key libero_all + +Environment variables can also be used, e.g. CHECKPOINT_PATH, +TRAIN_CONFIG_PATH, PORT, CUDA_ID, ACTION_HORIZON, WALLX_ENV. +EOF +} + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SOURCE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" + +if [[ -d "${SOURCE_ROOT}/wall_x" ]]; then + export PYTHONPATH="${SOURCE_ROOT}:${PYTHONPATH:-}" +fi + +PYTHON_BIN="${PYTHON_BIN:-python}" +CHECKPOINT_PATH="${CHECKPOINT_PATH:-}" +TRAIN_CONFIG_PATH="${TRAIN_CONFIG_PATH:-}" +PORT="${PORT:-32195}" +HOST="${HOST:-0.0.0.0}" +WALLX_ENV="${WALLX_ENV:-X2ROBOT}" +CUDA_ID="${CUDA_ID:-0}" +IMAGE_PASSING_MODE="${IMAGE_PASSING_MODE:-base64}" +ACTION_HORIZON="${ACTION_HORIZON:-32}" +ROBOT_TYPE="${ROBOT_TYPE:-desktop}" +ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${ROBOT_ACTION_INTERPOLATE_MULTIPLIER:-1}" +ROBOT_ACTION_END_RATIO="${ROBOT_ACTION_END_RATIO:-1.0}" +MODEL_DEVICE="${MODEL_DEVICE:-cuda}" +MAX_BATCH_SIZE="${MAX_BATCH_SIZE:-}" +DEFAULT_PROMPT="${DEFAULT_PROMPT:-}" +SERIALIZE_ACTIONS="${SERIALIZE_ACTIONS:-0}" +ENABLE_CUDA_GRAPH="${ENABLE_CUDA_GRAPH:-0}" +ENABLE_EXPERIMENTAL_ENGINE="${ENABLE_EXPERIMENTAL_ENGINE:-0}" +RTC_EXECUTION_HORIZON="${RTC_EXECUTION_HORIZON:-6}" +RTC_MAX_GUIDANCE_WEIGHT="${RTC_MAX_GUIDANCE_WEIGHT:-10.0}" +RTC_PREFIX_SCHEDULE="${RTC_PREFIX_SCHEDULE:-linear}" +DEBUG="${DEBUG:-0}" +DRY_RUN=0 +EXTRA_ARGS=() + +while [[ $# -gt 0 ]]; do + case "$1" in + -h|--help) + usage + exit 0 + ;; + --checkpoint-path) + CHECKPOINT_PATH="${2:?missing value for --checkpoint-path}" + shift 2 + ;; + --checkpoint-path=*) + CHECKPOINT_PATH="${1#*=}" + shift + ;; + --train-config-path) + TRAIN_CONFIG_PATH="${2:?missing value for --train-config-path}" + shift 2 + ;; + --train-config-path=*) + TRAIN_CONFIG_PATH="${1#*=}" + shift + ;; + --port) + PORT="${2:?missing value for --port}" + shift 2 + ;; + --port=*) + PORT="${1#*=}" + shift + ;; + --host) + HOST="${2:?missing value for --host}" + shift 2 + ;; + --host=*) + HOST="${1#*=}" + shift + ;; + --env) + WALLX_ENV="${2:?missing value for --env}" + shift 2 + ;; + --env=*) + WALLX_ENV="${1#*=}" + shift + ;; + --rtc-execution-horizon) + RTC_EXECUTION_HORIZON="${2:?missing value for --rtc-execution-horizon}" + shift 2 + ;; + --rtc-max-guidance-weight) + RTC_MAX_GUIDANCE_WEIGHT="${2:?missing value for --rtc-max-guidance-weight}" + shift 2 + ;; + --rtc-prefix-schedule) + RTC_PREFIX_SCHEDULE="${2:?missing value for --rtc-prefix-schedule}" + shift 2 + ;; + --cuda-id) + CUDA_ID="${2:?missing value for --cuda-id}" + shift 2 + ;; + --cuda-id=*) + CUDA_ID="${1#*=}" + shift + ;; + --image-passing-mode) + IMAGE_PASSING_MODE="${2:?missing value for --image-passing-mode}" + shift 2 + ;; + --image-passing-mode=*) + IMAGE_PASSING_MODE="${1#*=}" + shift + ;; + --action-horizon) + ACTION_HORIZON="${2:?missing value for --action-horizon}" + shift 2 + ;; + --action-horizon=*) + ACTION_HORIZON="${1#*=}" + shift + ;; + --robot-type) + ROBOT_TYPE="${2:?missing value for --robot-type}" + shift 2 + ;; + --robot-type=*) + ROBOT_TYPE="${1#*=}" + shift + ;; + --robot-action-interpolate-multiplier) + ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${2:?missing value for --robot-action-interpolate-multiplier}" + shift 2 + ;; + --robot-action-interpolate-multiplier=*) + ROBOT_ACTION_INTERPOLATE_MULTIPLIER="${1#*=}" + shift + ;; + --robot-action-end-ratio) + ROBOT_ACTION_END_RATIO="${2:?missing value for --robot-action-end-ratio}" + shift 2 + ;; + --robot-action-end-ratio=*) + ROBOT_ACTION_END_RATIO="${1#*=}" + shift + ;; + --model-device) + MODEL_DEVICE="${2:?missing value for --model-device}" + shift 2 + ;; + --model-device=*) + MODEL_DEVICE="${1#*=}" + shift + ;; + --max-batch-size) + MAX_BATCH_SIZE="${2:?missing value for --max-batch-size}" + shift 2 + ;; + --max-batch-size=*) + MAX_BATCH_SIZE="${1#*=}" + shift + ;; + --default-prompt) + DEFAULT_PROMPT="${2:?missing value for --default-prompt}" + shift 2 + ;; + --default-prompt=*) + DEFAULT_PROMPT="${1#*=}" + shift + ;; + --serialize-actions) + SERIALIZE_ACTIONS=1 + shift + ;; + --no-serialize-actions|--raw-actions) + SERIALIZE_ACTIONS=0 + shift + ;; + --enable-cuda-graph) + ENABLE_CUDA_GRAPH=1 + shift + ;; + --enable-experimental-engine) + ENABLE_EXPERIMENTAL_ENGINE=1 + shift + ;; + --debug) + DEBUG=1 + shift + ;; + --dry-run) + DRY_RUN=1 + shift + ;; + --) + shift + EXTRA_ARGS+=("$@") + break + ;; + *) + EXTRA_ARGS+=("$1") + shift + ;; + esac +done + +if [[ -z "${CHECKPOINT_PATH}" ]]; then + echo "error: --checkpoint-path is required." >&2 + echo >&2 + usage >&2 + exit 2 +fi + +# The CLI action horizon overrides the checkpoint config. Reject mismatches +# before loading a large model or handing an incompatible packet to Turtle2. +if [[ -z "${TRAIN_CONFIG_PATH}" && -f "${CHECKPOINT_PATH%/}/config.yml" ]]; then + TRAIN_CONFIG_PATH="${CHECKPOINT_PATH%/}/config.yml" +fi +if [[ -n "${TRAIN_CONFIG_PATH}" ]]; then + if ! "${PYTHON_BIN}" - "${TRAIN_CONFIG_PATH}" "${ACTION_HORIZON}" <<'PY' +import sys +import yaml + +path, requested = sys.argv[1], int(sys.argv[2]) +try: + with open(path, encoding="utf-8") as stream: + config = yaml.load(stream, Loader=yaml.FullLoader) or {} +except (OSError, yaml.YAMLError) as exc: + raise SystemExit(f"error: cannot read training config {path}: {exc}") + +task = config.get("task") or {} +data = config.get("data") or {} +configured = ( + task.get("action_horizon_flow") + or task.get("action_horizon") + or data.get("action_horizon_flow") + or data.get("action_horizon") +) +if configured is not None and requested != int(configured): + raise SystemExit( + f"error: action horizon {requested} disagrees with training config " + f"{path}: {configured}; use --action-horizon {configured}" + ) +PY + then + exit 2 + fi +fi + +export CUDA_VISIBLE_DEVICES="${CUDA_ID}" +export ENABLE_FAST_PREPROCESS="${ENABLE_FAST_PREPROCESS:-true}" + +CMD=( + "${PYTHON_BIN}" -m wall_x._vendor.harrix.serving.launch_serving_rtc + --env "${WALLX_ENV}" + --host "${HOST}" + --port "${PORT}" + --image-passing-mode "${IMAGE_PASSING_MODE}" + --rtc-execution-horizon "${RTC_EXECUTION_HORIZON}" + --rtc-max-guidance-weight "${RTC_MAX_GUIDANCE_WEIGHT}" + --rtc-prefix-attention-schedule "${RTC_PREFIX_SCHEDULE}" +) + +if [[ "${SERIALIZE_ACTIONS}" == "1" ]]; then + CMD+=(--serialize-actions) +else + CMD+=(--no-serialize-actions) +fi +if [[ -n "${MAX_BATCH_SIZE}" ]]; then + CMD+=(--max-batch-size "${MAX_BATCH_SIZE}") +fi +if [[ -n "${DEFAULT_PROMPT}" ]]; then + CMD+=(--default-prompt "${DEFAULT_PROMPT}") +fi +if [[ "${ENABLE_CUDA_GRAPH}" == "1" ]]; then + CMD+=(--enable-cuda-graph) +fi +if [[ "${ENABLE_EXPERIMENTAL_ENGINE}" == "1" ]]; then + CMD+=(--enable-experimental-engine) +fi +if [[ "${DEBUG}" == "1" ]]; then + CMD+=(--debug) +fi + +CMD+=( + model-config:server-model-config + --model-config.checkpoint-path "${CHECKPOINT_PATH}" + --model-config.action-horizon "${ACTION_HORIZON}" + --model-config.robot-type "${ROBOT_TYPE}" + --model-config.robot-action-interpolate-multiplier "${ROBOT_ACTION_INTERPOLATE_MULTIPLIER}" + --model-config.robot-action-end-ratio "${ROBOT_ACTION_END_RATIO}" + --model-config.model-device "${MODEL_DEVICE}" +) + +if [[ -n "${TRAIN_CONFIG_PATH}" ]]; then + CMD+=(--model-config.train-config-path "${TRAIN_CONFIG_PATH}") +fi + +CMD+=("${EXTRA_ARGS[@]}") + +printf 'Launching Wall-X RTC serving:\n' +printf ' %q' "${CMD[@]}" +printf '\n' + +if [[ "${DRY_RUN}" == "1" ]]; then + exit 0 +fi + +exec "${CMD[@]}" diff --git a/scripts/tcp_ws_bridge.py b/scripts/tcp_ws_bridge.py new file mode 100644 index 0000000..dcccde2 --- /dev/null +++ b/scripts/tcp_ws_bridge.py @@ -0,0 +1,833 @@ +#!/usr/bin/env python3 +"""TCP <-> WebSocket bridge: legacy Quantum-1 robot TCP protocol to +Wall-OSS-0.5 official WebSocket serving. + +Architecture: + Quantum-1 robot (legacy TCP client, e.g. `infer ip port`) + <-> legacy TCP on this bridge (default 30123) + this bridge + <-> Wall-OSS-0.5 WebSocket serving (default ws://127.0.0.1:32195) + +Legacy TCP protocol (from robot_controller.py): + robot -> bridge : [u32 len][state json], then 3x [u32 len][jpeg bytes] + (camera_left, camera_front, camera_right) + bridge -> robot : [u32 len][action json dict] + +Safety: default DRY-RUN. It converts and runs inference but does NOT send +actions back to the robot. Physical sending is gated by BOTH --allow-send +and ACTION_SEMANTICS_CONFIRMED (source-code flag) so a CLI typo alone cannot +enable motion. +""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import json +import logging +import struct +import sys + +import numpy as np +import cv2 +import msgpack +import msgpack_numpy as _m +import websockets +from scipy.spatial.transform import Rotation + +_m.patch() + +# Legacy wire order of the three cameras sent by the robot. +LEGACY_CAM_ORDER = ("camera_left", "camera_front", "camera_right") +# Forward all three legacy camera streams. The serving train config decides +# which of these it consumes (X2Robot commonly uses all three). +SERVE_CAM_KEYS = LEGACY_CAM_ORDER + +# Keep False until the arm action semantics (absolute vs relative) and the +# camera wire order have been confirmed against a real robot capture. +ACTION_SEMANTICS_CONFIRMED = True + +logger = logging.getLogger("bridge") + + +async def _recvall(reader, n): + buf = b"" + while len(buf) < n: + chunk = await reader.read(n - len(buf)) + if not chunk: + return None + buf += chunk + return buf + + +async def _recv_frame(reader): + size_b = await _recvall(reader, 4) + if size_b is None: + return None + size = struct.unpack(" 1: + peak = max(peak, float(np.max(np.abs(np.diff(delta, axis=0))))) + scale = min(1.0, limit / peak) if peak > 0.0 else 1.0 + if scale < 1.0: + limited[:, columns] = current[columns] + delta * scale + scales[name] = (peak, scale) + return limited, scales + + +def assess_right_arm_feedback( + initial_pose, + observed_pose, + commanded_pose, + *, + min_position_delta=0.0005, + min_rotation_delta=0.002, +): + """Compare the next robot state to the preceding right-arm command.""" + initial = np.asarray(initial_pose, dtype=np.float64).reshape(7) + observed = np.asarray(observed_pose, dtype=np.float64).reshape(7) + commanded = np.asarray(commanded_pose, dtype=np.float64).reshape(7) + if not ( + np.isfinite(initial).all() + and np.isfinite(observed).all() + and np.isfinite(commanded).all() + ): + raise ValueError("right-arm feedback poses must be finite 7D vectors") + commanded_delta = commanded - initial + observed_delta = observed - initial + command_requests_motion = ( + np.max(np.abs(commanded_delta[:3])) >= min_position_delta + or np.max(np.abs(commanded_delta[3:6])) >= min_rotation_delta + ) + observed_motion = ( + np.max(np.abs(observed_delta[:3])) >= min_position_delta + or np.max(np.abs(observed_delta[3:6])) >= min_rotation_delta + ) + return { + "commanded_delta": commanded_delta, + "observed_delta": observed_delta, + "command_requests_motion": command_requests_motion, + "observed_motion": observed_motion, + "missing_feedback": command_requests_motion and not observed_motion, + } + + +def prepare_robot_actions( + ws_action, + *, + state_follow1_pos=None, + state_follow2_pos=None, + state_head_pos=None, + state_lift=None, + state_car_pose=None, + action_horizon=32, + action_end_ratio=0.2, + action_interpolate_multiplier=32, + allow_base_motion=False, + hold_left_arm=True, + fixed_car_pose=None, + fixed_lift=None, + allow_constant_fallbacks=False, + max_position_delta=0.10, + max_rotation_delta=0.50, + clip_action_delta=False, + gripper_min=None, + gripper_max=None, +): + """Validate Wall-OSS-0.5 serialized actions for the legacy robot. + + The official Turtle serializer emits head/lift/base trajectories. The + bridge preserves head/lift, but holds the base at the robot-reported pose + by default for safety. Constant fallbacks are available only for dry-run + compatibility tests and must be explicitly requested. + """ + if not isinstance(ws_action, dict): + raise ValueError("WebSocket action response is not a dict") + if "follow1_pos" in ws_action and "follow2_pos" in ws_action: + right = _as_action_chunk(ws_action["follow2_pos"], "follow2_pos") + if hold_left_arm: + if state_follow1_pos is None: + raise ValueError("state_follow1_pos is required when holding the left arm") + current_left = _as_action_chunk( + np.asarray(state_follow1_pos, dtype=np.float64).reshape(1, 7), + "state.follow1_pos", + )[0] + # A right-arm-only checkpoint can still receive a synthetic + # follow1_pos from the generic serializer. Never actuate it. + left = np.repeat(current_left[None, :], right.shape[0], axis=0) + else: + left = _as_action_chunk(ws_action["follow1_pos"], "follow1_pos") + elif "follow2_pos" in ws_action: + # Single-arm serving responses contain the right arm reconstructed by + # Wall-X's official preprocessor. Hold the left arm at its current pose. + right = _as_action_chunk(ws_action["follow2_pos"], "follow2_pos") + if state_follow1_pos is None: + raise ValueError("state_follow1_pos is required for single-arm responses") + current_left = _as_action_chunk( + np.asarray(state_follow1_pos, dtype=np.float64).reshape(1, 7), + "state.follow1_pos", + )[0] + left = np.repeat(current_left[None, :], right.shape[0], axis=0) + else: + raise ValueError( + "Turtle2 requires a serialized follow2_pos trajectory; " + "start serving with --env X2ROBOT --robot-type turtle " + "--serialize-actions" + ) + if left.shape[0] != right.shape[0]: + raise ValueError("follow1_pos/follow2_pos trajectory length mismatch") + serialized_length = int(left.shape[0]) + source_length = serialized_length + # TurtleRobotPreprocessor stacks the observed state before the H model + # actions, then interpolates the result. Strip that one state row before + # applying the legacy infer.py timing policy, so its 20% window starts at + # the first predicted action rather than repeating the current pose. + if action_horizon <= 0: + raise ValueError("action_horizon must be positive") + if source_length == action_horizon + 1: + left = left[1:] + right = right[1:] + source_length -= 1 + elif source_length != action_horizon: + raise ValueError( + f"unexpected serialized action length {source_length}; expected " + f"{action_horizon} or {action_horizon + 1}" + ) + if not (0 < action_end_ratio <= 1): + raise ValueError("action_end_ratio must be in (0, 1]") + selected_length = int(action_end_ratio * source_length) + if selected_length <= 0: + raise ValueError( + f"action trim is empty: end_ratio={action_end_ratio} " + f"length={source_length}" + ) + logger.info( + "executing first %d of %d predicted action steps before interpolation", + selected_length, + source_length, + ) + left = left[:selected_length] + right = right[:selected_length] + if action_interpolate_multiplier < 1: + raise ValueError("action_interpolate_multiplier must be >= 1") + actions = { + "follow1_pos": _legacy_resample( + left, "follow1_pos", 1.0, action_interpolate_multiplier + ), + "follow2_pos": _legacy_resample( + right, "follow2_pos", 1.0, action_interpolate_multiplier + ), + } + for key, width, fallback in ( + ("head_pos", 2, [[0.0, -1.0] for _ in range(source_length)]), + ("lift", 1, [0.4 for _ in range(source_length)]), + ("car_pose", 3, [[0.0, 0.0, 0.0] for _ in range(source_length)]), + ): + values = _validate_optional_series(ws_action, key, serialized_length, width) + if values is None: + state_value = {"head_pos": state_head_pos, "lift": state_lift}.get(key) + if state_value is not None: + values = np.repeat(np.asarray(state_value, dtype=np.float64).reshape(1, -1), serialized_length, axis=0).tolist() + elif key == "car_pose" and state_car_pose is not None: + values = np.repeat(np.asarray(state_car_pose, dtype=np.float64).reshape(1, -1), serialized_length, axis=0).tolist() + elif not allow_constant_fallbacks: + raise ValueError( + f"WebSocket response missing {key!r}; start serving with " + "--env X2ROBOT --robot-type turtle --serialize-actions" + ) + values = fallback + elif len(values) == action_horizon + 1: + values = values[1:] + elif len(values) != action_horizon: + raise ValueError( + f"{key} has unexpected serialized length {len(values)}; expected " + f"{action_horizon} or {action_horizon + 1}" + ) + values = values[:selected_length] + # Preserve the old timing for non-base trajectories as well. + actions[key] = _legacy_resample( + values, key, 1.0, action_interpolate_multiplier + ) + + # The official Turtle serializer represents lift as (T, 1), but the + # legacy Turtle2 receiver assigns each row directly to ``lift_cmd`` and + # expects a scalar. Keep the wire format compatible with that receiver. + actions["lift"] = [float(row[0]) for row in actions["lift"]] + if fixed_lift is not None: + fixed_lift = float(fixed_lift) + if not np.isfinite(fixed_lift) or not 0.0 <= fixed_lift <= 0.47: + raise ValueError("fixed_lift must be within Turtle2 range [0.0, 0.47]") + actions["lift"] = [fixed_lift for _ in actions["lift"]] + + if fixed_car_pose is not None: + pose = np.asarray(fixed_car_pose, dtype=np.float64).reshape(-1) + if pose.shape != (3,) or not np.isfinite(pose).all(): + raise ValueError("fixed_car_pose must have shape (3,) and finite values") + # Turtle2 car_pose is a three-value [x, y, yaw] target. + actions["car_pose"] = [pose.tolist() for _ in range(len(actions["follow1_pos"]))] + elif not allow_base_motion: + # Turtle2 converts pose commands through relative_pose_to_absolute_pose + # before calling set_target_pose(). Sending the reported absolute pose + # again would therefore be interpreted as a relative displacement. + # Zero is the no-motion command; keep the observed pose only in logs. + actions["car_pose"] = [[0.0, 0.0, 0.0] for _ in actions["follow1_pos"]] + # The checkpoint-1 normalizer records the gripper in the robot's 0..4.5 + # units. Invalid negative values can make Turtle2 reject the whole command + # packet, so constrain only this scalar channel before sending. + gripper = np.asarray(actions["follow2_pos"], dtype=np.float64) + if gripper_min is not None and gripper_max is not None: + if not np.isfinite([gripper_min, gripper_max]).all() or gripper_min > gripper_max: + raise ValueError("invalid gripper limits") + before = gripper[:, 6].copy() + gripper[:, 6] = np.clip(gripper[:, 6], float(gripper_min), float(gripper_max)) + if not np.array_equal(before, gripper[:, 6]): + logger.warning( + "clipped gripper command range from [%.5g, %.5g] to [%.5g, %.5g]", + float(np.min(before)), float(np.max(before)), + float(np.min(gripper[:, 6])), float(np.max(gripper[:, 6])), + ) + actions["follow2_pos"] = gripper.tolist() + + if state_follow2_pos is not None: + current = np.asarray(state_follow2_pos, dtype=np.float64).reshape(7) + traj = np.asarray(actions["follow2_pos"], dtype=np.float64) + if clip_action_delta: + # Limit the complete packet relative to the observed pose while + # preserving the model trajectory's shape. Per-sample cumulative + # clipping creates a long synthetic ramp toward an unreachable + # pose and was the source of the repeatable fixed-pose stop. + clipped, scales = _limit_action_packet( + traj, current, max_position_delta, max_rotation_delta + ) + if scales["position"][1] < 1.0 or scales["rotation"][1] < 1.0: + logger.debug( + "scaled right-arm packet: position peak %.5g scale %.5g; " + "rotation peak %.5g scale %.5g", + scales["position"][0], scales["position"][1], + scales["rotation"][0], scales["rotation"][1], + ) + actions["follow2_pos"] = clipped.tolist() + traj = clipped + first_delta = np.abs(traj[0, :3] - current[:3]) + limit_eps = 1e-6 + if np.any(first_delta > max_position_delta + limit_eps): + raise ValueError( + f"right-arm first position delta {first_delta.tolist()} exceeds " + f"limit {max_position_delta} m" + ) + first_rot = np.abs(traj[0, 3:6] - current[3:6]) + if np.any(first_rot > max_rotation_delta + limit_eps): + raise ValueError( + f"right-arm first rotation delta {first_rot.tolist()} exceeds " + f"limit {max_rotation_delta} rad" + ) + if len(traj) > 1: + step_pos = np.max(np.abs(np.diff(traj[:, :3], axis=0)), axis=1) + step_rot = np.max(np.abs(np.diff(traj[:, 3:6], axis=0)), axis=1) + if np.any(step_pos > max_position_delta + limit_eps) or np.any(step_rot > max_rotation_delta + limit_eps): + if clip_action_delta: + # A trajectory can contain Euler-angle wrap discontinuities + # after preprocessing. In clip mode keep the connection + # alive and clamp the offending samples to the prior pose. + logger.warning( + "right-arm trajectory still has over-limit step after clipping; " + "holding offending samples" + ) + safe = traj.copy() + for i in range(1, len(safe)): + dp = safe[i, :3] - safe[i - 1, :3] + dr = safe[i, 3:6] - safe[i - 1, 3:6] + if np.any(np.abs(dp) > max_position_delta) or np.any( + np.abs(dr) > max_rotation_delta + ): + safe[i] = safe[i - 1] + actions["follow2_pos"] = safe.tolist() + else: + raise ValueError("right-arm trajectory contains an over-limit step") + return actions + + +async def _handle_client( + reader, + writer, + ws_url, + instruction, + allow_send, + allow_constant_fallbacks, + action_horizon, + action_end_ratio, + action_interpolate_multiplier, + allow_base_motion, + hold_left_arm, + fixed_car_pose, + fixed_lift, + max_action_cycles, + require_right_feedback, + clip_action_delta, + max_position_delta, + max_rotation_delta, + gripper_min=None, + gripper_max=None, +): + addr = writer.get_extra_info("peername") + logger.info("Robot connected: %s", addr) + + async with websockets.connect(ws_url, max_size=None) as ws: + await ws.recv() # consume server metadata + + # One observation per connection in dry-run; continuous loop when sending. + cycles = 0 + previous_right_plan = None + stop_after_feedback = False + while True: + state = await _recv_state(reader) + if state is None: + logger.info("Robot closed the connection") + return + if previous_right_plan is not None: + feedback = assess_right_arm_feedback( + previous_right_plan["initial"], + state["follow2_pos"], + previous_right_plan["target"], + ) + logger.info( + "right-arm feedback: observed_delta=%s commanded_delta=%s", + feedback["observed_delta"].tolist(), + feedback["commanded_delta"].tolist(), + ) + if feedback["missing_feedback"]: + logger.warning( + "right-arm feedback did not change after the prior packet; " + "bridge sent follow2_pos but the robot did not report EE motion. " + "Check /follow_pos_cmd_2 subscribers, controller enable state, " + "/follow2_pos_back, and /joint_information2 on arm-pc." + ) + if require_right_feedback: + logger.error( + "--require-right-feedback set; " + "stopping before another action packet" + ) + return + if stop_after_feedback: + logger.info("right-arm feedback check complete; stopping after one action packet") + return + previous_right_plan = None + # Turtle2 does not transmit velocity_decomposed. The official + # Turtle preprocessor accepts the field, so provide a neutral + # initial value rather than failing on a missing state key. This + # is deliberately conservative until recurrent velocity semantics + # are confirmed against a real checkpoint/robot run. + state.setdefault("velocity_decomposed", [0.0, 0.0, 0.0]) + + images = [await _recv_image(reader, i) for i in range(len(LEGACY_CAM_ORDER))] + left_img, front_img, right_img = images + logger.info( + "recv state keys=%s images=%sx%s", + list(state.keys()), + front_img.shape[:2], + right_img.shape[:2], + ) + + obs = { + "state": dict(state), + "views": { + SERVE_CAM_KEYS[0]: _img_to_b64(left_img), + SERVE_CAM_KEYS[1]: _img_to_b64(front_img), + SERVE_CAM_KEYS[2]: _img_to_b64(right_img), + }, + "instruction": instruction, + "infer_mode": "flow", + } + logger.info("forwarding to serving: %s", ws_url) + await ws.send(msgpack.packb(obs)) + raw_resp = await ws.recv() + if isinstance(raw_resp, str): + # The official server sends a traceback as a text frame before + # closing when inference fails. Preserve that diagnostic + # instead of masking it with msgpack's bytes-only error. + raise RuntimeError( + "serving returned a text error frame:\n" + raw_resp + ) + resp = msgpack.unpackb(raw_resp) + + actions = prepare_robot_actions( + resp, + state_follow1_pos=state.get("follow1_pos"), + state_follow2_pos=state.get("follow2_pos"), + state_head_pos=state.get("head_pos"), + state_lift=state.get("lift"), + state_car_pose=state.get("car_pose"), + action_horizon=action_horizon, + action_end_ratio=action_end_ratio, + action_interpolate_multiplier=action_interpolate_multiplier, + allow_base_motion=allow_base_motion, + hold_left_arm=hold_left_arm, + fixed_car_pose=fixed_car_pose, + fixed_lift=fixed_lift, + allow_constant_fallbacks=allow_constant_fallbacks, + clip_action_delta=clip_action_delta, + max_position_delta=max_position_delta, + max_rotation_delta=max_rotation_delta, + gripper_min=gripper_min, + gripper_max=gripper_max, + ) + logger.info( + "predicted follow1 T=%d last=%s", len(actions["follow1_pos"]), actions["follow1_pos"][-1] + ) + logger.info( + "predicted follow2 T=%d last=%s", len(actions["follow2_pos"]), actions["follow2_pos"][-1] + ) + try: + current_right = np.asarray(state["follow2_pos"], dtype=np.float64).reshape(7) + first_right = np.asarray(actions["follow2_pos"][0], dtype=np.float64) + last_right = np.asarray(actions["follow2_pos"][-1], dtype=np.float64) + logger.info( + "right-arm current=%s first=%s first_delta=%s last_delta=%s", + current_right.tolist(), first_right.tolist(), + (first_right - current_right).tolist(), + (last_right - current_right).tolist(), + ) + except (KeyError, ValueError): + logger.warning("could not compute right-arm current-to-first delta") + + if not allow_send: + logger.warning("DRY-RUN: action NOT sent to robot") + return + if not ACTION_SEMANTICS_CONFIRMED: + logger.error("ACTION_SEMANTICS_CONFIRMED=False; refusing to send to robot") + return + + logger.info( + "wire action T=%d left[first,last]=%s/%s right[first,last]=%s/%s " + "lift=%s car_pose=%s", + len(actions["follow1_pos"]), + actions["follow1_pos"][0], actions["follow1_pos"][-1], + actions["follow2_pos"][0], actions["follow2_pos"][-1], + actions["lift"][0], actions["car_pose"][0], + ) + payload = json.dumps(actions).encode("utf-8") + writer.write(struct.pack(" 0 and cycles >= max_action_cycles: + if require_right_feedback: + stop_after_feedback = True + logger.info( + "waiting for one robot feedback state before stopping after %d action cycle(s)", + cycles, + ) + continue + logger.warning("stopping after %d action cycle(s)", cycles) + return + + +async def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tcp-host", default="0.0.0.0") + parser.add_argument("--tcp-port", type=int, default=30123) + parser.add_argument("--ws-url", default="ws://127.0.0.1:32195") + parser.add_argument("--instruction", default="pick up the water bottel on the chair") + parser.add_argument( + "--action-horizon", + type=int, + default=32, + help="Model action horizon (Wall-OSS-0.5 default: 32).", + ) + parser.add_argument( + "--allow-send", + action="store_true", + help="Request sending actions back to robot (still gated by ACTION_SEMANTICS_CONFIRMED).", + ) + parser.add_argument( + "--allow-constant-fallbacks", + action="store_true", + help=( + "Use neutral head/lift/base trajectories when the serving response " + "omits them. Intended only for protocol dry-runs." + ), + ) + parser.add_argument( + "--action-end-ratio", + type=float, + default=0.2, + help="Legacy action trim ratio (default: 0.2).", + ) + parser.add_argument( + "--action-interpolate-multiplier", + type=int, + default=32, + help="Legacy linear interpolation multiplier (default: 32).", + ) + parser.add_argument( + "--allow-base-motion", + action="store_true", + help="Forward model-predicted car_pose instead of holding current pose.", + ) + parser.add_argument( + "--allow-left-arm-motion", + action="store_true", + help="Forward serializer-provided left-arm commands. Disabled by default for this right-arm-only checkpoint.", + ) + parser.add_argument( + "--fixed-car-pose", + type=float, + nargs=3, + metavar=("X", "Y", "YAW"), + default=None, + help=( + "Override every output car_pose with [X, Y, YAW]. " + "Use only after confirming Turtle2 relative-target semantics." + ), + ) + parser.add_argument( + "--fixed-lift", + type=float, + default=None, + help="Force lift target for every action frame (meters, range 0.0-0.47).", + ) + parser.add_argument( + "--max-action-cycles", type=int, default=1, + help="Maximum action packets per connection; 0 means unlimited.", + ) + parser.add_argument( + "--require-right-feedback", + action="store_true", + help="Stop before the next packet if the previous right-arm command produced no observed EE motion.", + ) + parser.add_argument( + "--clip-action-delta", action="store_true", + help=( + "Scale each right-arm action packet to safety limits instead of " + "rejecting it; limits apply to packet excursion and inter-sample steps." + ), + ) + parser.add_argument( + "--max-position-delta", type=float, default=0.10, + help="Maximum right-arm position excursion per inference packet (m).", + ) + parser.add_argument( + "--max-rotation-delta", type=float, default=0.50, + help="Maximum right-arm Euler rotation excursion per inference packet (rad).", + ) + parser.add_argument( + "--gripper-min", type=float, default=None, + help="Optional gripper lower bound; disabled by default.", + ) + parser.add_argument( + "--gripper-max", type=float, default=None, + help="Optional gripper upper bound; disabled by default.", + ) + parser.add_argument("--log-level", default="INFO") + args = parser.parse_args() + + logging.basicConfig( + level=getattr(logging, args.log_level.upper()), + format="%(asctime)s %(levelname)s %(name)s %(message)s", + ) + + async def client_connected(reader, writer): + try: + await _handle_client( + reader, + writer, + args.ws_url, + args.instruction, + args.allow_send, + args.allow_constant_fallbacks, + args.action_horizon, + args.action_end_ratio, + args.action_interpolate_multiplier, + args.allow_base_motion, + not args.allow_left_arm_motion, + args.fixed_car_pose, + args.fixed_lift, + args.max_action_cycles, + args.require_right_feedback, + args.clip_action_delta, + args.max_position_delta, + args.max_rotation_delta, + args.gripper_min, + args.gripper_max, + ) + except Exception as exc: + logger.error("handler error: %s: %s", type(exc).__name__, exc) + finally: + try: + writer.close() + await writer.wait_closed() + except Exception: + pass + logger.info("robot session closed") + + server = await asyncio.start_server(client_connected, args.tcp_host, args.tcp_port) + logger.info( + "bridge listening on %s:%s -> %s (allow_send=%s)", + args.tcp_host, args.tcp_port, args.ws_url, args.allow_send, + ) + async with server: + await server.serve_forever() + + +if __name__ == "__main__": + try: + raise SystemExit(asyncio.run(main())) + except KeyboardInterrupt: + sys.exit(0) diff --git a/scripts/tcp_ws_bridge_rtc.py b/scripts/tcp_ws_bridge_rtc.py new file mode 100644 index 0000000..f59a9e9 --- /dev/null +++ b/scripts/tcp_ws_bridge_rtc.py @@ -0,0 +1,872 @@ +#!/usr/bin/env python3 +"""TCP <-> WebSocket bridge: legacy Quantum-1 robot TCP protocol to +Wall-OSS-0.5 official WebSocket serving. + +Architecture: + Quantum-1 robot (legacy TCP client, e.g. `infer ip port`) + <-> legacy TCP on this bridge (default 30123) + this bridge + <-> Wall-OSS-0.5 WebSocket serving (default ws://127.0.0.1:32195) + +Legacy TCP protocol (from robot_controller.py): + robot -> bridge : [u32 len][state json], then 3x [u32 len][jpeg bytes] + (camera_left, camera_front, camera_right) + bridge -> robot : [u32 len][action json dict] + +Safety: default DRY-RUN. It converts and runs inference but does NOT send +actions back to the robot. Physical sending is gated by BOTH --allow-send +and ACTION_SEMANTICS_CONFIRMED (source-code flag) so a CLI typo alone cannot +enable motion. +""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import json +import logging +import struct +import sys + +import numpy as np +import cv2 +import msgpack +import msgpack_numpy as _m +import websockets +from scipy.spatial.transform import Rotation + +_m.patch() + +# Legacy wire order of the three cameras sent by the robot. +LEGACY_CAM_ORDER = ("camera_left", "camera_front", "camera_right") +# Forward all three legacy camera streams. The serving train config decides +# which of these it consumes (X2Robot commonly uses all three). +SERVE_CAM_KEYS = LEGACY_CAM_ORDER + +# Keep False until the arm action semantics (absolute vs relative) and the +# camera wire order have been confirmed against a real robot capture. +ACTION_SEMANTICS_CONFIRMED = True + +logger = logging.getLogger("bridge") + + +async def _recvall(reader, n): + buf = b"" + while len(buf) < n: + chunk = await reader.read(n - len(buf)) + if not chunk: + return None + buf += chunk + return buf + + +async def _recv_frame(reader): + size_b = await _recvall(reader, 4) + if size_b is None: + return None + size = struct.unpack(" 1: + peak = max(peak, float(np.max(np.abs(np.diff(delta, axis=0))))) + scale = min(1.0, limit / peak) if peak > 0.0 else 1.0 + if scale < 1.0: + limited[:, columns] = current[columns] + delta * scale + scales[name] = (peak, scale) + return limited, scales + + +def assess_right_arm_feedback( + initial_pose, + observed_pose, + commanded_pose, + *, + min_position_delta=0.0005, + min_rotation_delta=0.002, +): + """Compare the next robot state to the preceding right-arm command.""" + initial = np.asarray(initial_pose, dtype=np.float64).reshape(7) + observed = np.asarray(observed_pose, dtype=np.float64).reshape(7) + commanded = np.asarray(commanded_pose, dtype=np.float64).reshape(7) + if not ( + np.isfinite(initial).all() + and np.isfinite(observed).all() + and np.isfinite(commanded).all() + ): + raise ValueError("right-arm feedback poses must be finite 7D vectors") + commanded_delta = commanded - initial + observed_delta = observed - initial + command_requests_motion = ( + np.max(np.abs(commanded_delta[:3])) >= min_position_delta + or np.max(np.abs(commanded_delta[3:6])) >= min_rotation_delta + ) + observed_motion = ( + np.max(np.abs(observed_delta[:3])) >= min_position_delta + or np.max(np.abs(observed_delta[3:6])) >= min_rotation_delta + ) + return { + "commanded_delta": commanded_delta, + "observed_delta": observed_delta, + "command_requests_motion": command_requests_motion, + "observed_motion": observed_motion, + "missing_feedback": command_requests_motion and not observed_motion, + } + + +def prepare_robot_actions( + ws_action, + *, + state_follow1_pos=None, + state_follow2_pos=None, + state_head_pos=None, + state_lift=None, + state_car_pose=None, + action_horizon=32, + action_end_ratio=0.2, + action_interpolate_multiplier=32, + allow_base_motion=False, + hold_left_arm=True, + fixed_car_pose=None, + fixed_lift=None, + allow_constant_fallbacks=False, + max_position_delta=0.10, + max_rotation_delta=0.50, + clip_action_delta=False, + gripper_min=None, + gripper_max=None, +): + """Validate Wall-OSS-0.5 serialized actions for the legacy robot. + + The official Turtle serializer emits head/lift/base trajectories. The + bridge preserves head/lift, but holds the base at the robot-reported pose + by default for safety. Constant fallbacks are available only for dry-run + compatibility tests and must be explicitly requested. + """ + if not isinstance(ws_action, dict): + raise ValueError("WebSocket action response is not a dict") + if "follow1_pos" in ws_action and "follow2_pos" in ws_action: + right = _as_action_chunk(ws_action["follow2_pos"], "follow2_pos") + if hold_left_arm: + if state_follow1_pos is None: + raise ValueError("state_follow1_pos is required when holding the left arm") + current_left = _as_action_chunk( + np.asarray(state_follow1_pos, dtype=np.float64).reshape(1, 7), + "state.follow1_pos", + )[0] + # A right-arm-only checkpoint can still receive a synthetic + # follow1_pos from the generic serializer. Never actuate it. + left = np.repeat(current_left[None, :], right.shape[0], axis=0) + else: + left = _as_action_chunk(ws_action["follow1_pos"], "follow1_pos") + elif "follow2_pos" in ws_action: + # Single-arm serving responses contain the right arm reconstructed by + # Wall-X's official preprocessor. Hold the left arm at its current pose. + right = _as_action_chunk(ws_action["follow2_pos"], "follow2_pos") + if state_follow1_pos is None: + raise ValueError("state_follow1_pos is required for single-arm responses") + current_left = _as_action_chunk( + np.asarray(state_follow1_pos, dtype=np.float64).reshape(1, 7), + "state.follow1_pos", + )[0] + left = np.repeat(current_left[None, :], right.shape[0], axis=0) + else: + raise ValueError( + "Turtle2 requires a serialized follow2_pos trajectory; " + "start serving with --env X2ROBOT --robot-type turtle " + "--serialize-actions" + ) + if left.shape[0] != right.shape[0]: + raise ValueError("follow1_pos/follow2_pos trajectory length mismatch") + serialized_length = int(left.shape[0]) + source_length = serialized_length + # TurtleRobotPreprocessor stacks the observed state before the H model + # actions, then interpolates the result. Strip that one state row before + # applying the legacy infer.py timing policy, so its 20% window starts at + # the first predicted action rather than repeating the current pose. + if action_horizon <= 0: + raise ValueError("action_horizon must be positive") + if source_length == action_horizon + 1: + left = left[1:] + right = right[1:] + source_length -= 1 + elif source_length != action_horizon: + raise ValueError( + f"unexpected serialized action length {source_length}; expected " + f"{action_horizon} or {action_horizon + 1}" + ) + if not (0 < action_end_ratio <= 1): + raise ValueError("action_end_ratio must be in (0, 1]") + selected_length = int(action_end_ratio * source_length) + if selected_length <= 0: + raise ValueError( + f"action trim is empty: end_ratio={action_end_ratio} " + f"length={source_length}" + ) + logger.info( + "executing first %d of %d predicted action steps before interpolation", + selected_length, + source_length, + ) + left = left[:selected_length] + right = right[:selected_length] + if action_interpolate_multiplier < 1: + raise ValueError("action_interpolate_multiplier must be >= 1") + actions = { + "follow1_pos": _legacy_resample( + left, "follow1_pos", 1.0, action_interpolate_multiplier + ), + "follow2_pos": _legacy_resample( + right, "follow2_pos", 1.0, action_interpolate_multiplier + ), + } + for key, width, fallback in ( + ("head_pos", 2, [[0.0, -1.0] for _ in range(source_length)]), + ("lift", 1, [0.4 for _ in range(source_length)]), + ("car_pose", 3, [[0.0, 0.0, 0.0] for _ in range(source_length)]), + ): + values = _validate_optional_series(ws_action, key, serialized_length, width) + if values is None: + state_value = {"head_pos": state_head_pos, "lift": state_lift}.get(key) + if state_value is not None: + values = np.repeat(np.asarray(state_value, dtype=np.float64).reshape(1, -1), serialized_length, axis=0).tolist() + elif key == "car_pose" and state_car_pose is not None: + values = np.repeat(np.asarray(state_car_pose, dtype=np.float64).reshape(1, -1), serialized_length, axis=0).tolist() + elif not allow_constant_fallbacks: + raise ValueError( + f"WebSocket response missing {key!r}; start serving with " + "--env X2ROBOT --robot-type turtle --serialize-actions" + ) + values = fallback + elif len(values) == action_horizon + 1: + values = values[1:] + elif len(values) != action_horizon: + raise ValueError( + f"{key} has unexpected serialized length {len(values)}; expected " + f"{action_horizon} or {action_horizon + 1}" + ) + values = values[:selected_length] + # Preserve the old timing for non-base trajectories as well. + actions[key] = _legacy_resample( + values, key, 1.0, action_interpolate_multiplier + ) + + # The official Turtle serializer represents lift as (T, 1), but the + # legacy Turtle2 receiver assigns each row directly to ``lift_cmd`` and + # expects a scalar. Keep the wire format compatible with that receiver. + actions["lift"] = [float(row[0]) for row in actions["lift"]] + if fixed_lift is not None: + fixed_lift = float(fixed_lift) + if not np.isfinite(fixed_lift) or not 0.0 <= fixed_lift <= 0.47: + raise ValueError("fixed_lift must be within Turtle2 range [0.0, 0.47]") + actions["lift"] = [fixed_lift for _ in actions["lift"]] + + if fixed_car_pose is not None: + pose = np.asarray(fixed_car_pose, dtype=np.float64).reshape(-1) + if pose.shape != (3,) or not np.isfinite(pose).all(): + raise ValueError("fixed_car_pose must have shape (3,) and finite values") + # Turtle2 car_pose is a three-value [x, y, yaw] target. + actions["car_pose"] = [pose.tolist() for _ in range(len(actions["follow1_pos"]))] + elif not allow_base_motion: + # Turtle2 converts pose commands through relative_pose_to_absolute_pose + # before calling set_target_pose(). Sending the reported absolute pose + # again would therefore be interpreted as a relative displacement. + # Zero is the no-motion command; keep the observed pose only in logs. + actions["car_pose"] = [[0.0, 0.0, 0.0] for _ in actions["follow1_pos"]] + # The checkpoint-1 normalizer records the gripper in the robot's 0..4.5 + # units. Invalid negative values can make Turtle2 reject the whole command + # packet, so constrain only this scalar channel before sending. + gripper = np.asarray(actions["follow2_pos"], dtype=np.float64) + if gripper_min is not None and gripper_max is not None: + if not np.isfinite([gripper_min, gripper_max]).all() or gripper_min > gripper_max: + raise ValueError("invalid gripper limits") + before = gripper[:, 6].copy() + gripper[:, 6] = np.clip(gripper[:, 6], float(gripper_min), float(gripper_max)) + if not np.array_equal(before, gripper[:, 6]): + logger.warning( + "clipped gripper command range from [%.5g, %.5g] to [%.5g, %.5g]", + float(np.min(before)), float(np.max(before)), + float(np.min(gripper[:, 6])), float(np.max(gripper[:, 6])), + ) + actions["follow2_pos"] = gripper.tolist() + + if state_follow2_pos is not None: + current = np.asarray(state_follow2_pos, dtype=np.float64).reshape(7) + traj = np.asarray(actions["follow2_pos"], dtype=np.float64) + if clip_action_delta: + # Limit the complete packet relative to the observed pose while + # preserving the model trajectory's shape. Per-sample cumulative + # clipping creates a long synthetic ramp toward an unreachable + # pose and was the source of the repeatable fixed-pose stop. + clipped, scales = _limit_action_packet( + traj, current, max_position_delta, max_rotation_delta + ) + if scales["position"][1] < 1.0 or scales["rotation"][1] < 1.0: + logger.debug( + "scaled right-arm packet: position peak %.5g scale %.5g; " + "rotation peak %.5g scale %.5g", + scales["position"][0], scales["position"][1], + scales["rotation"][0], scales["rotation"][1], + ) + actions["follow2_pos"] = clipped.tolist() + traj = clipped + first_delta = np.abs(traj[0, :3] - current[:3]) + limit_eps = 1e-6 + if np.any(first_delta > max_position_delta + limit_eps): + raise ValueError( + f"right-arm first position delta {first_delta.tolist()} exceeds " + f"limit {max_position_delta} m" + ) + first_rot = np.abs(traj[0, 3:6] - current[3:6]) + if np.any(first_rot > max_rotation_delta + limit_eps): + raise ValueError( + f"right-arm first rotation delta {first_rot.tolist()} exceeds " + f"limit {max_rotation_delta} rad" + ) + if len(traj) > 1: + step_pos = np.max(np.abs(np.diff(traj[:, :3], axis=0)), axis=1) + step_rot = np.max(np.abs(np.diff(traj[:, 3:6], axis=0)), axis=1) + if np.any(step_pos > max_position_delta + limit_eps) or np.any(step_rot > max_rotation_delta + limit_eps): + if clip_action_delta: + # A trajectory can contain Euler-angle wrap discontinuities + # after preprocessing. In clip mode keep the connection + # alive and clamp the offending samples to the prior pose. + logger.warning( + "right-arm trajectory still has over-limit step after clipping; " + "holding offending samples" + ) + safe = traj.copy() + for i in range(1, len(safe)): + dp = safe[i, :3] - safe[i - 1, :3] + dr = safe[i, 3:6] - safe[i - 1, 3:6] + if np.any(np.abs(dp) > max_position_delta) or np.any( + np.abs(dr) > max_rotation_delta + ): + safe[i] = safe[i - 1] + actions["follow2_pos"] = safe.tolist() + else: + raise ValueError("right-arm trajectory contains an over-limit step") + return actions + + +async def _handle_client( + reader, + writer, + ws_url, + instruction, + allow_send, + allow_constant_fallbacks, + action_horizon, + action_end_ratio, + action_interpolate_multiplier, + rtc_execution_horizon, + allow_base_motion, + hold_left_arm, + fixed_car_pose, + fixed_lift, + max_action_cycles, + require_right_feedback, + clip_action_delta, + max_position_delta, + max_rotation_delta, + gripper_min=None, + gripper_max=None, +): + addr = writer.get_extra_info("peername") + logger.info("Robot connected: %s", addr) + + async with websockets.connect(ws_url, max_size=None) as ws: + await ws.recv() # consume server metadata + + # One observation per connection in dry-run; continuous loop when sending. + cycles = 0 + previous_right_plan = None + stop_after_feedback = False + while True: + state = await _recv_state(reader) + if state is None: + logger.info("Robot closed the connection") + return + if previous_right_plan is not None: + feedback = assess_right_arm_feedback( + previous_right_plan["initial"], + state["follow2_pos"], + previous_right_plan["target"], + ) + logger.info( + "right-arm feedback: observed_delta=%s commanded_delta=%s", + feedback["observed_delta"].tolist(), + feedback["commanded_delta"].tolist(), + ) + if feedback["missing_feedback"]: + logger.warning( + "right-arm feedback did not change after the prior packet; " + "bridge sent follow2_pos but the robot did not report EE motion. " + "Check /follow_pos_cmd_2 subscribers, controller enable state, " + "/follow2_pos_back, and /joint_information2 on arm-pc." + ) + if require_right_feedback: + logger.error( + "--require-right-feedback set; " + "stopping before another action packet" + ) + return + if stop_after_feedback: + logger.info("right-arm feedback check complete; stopping after one action packet") + return + previous_right_plan = None + # Turtle2 does not transmit velocity_decomposed. The official + # Turtle preprocessor accepts the field, so provide a neutral + # initial value rather than failing on a missing state key. This + # is deliberately conservative until recurrent velocity semantics + # are confirmed against a real checkpoint/robot run. + state.setdefault("velocity_decomposed", [0.0, 0.0, 0.0]) + + images = [await _recv_image(reader, i) for i in range(len(LEGACY_CAM_ORDER))] + left_img, front_img, right_img = images + logger.info( + "recv state keys=%s images=%sx%s", + list(state.keys()), + front_img.shape[:2], + right_img.shape[:2], + ) + + rtc_feedback = dict(state.pop("_rtc", {}) or {}) + request_id = cycles + 1 + obs = { + "state": dict(state), + "views": { + SERVE_CAM_KEYS[0]: _img_to_b64(left_img), + SERVE_CAM_KEYS[1]: _img_to_b64(front_img), + SERVE_CAM_KEYS[2]: _img_to_b64(right_img), + }, + "instruction": instruction, + "infer_mode": "flow", + "rtc": { + "session_id": str(rtc_feedback.get("session_id", addr[0])), + "request_id": request_id, + "consumed_model_steps": int( + rtc_feedback.get("consumed_model_steps", 0) + ), + "inference_delay_steps": int( + rtc_feedback.get("inference_delay_steps", 0) + ), + "execution_horizon": int(rtc_execution_horizon), + "reset": bool(rtc_feedback.get("reset", cycles == 0)), + }, + } + logger.info("forwarding to serving: %s", ws_url) + await ws.send(msgpack.packb(obs)) + raw_resp = await ws.recv() + if isinstance(raw_resp, str): + # The official server sends a traceback as a text frame before + # closing when inference fails. Preserve that diagnostic + # instead of masking it with msgpack's bytes-only error. + raise RuntimeError( + "serving returned a text error frame:\n" + raw_resp + ) + resp = msgpack.unpackb(raw_resp) + response_rtc = dict(resp.get("_rtc", {}) or {}) + if int(response_rtc.get("request_id", request_id)) != request_id: + raise RuntimeError( + "RTC serving returned a stale/mismatched request_id: " + f"expected {request_id}, got {response_rtc.get('request_id')}" + ) + + actions = prepare_robot_actions( + resp, + state_follow1_pos=state.get("follow1_pos"), + state_follow2_pos=state.get("follow2_pos"), + state_head_pos=state.get("head_pos"), + state_lift=state.get("lift"), + state_car_pose=state.get("car_pose"), + action_horizon=action_horizon, + action_end_ratio=action_end_ratio, + action_interpolate_multiplier=action_interpolate_multiplier, + allow_base_motion=allow_base_motion, + hold_left_arm=hold_left_arm, + fixed_car_pose=fixed_car_pose, + fixed_lift=fixed_lift, + allow_constant_fallbacks=allow_constant_fallbacks, + clip_action_delta=clip_action_delta, + max_position_delta=max_position_delta, + max_rotation_delta=max_rotation_delta, + gripper_min=gripper_min, + gripper_max=gripper_max, + ) + logger.info( + "predicted follow1 T=%d last=%s", len(actions["follow1_pos"]), actions["follow1_pos"][-1] + ) + logger.info( + "predicted follow2 T=%d last=%s", len(actions["follow2_pos"]), actions["follow2_pos"][-1] + ) + try: + current_right = np.asarray(state["follow2_pos"], dtype=np.float64).reshape(7) + first_right = np.asarray(actions["follow2_pos"][0], dtype=np.float64) + last_right = np.asarray(actions["follow2_pos"][-1], dtype=np.float64) + logger.info( + "right-arm current=%s first=%s first_delta=%s last_delta=%s", + current_right.tolist(), first_right.tolist(), + (first_right - current_right).tolist(), + (last_right - current_right).tolist(), + ) + except (KeyError, ValueError): + logger.warning("could not compute right-arm current-to-first delta") + + if not allow_send: + logger.warning("DRY-RUN: action NOT sent to robot") + return + if not ACTION_SEMANTICS_CONFIRMED: + logger.error("ACTION_SEMANTICS_CONFIRMED=False; refusing to send to robot") + return + + logger.info( + "wire action T=%d left[first,last]=%s/%s right[first,last]=%s/%s " + "lift=%s car_pose=%s", + len(actions["follow1_pos"]), + actions["follow1_pos"][0], actions["follow1_pos"][-1], + actions["follow2_pos"][0], actions["follow2_pos"][-1], + actions["lift"][0], actions["car_pose"][0], + ) + actions["_rtc"] = { + "session_id": obs["rtc"]["session_id"], + "request_id": request_id, + "guided": bool(resp.get("_rtc", {}).get("guided", False)), + "model_horizon": action_horizon, + "interpolate_multiplier": action_interpolate_multiplier, + "execution_horizon": obs["rtc"]["execution_horizon"], + } + payload = json.dumps(actions).encode("utf-8") + writer.write(struct.pack(" 0 and cycles >= max_action_cycles: + if require_right_feedback: + stop_after_feedback = True + logger.info( + "waiting for one robot feedback state before stopping after %d action cycle(s)", + cycles, + ) + continue + logger.warning("stopping after %d action cycle(s)", cycles) + return + + +async def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tcp-host", default="0.0.0.0") + parser.add_argument("--tcp-port", type=int, default=30123) + parser.add_argument("--ws-url", default="ws://127.0.0.1:32195") + parser.add_argument("--instruction", default="pick up the water bottel on the chair") + parser.add_argument( + "--action-horizon", + type=int, + default=32, + help="Model action horizon (Wall-OSS-0.5 default: 32).", + ) + parser.add_argument( + "--allow-send", + action="store_true", + help="Request sending actions back to robot (still gated by ACTION_SEMANTICS_CONFIRMED).", + ) + parser.add_argument( + "--allow-constant-fallbacks", + action="store_true", + help=( + "Use neutral head/lift/base trajectories when the serving response " + "omits them. Intended only for protocol dry-runs." + ), + ) + parser.add_argument( + "--action-end-ratio", + type=float, + default=1.0, + help="RTC keeps the full model horizon as a fallback queue (default: 1.0).", + ) + parser.add_argument( + "--action-interpolate-multiplier", + type=int, + default=32, + help="Legacy linear interpolation multiplier (default: 32).", + ) + parser.add_argument( + "--rtc-execution-horizon", + type=int, + default=6, + help="Model steps consumed before requesting the next RTC chunk (default: 6).", + ) + parser.add_argument( + "--allow-base-motion", + action="store_true", + help="Forward model-predicted car_pose instead of holding current pose.", + ) + parser.add_argument( + "--allow-left-arm-motion", + action="store_true", + help="Forward serializer-provided left-arm commands. Disabled by default for this right-arm-only checkpoint.", + ) + parser.add_argument( + "--fixed-car-pose", + type=float, + nargs=3, + metavar=("X", "Y", "YAW"), + default=None, + help=( + "Override every output car_pose with [X, Y, YAW]. " + "Use only after confirming Turtle2 relative-target semantics." + ), + ) + parser.add_argument( + "--fixed-lift", + type=float, + default=None, + help="Force lift target for every action frame (meters, range 0.0-0.47).", + ) + parser.add_argument( + "--max-action-cycles", type=int, default=1, + help="Maximum action packets per connection; 0 means unlimited.", + ) + parser.add_argument( + "--require-right-feedback", + action="store_true", + help="Stop before the next packet if the previous right-arm command produced no observed EE motion.", + ) + parser.add_argument( + "--clip-action-delta", action="store_true", + help=( + "Scale each right-arm action packet to safety limits instead of " + "rejecting it; limits apply to packet excursion and inter-sample steps." + ), + ) + parser.add_argument( + "--max-position-delta", type=float, default=0.10, + help="Maximum right-arm position excursion per inference packet (m).", + ) + parser.add_argument( + "--max-rotation-delta", type=float, default=0.50, + help="Maximum right-arm Euler rotation excursion per inference packet (rad).", + ) + parser.add_argument( + "--gripper-min", type=float, default=None, + help="Optional gripper lower bound; disabled by default.", + ) + parser.add_argument( + "--gripper-max", type=float, default=None, + help="Optional gripper upper bound; disabled by default.", + ) + parser.add_argument("--log-level", default="INFO") + args = parser.parse_args() + + if not 1 <= args.rtc_execution_horizon <= args.action_horizon: + parser.error("--rtc-execution-horizon must be between 1 and --action-horizon") + + logging.basicConfig( + level=getattr(logging, args.log_level.upper()), + format="%(asctime)s %(levelname)s %(name)s %(message)s", + ) + + async def client_connected(reader, writer): + try: + await _handle_client( + reader, + writer, + args.ws_url, + args.instruction, + args.allow_send, + args.allow_constant_fallbacks, + args.action_horizon, + args.action_end_ratio, + args.action_interpolate_multiplier, + args.rtc_execution_horizon, + args.allow_base_motion, + not args.allow_left_arm_motion, + args.fixed_car_pose, + args.fixed_lift, + args.max_action_cycles, + args.require_right_feedback, + args.clip_action_delta, + args.max_position_delta, + args.max_rotation_delta, + args.gripper_min, + args.gripper_max, + ) + except Exception as exc: + logger.error("handler error: %s: %s", type(exc).__name__, exc) + finally: + try: + writer.close() + await writer.wait_closed() + except Exception: + pass + logger.info("robot session closed") + + server = await asyncio.start_server(client_connected, args.tcp_host, args.tcp_port) + logger.info( + "bridge listening on %s:%s -> %s (allow_send=%s)", + args.tcp_host, args.tcp_port, args.ws_url, args.allow_send, + ) + async with server: + await server.serve_forever() + + +if __name__ == "__main__": + try: + raise SystemExit(asyncio.run(main())) + except KeyboardInterrupt: + sys.exit(0) diff --git a/setup.py b/setup.py index ca6b4f8..6d50334 100644 --- a/setup.py +++ b/setup.py @@ -9,6 +9,8 @@ CSRC_DIR = OPS_DIR / "csrc" PUBLIC_SCRIPTS = [ "scripts/compute_norm_stats.py", + "scripts/convert_lerobot_v3_to_v21.py", + "scripts/merge_lerobot_v21_datasets.py", "scripts/draw_openloop_plot.py", "scripts/fake_inference.py", "scripts/infer_libero.py", diff --git a/tests/test_lora_scale.py b/tests/test_lora_scale.py new file mode 100644 index 0000000..5e3b67b --- /dev/null +++ b/tests/test_lora_scale.py @@ -0,0 +1,45 @@ +"""Keep serving LoRA merge consistent with exported training parameters.""" + +import json + +import pytest +import torch + +from wall_x._vendor.harrix.utils import ckpt_load + + +def _tiny_lora(): + stem = "model.base_model.model.proj" + return { + stem + ".base_layer.weight": torch.zeros((2, 2)), + stem + ".lora_A.default.weight": torch.eye(2), + stem + ".lora_B.default.weight": torch.eye(2), + } + + +def test_checkpoint_local_lora_config_controls_merge(tmp_path): + (tmp_path / "lora_config.json").write_text( + json.dumps({"lora_r": 2, "lora_alpha": 6}), encoding="utf-8" + ) + tensors = _tiny_lora() + scale = ckpt_load.resolve_lora_scale(str(tmp_path), {}, tensors) + merged = ckpt_load.reshape_compatible_state_dict( + tensors, {"model.proj.weight": torch.zeros((2, 2))}, lora_scale=scale + ) + assert scale == 3 + torch.testing.assert_close(merged["model.proj.weight"], 3 * torch.eye(2)) + + +def test_checkpoint_lora_rank_mismatch_is_rejected(tmp_path): + (tmp_path / "lora_config.json").write_text( + json.dumps({"lora_r": 4, "lora_alpha": 8}), encoding="utf-8" + ) + with pytest.raises(ValueError, match="rank"): + ckpt_load.resolve_lora_scale(str(tmp_path), {}, _tiny_lora()) + + +def test_legacy_checkpoint_without_lora_metadata_uses_previous_scale(tmp_path): + messages = [] + scale = ckpt_load.resolve_lora_scale(str(tmp_path), {}, _tiny_lora(), log_fn=messages.append) + assert scale == 2 + assert any("metadata" in message for message in messages) diff --git a/tests/test_robot_preprocessor.py b/tests/test_robot_preprocessor.py new file mode 100644 index 0000000..1b80d91 --- /dev/null +++ b/tests/test_robot_preprocessor.py @@ -0,0 +1,29 @@ +"""Online robot preprocessing checks shared by Turtle serving.""" + +from types import SimpleNamespace + +import numpy as np + +from wall_x._vendor.harrix.serving._wallx_infer.robot import Robot + + +def test_dof_mask_disables_virtual_action_padding(): + wrapper = SimpleNamespace( + config=SimpleNamespace( + action_horizon=10, + train_config={ + "dof_config": { + "master_right_ee_cartesian_pos": 3, + "master_right_ee_rotation": 3, + "master_right_gripper": 1, + "action_padding": 19, + } + }, + ) + ) + + mask = Robot._get_dof_mask(wrapper) + + assert mask.shape == (1, 10, 26) + np.testing.assert_array_equal(mask[:, :, :7], 1) + np.testing.assert_array_equal(mask[:, :, 7:], 0) diff --git a/tests/test_rtc_wallx.py b/tests/test_rtc_wallx.py new file mode 100644 index 0000000..74fef80 --- /dev/null +++ b/tests/test_rtc_wallx.py @@ -0,0 +1,149 @@ +"""Offline tests for the Wall-X RTC copies.""" + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import torch + +from wall_x._vendor.harrix.serving.rtc_wallx import ( + WallXRTCConfig, + WallXRTCProcessor, +) +from wall_x._vendor.harrix.serving.policy.wall_x_policy_rtc import WallXPolicy +from wall_x._vendor.x2robot_utils import geometry as geom + + +ROOT = Path(__file__).parents[1] + + +def _load_rtc_bridge(): + path = ROOT / "scripts" / "tcp_ws_bridge_rtc.py" + spec = importlib.util.spec_from_file_location("tcp_ws_bridge_rtc", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def test_linear_prefix_weights(): + processor = WallXRTCProcessor( + WallXRTCConfig(execution_horizon=6, prefix_attention_schedule="linear") + ) + weights = processor.get_prefix_weights(start=2, end=6, total=10) + + torch.testing.assert_close(weights[:2], torch.ones(2)) + assert torch.all(weights[2:6] < 1) + assert torch.all(weights[2:6] > 0) + torch.testing.assert_close(weights[6:], torch.zeros(4)) + + +def test_no_prefix_preserves_wallx_velocity(): + processor = WallXRTCProcessor(WallXRTCConfig()) + x = torch.randn(1, 8, 4) + + result = processor.guide_increasing_flow( + x_t=x, + time=torch.tensor(0.4), + predict_velocity=lambda value: torch.ones_like(value), + prev_chunk_left_over=None, + ) + + torch.testing.assert_close(result, torch.ones_like(x)) + + +def test_guidance_changes_velocity_toward_prefix(): + processor = WallXRTCProcessor( + WallXRTCConfig( + execution_horizon=6, + max_guidance_weight=10.0, + prefix_attention_schedule="ones", + ) + ) + x = torch.zeros(1, 8, 3) + prefix = torch.ones(1, 8, 3) + + result = processor.guide_increasing_flow( + x_t=x, + time=torch.tensor(0.5), + predict_velocity=lambda value: value * 0, + prev_chunk_left_over=prefix, + inference_delay=0, + execution_horizon=6, + ) + + assert torch.all(result[:, :6] > 0) + torch.testing.assert_close(result[:, 6:], torch.zeros_like(result[:, 6:])) + + +def test_rtc_bridge_keeps_full_horizon_before_interpolation(): + bridge = _load_rtc_bridge() + horizon = 32 + multiplier = 6 + right = np.zeros((horizon, 7), dtype=np.float64) + right[:, 0] = np.arange(horizon, dtype=np.float64) + + actions = bridge.prepare_robot_actions( + {"follow2_pos": right.tolist()}, + state_follow1_pos=np.zeros(7), + state_follow2_pos=np.zeros(7), + state_head_pos=[0.0, -1.0], + state_lift=[0.3], + state_car_pose=np.zeros(3), + action_horizon=horizon, + action_end_ratio=1.0, + action_interpolate_multiplier=multiplier, + max_position_delta=100.0, + max_rotation_delta=100.0, + ) + + assert len(actions["follow2_pos"]) == horizon * multiplier + assert actions["follow2_pos"][-1][0] == 31.0 + +def test_relative_prefix_is_reanchored_in_normalized_model_layout(): + class IdentityNormalizer: + @staticmethod + def normalize_data(value, dataset_names): + assert dataset_names == ['test'] + return value + + policy = WallXPolicy.__new__(WallXPolicy) + policy.config = SimpleNamespace( + model_device='cpu', + train_config={ + 'dof_config': { + 'action_padding': 16, + 'follow_right_ee_cartesian_pos_relative': 3, + 'follow_right_ee_rotation_6D_relative': 6, + 'follow_right_gripper': 1, + } + }, + ) + policy.model_wrapper = SimpleNamespace( + normalizer_action=IdentityNormalizer(), norm_key='test' + ) + current = np.array([0.4, -0.2, 0.3, 0.1, -0.2, 0.3, 0.5]) + target = np.array([0.45, -0.1, 0.28, 0.1, -0.2, 0.3, 0.7]) + policy._rtc_sessions = { + 'robot': {'follow2_pos': np.stack([current, target]), 'chunk_id': 1} + } + + prefix = policy._rtc_build_normalized_prefix( + {'follow2_pos': current.tolist()}, + {'session_id': 'robot', 'consumed_model_steps': 1}, + ) + + assert prefix.shape == (1, 1, 26) + torch.testing.assert_close(prefix[0, 0, :16], torch.zeros(16)) + torch.testing.assert_close( + prefix[0, 0, 16:19], + torch.tensor(target[:3] - current[:3], dtype=torch.float32), + ) + expected_identity_6d = torch.tensor( + geom.euler_to_matrix_zyx_6d_nb(np.zeros((1, 3)))[0], + dtype=torch.float32, + ) + torch.testing.assert_close(prefix[0, 0, 19:25], expected_identity_6d) + assert np.isclose(prefix[0, 0, 25].item(), target[6]) + diff --git a/tests/test_run_serving_contract.py b/tests/test_run_serving_contract.py new file mode 100644 index 0000000..1a92626 --- /dev/null +++ b/tests/test_run_serving_contract.py @@ -0,0 +1,40 @@ +"""Check that the launch command rejects a mismatched model action length.""" + +import os +from pathlib import Path +import subprocess +import pytest + + +ROOT = Path(__file__).parents[1] +SCRIPT = ROOT / "scripts" / "run_serving.sh" + + +def _launch(horizon): + if "WALLX_TEST_CHECKPOINT" not in os.environ: + pytest.skip("set WALLX_TEST_CHECKPOINT to run checkpoint-specific contract tests") + checkpoint = Path(os.environ["WALLX_TEST_CHECKPOINT"]) + env = os.environ.copy() + env["PYTHON_BIN"] = os.environ.get("WALLX_TEST_PYTHON", "python") + return subprocess.run( + [ + "bash", str(SCRIPT), + "--checkpoint-path", str(checkpoint), + "--train-config-path", str(checkpoint / "config.yml"), + "--action-horizon", str(horizon), + "--dry-run", + ], + cwd=ROOT, env=env, text=True, capture_output=True, check=False, + ) + + +def test_launch_rejects_mismatch_to_training_horizon(): + result = _launch(10) + assert result.returncode != 0 + assert "action horizon" in (result.stderr + result.stdout).lower() + + +def test_launch_accepts_training_horizon(): + result = _launch(32) + assert result.returncode == 0, result.stderr + assert "--model-config.action-horizon 32" in result.stdout diff --git a/tests/test_tcp_ws_bridge.py b/tests/test_tcp_ws_bridge.py new file mode 100644 index 0000000..73f42ea --- /dev/null +++ b/tests/test_tcp_ws_bridge.py @@ -0,0 +1,137 @@ +"""Safety and protocol checks for the Turtle2 action bridge.""" + +import importlib.util +from pathlib import Path + +import numpy as np +import pytest + + +_BRIDGE_PATH = Path(__file__).parents[1] / "scripts" / "tcp_ws_bridge.py" +_SPEC = importlib.util.spec_from_file_location("tcp_ws_bridge", _BRIDGE_PATH) +bridge = importlib.util.module_from_spec(_SPEC) +assert _SPEC.loader is not None +_SPEC.loader.exec_module(bridge) + + + + +def test_packet_limit_scales_chunk_without_cumulative_ramp(): + current = np.zeros(7) + traj = np.array( + [ + [0.10, 0.04, -0.02, 0.20, -0.10, 0.05, 0.0], + [0.20, 0.08, -0.04, 0.40, -0.20, 0.10, 0.1], + [0.40, 0.16, -0.08, 0.80, -0.40, 0.20, 0.2], + ] + ) + + limited, scales = bridge._limit_action_packet(traj, current, 0.005, 0.0125) + + np.testing.assert_allclose(limited[:, :3], traj[:, :3] * 0.005 / 0.40) + np.testing.assert_allclose(limited[:, 3:6], traj[:, 3:6] * 0.0125 / 0.80) + assert np.max(np.abs(limited[:, :3] - current[:3])) <= 0.005 + 1e-12 + assert np.max(np.abs(limited[:, 3:6] - current[3:6])) <= 0.0125 + 1e-12 + assert np.max(np.abs(np.diff(limited[:, :3], axis=0))) <= 0.005 + 1e-12 + assert np.max(np.abs(np.diff(limited[:, 3:6], axis=0))) <= 0.0125 + 1e-12 + assert scales["position"][1] < 1.0 + assert scales["rotation"][1] < 1.0 + + +def test_prepare_robot_actions_keeps_packet_endpoint_small(): + traj = np.array( + [ + [0.10, 0.00, 0.00, 0.20, 0.00, 0.00, 0.0], + [0.20, 0.00, 0.00, 0.40, 0.00, 0.00, 0.0], + [0.30, 0.00, 0.00, 0.60, 0.00, 0.00, 0.0], + [0.40, 0.00, 0.00, 0.80, 0.00, 0.00, 0.0], + ] + ) + state = np.zeros(7) + actions = bridge.prepare_robot_actions( + {"follow1_pos": traj.tolist(), "follow2_pos": traj.tolist()}, + state_follow1_pos=state, + state_follow2_pos=state, + state_head_pos=[0.0, -1.0], + state_lift=[0.4], + state_car_pose=np.zeros(3), + action_horizon=4, + action_end_ratio=1.0, + action_interpolate_multiplier=1, + max_position_delta=0.005, + max_rotation_delta=0.0125, + clip_action_delta=True, + ) + + right = np.asarray(actions["follow2_pos"]) + assert right[-1, 0] <= 0.005 + 1e-12 + assert right[-1, 3] <= 0.0125 + 1e-12 + assert right[-1, 0] < 0.01 # old cumulative clipping reached 4 * 0.005 + + +def test_disabled_base_motion_sends_relative_zero(): + actions = bridge.prepare_robot_actions( + {"follow1_pos": [[0.0] * 7] * 2, "follow2_pos": [[0.0] * 7] * 2}, + state_follow1_pos=[0.0] * 7, state_follow2_pos=[0.0] * 7, + state_head_pos=[0.0, -1.0], state_lift=[0.1], state_car_pose=[1.2, -0.3, 0.7], + action_horizon=2, action_end_ratio=1.0, action_interpolate_multiplier=1, + max_position_delta=0.005, max_rotation_delta=0.0125, clip_action_delta=True, + ) + assert actions["car_pose"] == [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] + + +def test_right_arm_only_mode_holds_serializer_left_arm(): + state_left = np.array([1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.4]) + right = np.array( + [ + [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.7], + [0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.8], + ] + ) + synthetic_left = np.full((2, 7), 99.0) + actions = bridge.prepare_robot_actions( + {"follow1_pos": synthetic_left.tolist(), "follow2_pos": right.tolist()}, + state_follow1_pos=state_left, + state_follow2_pos=np.zeros(7), + state_head_pos=[0.0, -1.0], + state_lift=[0.4], + state_car_pose=np.zeros(3), + action_horizon=2, + action_end_ratio=1.0, + action_interpolate_multiplier=1, + ) + np.testing.assert_allclose(actions["follow1_pos"], np.repeat(state_left[None, :], 2, axis=0)) + np.testing.assert_allclose(actions["follow2_pos"], right) + + +def test_right_arm_feedback_reports_missing_motion(): + initial = np.zeros(7) + command = np.array([0.003, 0.0, 0.0, 0.0, 0.01, 0.0, 0.5]) + missing = bridge.assess_right_arm_feedback(initial, initial, command) + assert missing["command_requests_motion"] + assert missing["missing_feedback"] + + observed = np.array([0.001, 0.0, 0.0, 0.0, 0.003, 0.0, 0.1]) + moved = bridge.assess_right_arm_feedback(initial, observed, command) + assert moved["observed_motion"] + assert not moved["missing_feedback"] + + +def test_raw_action_response_is_rejected_before_robot_packet(): + """The unsupported raw path must not reinterpret virtual padding as pose.""" + raw = np.zeros((2, 26)) + raw[:, :10] = [0.01, 0.02, 0.03, 1, 0, 0, 0, 1, 0, 0.8] + raw[:, 16:26] = [0.20, 0.10, 0.05, 1, 0, 0, 0, 1, 0, 0.2] + + with pytest.raises(ValueError, match="--serialize-actions"): + bridge.prepare_robot_actions( + {"predict_action": raw.tolist()}, + state_follow1_pos=np.zeros(7), + state_follow2_pos=np.zeros(7), + state_head_pos=[0, -1], + state_lift=[0.1], + state_car_pose=np.zeros(3), + action_horizon=2, + action_end_ratio=1, + action_interpolate_multiplier=1, + ) diff --git a/tests/test_turtle2_inference_input.py b/tests/test_turtle2_inference_input.py new file mode 100644 index 0000000..4d3e74a --- /dev/null +++ b/tests/test_turtle2_inference_input.py @@ -0,0 +1,47 @@ +"""Checkpoint-specific multimodal prompt contract, without loading model weights.""" + +import logging +import os +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig +from wall_x._vendor.harrix.serving._wallx_infer.model_wrapper import WallxModelWrapper +from wall_x.trainer.trainer_utils import load_wallx_processors + + +def test_three_square_cameras_preserve_all_image_tokens(): + """Truncating an image placeholder must never reach the VLA as valid input.""" + if "WALLX_TEST_CHECKPOINT" not in os.environ: + pytest.skip("set WALLX_TEST_CHECKPOINT to run checkpoint-specific input tests") + checkpoint = Path(os.environ["WALLX_TEST_CHECKPOINT"]) + config = InferConfig( + checkpoint_path=str(checkpoint), + train_config_path=str(checkpoint / "config.yml"), + model_device="cpu", action_horizon=32, + ) + processor = load_wallx_processors(config.train_config, device="cpu")["processor"] + wrapper = WallxModelWrapper.__new__(WallxModelWrapper) + wrapper.config = config + wrapper.model = SimpleNamespace(processor=processor) + wrapper.cam_names = config.cam_names + wrapper.norm_key = "pick_paper_lerobot_v21_3cam" + wrapper.logger = logging.getLogger(__name__) + wrapper.role_start_symbol = "<|im_start|>" + wrapper.role_end_symbol = "<|im_end|>" + wrapper.vision_start_symbol = "<|vision_start|>" + wrapper.vision_end_symbol = "<|vision_end|>" + wrapper.image_pad_symbol = "<|image_pad|>" + wrapper.propri_symbol = "<|propri|>" + wrapper.action_symbol = "<|action|>" + prefix, postfix = wrapper.get_text_for_action("pick up the paper towel") + image = np.full((448, 448, 3), 127, dtype=np.uint8) + observation = [{name: image for name in wrapper.cam_names}] + inputs = wrapper.construct_model_input(observation, prefix, postfix) + merge = processor.image_processor.merge_size ** 2 + expected = int((inputs["image_grid_thw"].prod(dim=1) // merge).sum()) + actual = int((inputs["input_ids"] == processor.tokenizer.convert_tokens_to_ids("<|image_pad|>")).sum()) + assert actual == expected == 768 diff --git a/wall_x/_vendor/harrix/adapters/qwen_vlact.py b/wall_x/_vendor/harrix/adapters/qwen_vlact.py index a7f2118..583622a 100644 --- a/wall_x/_vendor/harrix/adapters/qwen_vlact.py +++ b/wall_x/_vendor/harrix/adapters/qwen_vlact.py @@ -57,9 +57,9 @@ logger = logging.getLogger(__name__) _PUBLIC_CAMERA_LABELS = { - "face_view": "front view", - "right_wrist_view": "right wrist view", - "left_wrist_view": "left wrist view", + "face_view": "front_view", + "right_wrist_view": "right_wrist_view", + "left_wrist_view": "left_wrist_view", } diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py b/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py index 187ea8a..314da4d 100644 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py @@ -19,10 +19,12 @@ from wall_x._vendor.x2robot_utils.grounding import ( ) from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig -from wall_x._vendor.harrix.utils.ckpt_load import reshape_compatible_state_dict +from wall_x._vendor.harrix.utils.ckpt_load import ( + reshape_compatible_state_dict, + resolve_lora_scale, +) from wall_x._vendor.harrix.utils.train_config import ( resolve_camera_label, - resolve_max_length, resolve_state_bins, resolve_use_state_string_representation, ) @@ -165,8 +167,17 @@ class WallxModelWrapper: "The weights is fused, skipping conversion.", ) + lora_scale = resolve_lora_scale( + checkpoint_path, + self.config.train_config, + state_dict, + log_fn=self.logger.warning, + ) state_dict = reshape_compatible_state_dict( - state_dict, self.model.state_dict(), log_fn=self.logger.info + state_dict, + self.model.state_dict(), + log_fn=self.logger.info, + lora_scale=lora_scale, ) msg = self.model.load_state_dict(state_dict, strict=False) self.model.set_normalizer( @@ -319,9 +330,11 @@ class WallxModelWrapper: images=image_inputs, videos=None, padding=True, - truncation=True, + # Keep every image placeholder during online inference. The + # training text limit can cut them for three large images. + truncation=False, return_tensors="pt", - max_length=resolve_max_length(self.config.train_config), + max_length=None, pad_to_128_multiple=False, pad_prefix_to_same_length=pad_prefix, norm_state=( diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper_rtc.py b/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper_rtc.py new file mode 100644 index 0000000..af03dda --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper_rtc.py @@ -0,0 +1,802 @@ +import os +import torch +import copy +from safetensors.torch import load_file +import numpy as np +from PIL import Image +from qwen_vl_utils.vision_process import smart_resize +from transformers import BatchFeature + +from wall_x.trainer.trainer_utils import load_wallx_processors + +from wall_x._vendor.x2robot_utils.text_templates import ( + preprocesser_call, + get_prologue_with_embodied_information, +) +from wall_x._vendor.x2robot_utils.grounding import ( + reverse_grounding_points, + extract_grounding_points, +) + +from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig +from wall_x._vendor.harrix.utils.ckpt_load import ( + reshape_compatible_state_dict, + resolve_lora_scale, +) +from wall_x._vendor.harrix.utils.train_config import ( + resolve_camera_label, + resolve_state_bins, + resolve_use_state_string_representation, +) +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger +from wall_x.utils.timers import timer, ScopeTimer + + +ENABLE_FAST_PREPROCESS = os.getenv("ENABLE_FAST_PREPROCESS", "False").lower() == "true" + + +def move_to_cuda(obj, device="cuda"): + if isinstance(obj, torch.Tensor): + return obj.to(device) + elif isinstance(obj, (dict, BatchFeature)): + return {k: move_to_cuda(v, device) for k, v in obj.items()} + elif isinstance(obj, list): + return [move_to_cuda(v, device) for v in obj] + elif isinstance(obj, tuple): + return tuple(move_to_cuda(v, device) for v in obj) + else: + return obj + + +class WallxModelWrapper: + def __init__(self, config: InferConfig, rtc_config=None): + from wall_x._vendor.harrix.serving.rtc_wallx import ( + WallXRTCConfig, + WallXRTCProcessor, + ) + + self.config = config + self.rtc_config = rtc_config or WallXRTCConfig() + self.rtc_processor = WallXRTCProcessor(self.rtc_config) + self.logger = InferLogger.get_model_logger("WallxModelWrapper") + self.norm_key = self.config.norm_key + self._register_normalizers() + self.logger.info(f"normalizers {self.norm_key} registered") + self._load_processor() + self._load_model() + self.load_ckpt() + for parameter in self.model.parameters(): + parameter.requires_grad_(False) + self.logger.info( + "model %s loaded; parameters frozen for RTC input-VJP inference", + self.config.checkpoint_path, + ) + + self.norm_key = self.config.norm_key + self._register_normalizers() + self.logger.info(f"normalizers {self.norm_key} registered") + + # Initialize robot_type_id (for v3.1 delta tokenizer) + self.robot_type_id = None + if self.tokenizer_mixin is not None: + robot_type = getattr(self.config, "robot_type", None) + if robot_type: + self.tokenizer_mixin.init_inference(robot_type=robot_type) + if hasattr(self.tokenizer_mixin, "robot_type_id"): + self.robot_type_id = self.tokenizer_mixin.robot_type_id + if self.robot_type_id is not None: + self.logger.info( + f"Initialized robot_type_id: {self.robot_type_id} from robot_type: {robot_type}" + ) + + self.role_start_symbol = "<|im_start|>" + self.role_end_symbol = "<|im_end|>" + self.vision_start_symbol = "<|vision_start|>" + self.vision_end_symbol = "<|vision_end|>" + self.image_pad_symbol = "<|image_pad|>" + self.propri_symbol = "<|propri|>" + self.action_symbol = "<|action|>" + + self.cam_names = self.config.cam_names + + def _camera_name_mapping(self): + return self.config.train_config.get("data", {}).get("camera_name_mapping") + + def _load_processor(self): + # Load tokenizer on model_device for inference + device = getattr(self.config, "model_device", "cuda") + processors_dict = load_wallx_processors(self.config.train_config, device=device) + self.processor = processors_dict["processor"] + self.action_mapper = processors_dict["action_mapper"] + self.tokenizer_mixin = processors_dict.get("tokenizer_mixin") + + def _load_model(self): + from wall_x.trainer.adapters import resolve_adapter + from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl_act_rtc import ( + Qwen2_5_VLMoEForAction, + ) + + model_type = self.config.train_config["model_type"] + if model_type != "qwen2_5": + raise ValueError( + f"RTC serving currently supports model_type='qwen2_5', got {model_type!r}" + ) + adapter_cls = resolve_adapter(model_type) + ModelClass = Qwen2_5_VLMoEForAction + self.ModelClass = ModelClass + + self.logger.info( + "initializing RTC model: %s (%s)", model_type, ModelClass.__name__ + ) + self.model = ModelClass( + self.config.model_config, + self.processor, + self.tokenizer_mixin, + ) + adapter_cls.log_attention_implementation(self.logger, self.model) + self.model_type = model_type + + self.logger.info("resizing model token embeddings") + self.model.resize_token_embeddings(len(self.processor.tokenizer)) + self.logger.info("token embedding resize done") + self.logger.info("casting selected params to bfloat16") + self.model.to_bfloat16_for_selected_params() + self.logger.info("bfloat16 cast done") + + def load_ckpt(self, checkpoint_path: str = None): + if checkpoint_path is None: + checkpoint_path = self.config.checkpoint_path + + if os.path.exists(os.path.join(checkpoint_path, "global_step.pth")): + global_step = torch.load(os.path.join(checkpoint_path, "global_step.pth"))[ + "global_step" + ] + self.logger.info(f"checkpoint global_step: {global_step}") + + fsdp_ckpt = os.path.join(checkpoint_path, "pytorch_model_fsdp.bin") + safetensor_ckpt = os.path.join(checkpoint_path, "model.safetensors") + if os.path.exists(fsdp_ckpt): + self.logger.info(f"Loading FSDP checkpoint: {fsdp_ckpt}") + state_dict = torch.load(fsdp_ckpt, map_location="cpu") + # Unwrap outer dict (e.g. {'state_dict': {...}}) + if isinstance(state_dict, dict) and "state_dict" in state_dict: + self.logger.info( + "Using nested state_dict inside pytorch_model_fsdp.bin" + ) + state_dict = state_dict["state_dict"] + elif os.path.exists(safetensor_ckpt): + self.logger.info(f"loading safetensors checkpoint: {safetensor_ckpt}") + state_dict = load_file(safetensor_ckpt, device="cpu") + else: + raise FileNotFoundError( + f"ERROR: No checkpoint found under {checkpoint_path}. " + "Expecting either pytorch_model_fsdp.bin or model.safetensors." + ) + + # Qwen models need fused weight conversion + if not self.ModelClass.is_fused(state_dict): + self.logger.info( + "Converting non-fused weights to fused format...", + ) + state_dict = self.ModelClass.convert_to_fused(state_dict) + else: + self.logger.info( + "The weights is fused, skipping conversion.", + ) + + lora_scale = resolve_lora_scale( + checkpoint_path, + self.config.train_config, + state_dict, + log_fn=self.logger.warning, + ) + state_dict = reshape_compatible_state_dict( + state_dict, + self.model.state_dict(), + log_fn=self.logger.info, + lora_scale=lora_scale, + ) + msg = self.model.load_state_dict(state_dict, strict=False) + self.model.set_normalizer( + copy.deepcopy(self.normalizer_action), + copy.deepcopy(self.normalizer_propri), + ) + self.logger.info(f"load_state_dict result: {msg}") + self.model.eval() + self.model.to(self.config.model_device) + self.model.to_bfloat16_for_selected_params() + + if hasattr(self.model, "load_optimized_weights"): + self.model.load_optimized_weights(state_dict) + + def _register_normalizers(self): + from wall_x._vendor.harrix.utils.normalizer import build_normalizers + + self.normalizer_action, self.normalizer_propri, resolved = build_normalizers( + self.config.checkpoint_path, + self.config.train_config, + self.config.norm_key, + ) + self.norm_key = resolved + self.config.norm_key = resolved + self._log_norm_debug(self.norm_key) + + def _log_norm_debug(self, norm_key): + if not hasattr(self, "normalizer_action") or not hasattr( + self, "normalizer_propri" + ): + self.logger.warning("[NormDebug] normalizer is not initialized yet") + return + + if ( + norm_key in self.normalizer_action.min + and norm_key in self.normalizer_propri.min + ): + self.logger.debug( + "[NormDebug] norm_key=%s action(min/delta) shape=%s/%s", + norm_key, + tuple(self.normalizer_action.min[norm_key].shape), + tuple(self.normalizer_action.delta[norm_key].shape), + ) + self.logger.debug( + "[NormDebug] action min(all)=%s delta(all)=%s", + self.normalizer_action.min[norm_key].detach().cpu().tolist(), + self.normalizer_action.delta[norm_key].detach().cpu().tolist(), + ) + self.logger.debug( + "[NormDebug] propri(min/delta) shape=%s/%s", + tuple(self.normalizer_propri.min[norm_key].shape), + tuple(self.normalizer_propri.delta[norm_key].shape), + ) + self.logger.debug( + "[NormDebug] propri min(all)=%s delta(all)=%s", + self.normalizer_propri.min[norm_key].detach().cpu().tolist(), + self.normalizer_propri.delta[norm_key].detach().cpu().tolist(), + ) + else: + self.logger.warning( + "[NormDebug] norm_key=%s not found in normalizer", norm_key + ) + + @timer + def construct_model_input( + self, observation, prefix_text, postfix_text, pad_prefix=False + ): + batch_size = len(observation) + dataset_names = [self.norm_key] * batch_size + self.logger.debug("[NormDebug] dataset_names=%s", dataset_names) + + additional_inputs = {} + + # -------- proprioception / masks (batch) -------- + agent_pos_list = [] + agent_pos_mask_list = [] + dof_mask_list = [] + for obs in observation: + if "robot_state_action_data" in obs: + robot_state_action_data = obs["robot_state_action_data"] + + agent_pos = torch.from_numpy(robot_state_action_data.agent_pos) + # Normalize to [1, T, D] for batch cat + if agent_pos.dim() == 2: + agent_pos = agent_pos.unsqueeze(0) + agent_pos_list.append(agent_pos) + + agent_pos_mask = torch.from_numpy( + robot_state_action_data.agent_pos_mask + ) + if agent_pos_mask.dim() == 2: + agent_pos_mask = agent_pos_mask.unsqueeze(0) + agent_pos_mask_list.append(agent_pos_mask) + + dof_mask = torch.from_numpy(robot_state_action_data.dof_mask) + if dof_mask.dim() == 1: + dof_mask = dof_mask.unsqueeze(0) + dof_mask_list.append(dof_mask) + + # cat: [B, T, D] / [B, ...] + if len(agent_pos_list) > 0: + agent_pos = torch.cat(agent_pos_list, dim=0) + agent_pos_mask = torch.cat(agent_pos_mask_list, dim=0) + dof_mask = torch.cat(dof_mask_list, dim=0) + + if self.normalizer_propri is not None: + agent_pos = self.normalizer_propri.normalize_data( + agent_pos, dataset_names + ) + + additional_inputs["proprioception"] = agent_pos.detach() + additional_inputs["agent_pos_mask"] = agent_pos_mask + additional_inputs["dof_mask"] = dof_mask + + # -------- images (flattened, in placeholder scan order) -------- + with ScopeTimer("resize_images"): + # TODO[KC]: optimize this using batch processing + image_inputs = [] + all_image_sizes = [] + if ENABLE_FAST_PREPROCESS: + for obs in observation: + current_image_inputs = self._resize_images_fast(obs) + image_inputs.extend(current_image_inputs) + # tensor shape is (H, W, C), convert to (W, H) to match PIL.size format + all_image_sizes.extend( + [ + (image_i.shape[1], image_i.shape[0]) + for image_i in current_image_inputs + ] + ) + else: + for obs in observation: + current_image_inputs = self._resize_images(obs) + image_inputs.extend(current_image_inputs) + # tensor shape is (H, W, C), convert to (W, H) to match PIL.size format + all_image_sizes.extend( + [ + (image_i.shape[1], image_i.shape[0]) + for image_i in current_image_inputs + ] + ) + + additional_inputs["image_size"] = all_image_sizes + + with ScopeTimer("preprocesser_call"): + inputs = preprocesser_call( + processor=self.model.processor, + prefix_text=prefix_text, + postfix_text=postfix_text, + images=image_inputs, + videos=None, + padding=True, + # Keep every image placeholder during online inference. The + # training text limit can cut them for three large images. + truncation=False, + return_tensors="pt", + max_length=None, + pad_to_128_multiple=False, + pad_prefix_to_same_length=pad_prefix, + norm_state=( + additional_inputs["proprioception"] + if "proprioception" in additional_inputs + and resolve_use_state_string_representation( + self.config.train_config + ) + else None + ), + agent_pos_mask=( + additional_inputs["agent_pos_mask"] + if "agent_pos_mask" in additional_inputs + else None + ), + state_augmentation_prob=0.0, + state_drop_prob=0.0, + state_augmentation_ratio=0.0, + state_bins=resolve_state_bins(self.config.train_config), + inference_mode=True, + ) + + with ScopeTimer("convert_action_token_id and post "): + action_token_id = self.model.processor.tokenizer.convert_tokens_to_ids( + "<|action|>" + ) + flow_action_mask = inputs["input_ids"] == action_token_id + additional_inputs["moe_token_types"] = flow_action_mask + additional_inputs["dataset_names"] = dataset_names + + inputs.update(additional_inputs) + inputs = move_to_cuda(inputs, self.config.model_device) + + return inputs + + def get_text_for_dllm_action(self, instruction): + if ( + self.config.train_config["data"].get("use_embodied_system_prompt_ratio", 0) + > 0 + ): + if self.norm_key != "x2_normal" and self.norm_key != "ex_normal": + self.config.robot_id = 0 + cam_name_mapping = {cam_name: cam_name for cam_name in self.cam_names} + prologue = get_prologue_with_embodied_information( + dataset_name=self.norm_key, + cam_mapping=cam_name_mapping, + robot_id=self.config.robot_id, + uid="", + config=self.config.data_config, + ) + else: + prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n" + user_request = f"{self.role_start_symbol}user\nObservation:" + camera_name_mapping = self._camera_name_mapping() + for cam_name in self.cam_names: + user_request += ( + f" {resolve_camera_label(cam_name, camera_name_mapping)}:" + f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}" + ) + user_request += "\nInstruction:" + text_prompt = f"\nPredict the next action in robot action.\nProprioception: {self.propri_symbol}\n" + user_message = ( + f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n" + ) + placeholder_seq = self.tokenizer_mixin.get_placeholder_for_dllm() + ar_token = "".join(placeholder_seq) + "<|im_end|>\n" + assistant_message = f"{self.role_start_symbol}assistant\n{ar_token}" + flow_action = f"{self.action_symbol * self.config.action_horizon}" + + prefix_text = prologue + user_message + assistant_message + postfix_text = flow_action + + return prefix_text, postfix_text + + @timer + def get_text_for_action(self, instruction): + if ( + self.config.train_config["data"].get("use_embodied_system_prompt_ratio", 0) + > 0 + ): + if self.norm_key not in ["x2_normal", "ex_normal"]: + self.config.robot_id = 0 + cam_name_mapping = {cam_name: cam_name for cam_name in self.cam_names} + prologue = get_prologue_with_embodied_information( + dataset_name=self.norm_key, + cam_mapping=cam_name_mapping, + robot_id=self.config.robot_id, + uid="", + config=self.config.data_config, + ) + else: + prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n" + user_request = f"{self.role_start_symbol}user\nObservation:" + camera_name_mapping = self._camera_name_mapping() + for cam_name in self.cam_names: + user_request += ( + f" {resolve_camera_label(cam_name, camera_name_mapping)}:" + f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}" + ) + user_request += "\nInstruction:" + text_prompt = f"\nPredict the next action in robot action.\nProprioception: {self.propri_symbol}\n" + user_message = ( + f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n" + ) + assistant_message = f"{self.role_start_symbol}assistant\n" + flow_action = f"{self.action_symbol * self.config.action_horizon}" + + prefix_text = prologue + user_message + assistant_message + postfix_text = flow_action + + return prefix_text, postfix_text + + @timer + def get_text_for_subtask_generation(self, instruction): + prologue = f"{self.role_start_symbol}system\nYou are a helpful assistant.{self.role_end_symbol}\n" + user_request = f"{self.role_start_symbol}user\nObservation:" + camera_name_mapping = self._camera_name_mapping() + for cam_name in self.cam_names: + user_request += ( + f" {resolve_camera_label(cam_name, camera_name_mapping)}:" + f" {self.vision_start_symbol}{self.image_pad_symbol}{self.vision_end_symbol}" + ) + user_request += "\nInstruction:" + text_prompt = "\nPredict the next action in language.\n" + user_message = ( + f"{user_request} {instruction}{text_prompt}{self.role_end_symbol}\n" + ) + assistant_message = f"{self.role_start_symbol}assistant\n" + + prefix_text = prologue + user_message + assistant_message + postfix_text = "" + + return prefix_text, postfix_text + + @timer + def _resize_images(self, observation): + image_inputs = [] + for key in self.cam_names: + if key not in observation: + continue + current_obs = observation[key] + if isinstance(current_obs, np.ndarray): + img_pil = Image.fromarray(current_obs) + elif isinstance(current_obs, Image.Image): + img_pil = current_obs + else: + raise ValueError(f"Unsupported image type: {type(current_obs)}") + orig_width, orig_height = img_pil.size + + target_size = self.config.data_config.resolution.get(key, -1) + if target_size != -1: + # Aspect-ratio-preserving resize + if orig_width > orig_height: # landscape + new_width = target_size + new_height = int(target_size * orig_height / orig_width) + else: # portrait + new_height = target_size + new_width = int(target_size * orig_width / orig_height) + img_pil = img_pil.resize((new_width, new_height)) + + # Apply smart resize (Qwen logic) + current_width, current_height = img_pil.size + resized_height, resized_width = smart_resize( + current_height, + current_width, + factor=self.config.data_config.image_factor, + min_pixels=self.config.data_config.min_pixels, + max_pixels=self.config.data_config.max_pixels, + ) + resized_img = img_pil.resize((resized_width, resized_height)) + resized_img = torch.from_numpy(np.array(resized_img)).to( + self.config.model_device + ) + image_inputs.append(resized_img) + + return image_inputs + + def _resize_images_fast(self, observation): + import cv2 + + image_inputs = [] + for key in self.cam_names: + if key not in observation: + continue + current_obs = observation[key] + orig_height, orig_width, _ = current_obs.shape + + target_size = self.config.data_config.resolution.get(key, -1) + current_width, current_height = orig_width, orig_height + if target_size != -1: + # Aspect-ratio-preserving resize + if orig_width > orig_height: # landscape + new_width = target_size + new_height = int(target_size * orig_height / orig_width) + else: # portrait + new_height = target_size + new_width = int(target_size * orig_width / orig_height) + current_width = new_width + current_height = new_height + + # Apply smart resize (Qwen logic) + resized_height, resized_width = smart_resize( + current_height, + current_width, + factor=self.config.data_config.image_factor, # FIXME + min_pixels=self.config.data_config.min_pixels, # FIXME + max_pixels=self.config.data_config.max_pixels, # FIXME + ) + + resized_img = cv2.resize( + current_obs, + (resized_width, resized_height), + interpolation=cv2.INTER_CUBIC, + ) + resized_img = torch.from_numpy(resized_img).to(self.config.model_device) + image_inputs.append(resized_img) + + return image_inputs + + def infer_flow_action( + self, + observation, + instruction, + *, + prev_chunk_left_over=None, + inference_delay=0, + execution_horizon=None, + ): + self.logger.info("generating flow action") + self.logger.info(f"flow action instruction: {instruction}") + + prefix_text, postfix_text = self.get_text_for_action(instruction) + model_input = self.construct_model_input( + [observation], [prefix_text], [postfix_text] + ) + + padding = ( + torch.zeros_like( + self.normalizer_action.delta[model_input["dataset_names"][0]] + ) + .unsqueeze(0) + .to("cpu") + ) + padding_action = self.normalizer_action.normalize_data( + padding, model_input["dataset_names"] + ).to(model_input["input_ids"].device) + + self.logger.info( + "generate_flow_action start (horizon=%s, flow_steps=%s, device=%s)", + self.config.action_horizon, + self.config.num_inference_timesteps, + self.config.model_device, + ) + if torch.cuda.is_available(): + try: + free_b, total_b = torch.cuda.mem_get_info( + torch.device(self.config.model_device) + ) + self.logger.info( + "CUDA mem before flow: free=%.2f GiB / total=%.2f GiB", + free_b / (1024**3), + total_b / (1024**3), + ) + except Exception as e: + self.logger.warning("CUDA mem_get_info failed: %s", e) + + with ScopeTimer("generate_flow_action"): + model_output = self.model.generate_flow_action( + action_horizon=self.config.action_horizon, + action_dim=self.config.action_dim, + num_inference_timesteps=self.config.num_inference_timesteps, + padding_action=padding_action, + rtc_processor=self.rtc_processor, + prev_chunk_left_over=prev_chunk_left_over, + inference_delay=inference_delay, + execution_horizon=execution_horizon, + **model_input, + ) + + self.logger.info("flow action generation done") + + model_output["robot_state_action_data"] = observation["robot_state_action_data"] + model_output["robot_state_action_data"].save_action_data( + model_output["predict_action"] + ) + self.logger.info("saved flow action to robot_state_action_data") + return model_output + + def infer_flow_action_batch(self, observations, instructions): + """ + Batch flow action inference: + - observations: List[Dict], same format as single inference + - instructions: List[str], aligned with observations + Returns List[model_output] of length batch size + """ + assert len(observations) == len( + instructions + ), "observations and instructions must have the same length" + batch_size = len(observations) + + prefix_list = [] + postfix_list = [] + for ins in instructions: + prefix_text, postfix_text = self.get_text_for_action(ins) + prefix_list.append(prefix_text) + postfix_list.append(postfix_text) + + # Build batch input in one preprocesser_call; avoid per-sample cat + batch_inputs = self.construct_model_input( + observations, prefix_list, postfix_list + ) + + padding_list = [] + for ds_name in batch_inputs["dataset_names"]: + padding = ( + torch.zeros_like(self.normalizer_action.delta[ds_name]) + .unsqueeze(0) + .to("cpu") + ) + padding_list.append(padding) + padding = torch.cat(padding_list, dim=0) + padding_action = self.normalizer_action.normalize_data( + padding, batch_inputs["dataset_names"] + ).to(batch_inputs["input_ids"].device) + + with ScopeTimer("generate_flow_action_batch"): + model_output = self.model.generate_flow_action( + action_horizon=self.config.action_horizon, + action_dim=self.config.action_dim, + num_inference_timesteps=self.config.num_inference_timesteps, + padding_action=padding_action, + **batch_inputs, + ) + + predict_action = model_output["predict_action"] # [B, H, D] + + outputs = [] + for i in range(batch_size): + single_action = predict_action[i : i + 1] + single_output = { + "predict_action": single_action, + "robot_state_action_data": observations[i]["robot_state_action_data"], + } + single_output["robot_state_action_data"].save_action_data( + single_output["predict_action"] + ) + outputs.append(single_output) + + return outputs + + def infer_ar_action(self, observation, instruction): + self.logger.info("generating ar action") + self.logger.info(f"ar action instruction: {instruction}") + + prefix_text, _ = self.get_text_for_action(instruction) + model_input = self.construct_model_input([observation], [prefix_text], [""]) + + model_output = self.model.generate_ar_action( + action_horizon=self.config.action_horizon, + action_dim=self.config.ar_action_dim, + num_inference_timesteps=self.config.num_inference_timesteps, + robot_type_id=self.robot_type_id, + **model_input, + ) + self.logger.info("ar action generation done") + + model_output["robot_state_action_data"] = observation["robot_state_action_data"] + model_output["robot_state_action_data"].save_action_data( + model_output["predict_action"] + ) + self.logger.info("saved ar action to robot_state_action_data") + return model_output + + def infer_subtask(self, observation, instruction): + self.logger.info("generating subtask") + self.logger.info(f"subtask instruction: {instruction}") + + prefix_text, postfix_text = self.get_text_for_subtask_generation(instruction) + model_input = self.construct_model_input([observation], [prefix_text], [""]) + + model_output = self.model.generate_text(**model_input) + subtask = model_output["predict_output_text"][0].split("<|im_end|>")[0].strip() + self.logger.info(f"subtask generation done, subtask: {subtask}") + return subtask + + def infer_vqa(self, observation, instruction): + self.logger.info("generating vqa answer") + + if isinstance(observation["multi_modal"], list): + orig_size = observation["multi_modal"][0].size + else: + orig_size = observation["multi_modal"].size + + prefix_text = instruction + model_input = self.construct_model_input([observation], [prefix_text], [""]) + + model_output = self.model.generate_text(**model_input) + answer = model_output["predict_output_text"][0].split("<|im_end|>")[0].strip() + self.logger.info(f"vqa answer done, answer: {answer}") + answer = reverse_grounding_points( + answer, + orig_size[1], + orig_size[0], + model_input["image_size"][0][1], + model_input["image_size"][0][0], + self.config.data_config.model_type, + ) + points = extract_grounding_points(answer) + + return { + "answer": answer, + "points": points, + } + + def infer_dllm_action( + self, observation, instruction, use_ar_action=False, dataset_name="x2_normal" + ): + assert self.model_type in ["qwen2_5"], "DLLM only supports qwen2_5" + prefix_text, postfix_text = self.get_text_for_dllm_action(instruction) + model_input = self.construct_model_input( + [observation], [prefix_text], [postfix_text], pad_prefix=True + ) + model_input = self.model.update_infer_dllm_position_mask(model_input) + total_ar_step = self.tokenizer_mixin.inference_ar_steps_for_dllm + cnt = 0 + while cnt < 3: # fast mode: at most 3 retries + model_output = self.model.generate_dllm_action( + action_horizon=self.config.action_horizon, + action_dim=self.config.action_dim, + ar_action_dim=self.config.ar_action_dim, + num_inference_timesteps=self.config.num_inference_timesteps, + use_ar_action=use_ar_action, + total_ar_step=total_ar_step, + robot_type_id=self.robot_type_id, # for v3.1 delta decode + **model_input, + ) + if model_output["predict_action"] is not None: + break + cnt += 1 + self.logger.warning(f"dllm action generation failed, retry {cnt}") + self.logger.info("dllm action generation done") + + model_output["robot_state_action_data"] = observation["robot_state_action_data"] + model_output["robot_state_action_data"].save_action_data( + model_output["predict_action"] + ) + self.logger.info("saved dllm action to robot_state_action_data") + return model_output diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py b/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py index 641bbc4..1541525 100644 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py @@ -312,9 +312,16 @@ class Robot(ABC): return state, views def _get_dof_mask(self): - dof_config = self.config.train_config["dof_config"] + from wall_x._vendor.harrix.utils.train_config import resolve_dof_config + + dof_config = resolve_dof_config(self.config.train_config) total_dof = sum(dof_config.values()) dof_mask = np.ones((1, self.config.action_horizon, total_dof)) + start_idx = 0 + for key, dof_size in dof_config.items(): + if key == "action_padding": + dof_mask[:, :, start_idx : start_idx + dof_size] = 0 + start_idx += dof_size return dof_mask @staticmethod diff --git a/wall_x/_vendor/harrix/serving/launch_serving_rtc.py b/wall_x/_vendor/harrix/serving/launch_serving_rtc.py new file mode 100644 index 0000000..913c711 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/launch_serving_rtc.py @@ -0,0 +1,319 @@ +#!/usr/bin/env python3 +""" +Server script for Wall-X model. + +This script serves a Wall-X model using a websocket server, allowing +clients to connect and get action predictions from observations. + +""" + +import dataclasses +import enum +import inspect +import logging +import os +import socket +import sys +import yaml +import traceback +import tyro + +from wall_x._vendor.harrix.serving.websocket_policy_server import WebsocketPolicyServer + + +def _server_model_config_to_infer_kwargs(model_config) -> dict: + from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig + + infer_params = inspect.signature(InferConfig.__init__).parameters + return { + k: v + for k, v in vars(model_config).items() + if v is not None and k in infer_params + } + + +def get_wallx_policy( + model_config, image_passing_mode, serialize_actions=True, rtc_config=None +): + from wall_x._vendor.harrix.serving.policy.wall_x_policy_rtc import WallXPolicy + from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig + + config = InferConfig(**_server_model_config_to_infer_kwargs(model_config)) + return WallXPolicy( + config=config, + image_passing_mode=image_passing_mode, + serialize_actions=serialize_actions, + rtc_config=rtc_config, + ) + + + + +logger = logging.getLogger(__name__) + + +class EnvMode(enum.Enum): + """Supported environments/datasets.""" + + X2ROBOT = "x2robot" + LIBERO = "libero" + + +@dataclasses.dataclass +class ServerModelConfig: + """Configuration for loading a Wall-X model.""" + + checkpoint_path: str | None = None + train_config_path: str | None = None + # robot_host: str = '0.0.0.0' + # robot_port: int = 33723 + robot_type: str = "desktop" # ["desktop", "turtle"] + robot_action_start_ratio: float = ( + 0.0 # proportion of action execution to start from + ) + robot_action_end_ratio: float = 1.0 # proportion of action execution to end at + robot_action_interpolate_multiplier: int = 10 # action interpolation multiplier + robot_use_joint_angle_control: bool = ( + False # use joint angle control (model must predict joints) + ) + turtle_as_desktop: bool = ( + False # use turtle platform as desktop with fixed base/head/camera/height + ) + action_horizon: int = 32 # specify the correct horizon for the model + action_dim: int | None = None + model_device: str = "cuda" + num_inference_timesteps: int = 10 + num_inference_steps: int | None = None + cfg_scale: float | None = None + seed: int | None = None + save_video_dir: str = "./videos" + + # Please specify explicitly if the checkpoint was not trained on the x2robot dataset. + norm_key: str | None = None + # Model cameras; None = infer from train config ``data.key_mappings.camera``. + cam_names: list[str] | None = None + # Robot camera keys in incoming websocket observations. + camera_front_key: str = "camera_front" + camera_left_key: str = "camera_left" + camera_right_key: str = "camera_right" + # Serving prompt controls. If the client request has no instruction, + # default_instruction is used. prompt_template follows train-config semantics. + default_instruction: str | None = None + prompt_template: str | None = None + qwen25_prompt_template: str | None = None + prompt_priority_order: str | None = None + +@dataclasses.dataclass +class Args: + """Arguments for the serve_wall_x script.""" + + # Environment mode (used for default configurations) + env: EnvMode = EnvMode.X2ROBOT + + # Model configuration. If not provided, uses default config for the environment + model_config: ServerModelConfig | None = None + + # Default text prompt to use if not provided in observation + default_prompt: str | None = None + + # Port to serve the policy on + port: int = 43007 + + # Host to bind the server to + host: str = "0.0.0.0" + + # Enable debug logging + debug: bool = False + + # Image passing mode + image_passing_mode: str = "base64" # ["numpy", "base64"] + + # Model type + model_type: str = "wallx" # OSS supports qwen2.5 Wall-X only + + # Serialize actions via robot_preprocessor (True for robot control, False for raw output) + serialize_actions: bool = True + + # -- Dynamic batching ----------------------------------------- + # Set max_batch_size to enable dynamic batching. None = single mode. + max_batch_size: int | None = None + max_wait_time_ms: float = 0 + max_queue_size: int = 100 + timeout_ms: float = 30000 + + # -- Real-Time Chunking --------------------------------------- + rtc_execution_horizon: int = 6 + rtc_max_guidance_weight: float = 10.0 + rtc_prefix_attention_schedule: str = "linear" + + # -- Engine flags --------------------------------------------- + enable_experimental_engine: bool = False + enable_cuda_graph: bool = False + + +# Default model configurations for each environment +DEFAULT_CONFIGS: dict[EnvMode, ServerModelConfig] = { + EnvMode.X2ROBOT: ServerModelConfig( + checkpoint_path=None, + train_config_path=None, + robot_action_start_ratio=0.0, + robot_action_end_ratio=1.0, + robot_action_interpolate_multiplier=10, + robot_use_joint_angle_control=False, + turtle_as_desktop=False, + action_horizon=32, + action_dim=None, + model_device="cuda", + num_inference_timesteps=10, + ), + EnvMode.LIBERO: ServerModelConfig( + checkpoint_path=None, + train_config_path=None, + robot_type="desktop", + robot_action_start_ratio=0.0, + robot_action_end_ratio=1.0, + robot_action_interpolate_multiplier=1, + action_horizon=10, + action_dim=None, + model_device="cuda", + num_inference_timesteps=10, + cam_names=None, + ), +} + + +def get_model_config(args: Args) -> ServerModelConfig: + """Get model configuration from args or defaults.""" + if args.model_config is not None: + return args.model_config + + if config := DEFAULT_CONFIGS.get(args.env): + logger.info(f"Using default configuration for {args.env.value}") + return config + + raise ValueError( + f"No default configuration for {args.env.value}. " + f"Please provide --model-config with model_path and action_tokenizer_path." + ) + + +def create_policy(args: Args): + """Create a policy from the given arguments.""" + model_config = get_model_config(args) + if args.model_type != "wallx": + raise ValueError( + f"Unsupported model type: {args.model_type!r}. " + "The public package only supports model_type='wallx'." + ) + from wall_x._vendor.harrix.serving.rtc_wallx import WallXRTCConfig + + if args.max_batch_size is not None: + raise ValueError( + "RTC serving keeps per-session chunk state and does not support " + "dynamic batching; omit --max-batch-size." + ) + rtc_config = WallXRTCConfig( + execution_horizon=args.rtc_execution_horizon, + max_guidance_weight=args.rtc_max_guidance_weight, + prefix_attention_schedule=args.rtc_prefix_attention_schedule, + ) + policy = get_wallx_policy( + model_config, + args.image_passing_mode, + args.serialize_actions, + rtc_config=rtc_config, + ) + return policy + + +def main(args: Args) -> None: + """Main function to start the Wall-X model server.""" + log_level = logging.DEBUG if args.debug else logging.INFO + logging.basicConfig( + level=log_level, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) + + # Set engine environment variables + if args.enable_experimental_engine: + os.environ["ENABLE_EXPERIMENTAL_INFERENCE_ENGINE"] = "true" + logger.info("ENABLE_EXPERIMENTAL_INFERENCE_ENGINE=true") + + if args.enable_cuda_graph: + os.environ["ENABLE_CUDA_GRAPH"] = "true" + logger.info("ENABLE_CUDA_GRAPH=true") + + logger.info("Starting model server") + logger.info(f"Model type: {args.model_type}") + logger.info(f"Environment: {args.env.value}") + logger.info(f"Port: {args.port}") + logger.info(f"Host: {args.host}") + logger.info(f"Serialize actions: {args.serialize_actions}") + logger.info( + "RTC: execution_horizon=%d max_guidance_weight=%.3f schedule=%s", + args.rtc_execution_horizon, + args.rtc_max_guidance_weight, + args.rtc_prefix_attention_schedule, + ) + + # Create policy + try: + policy = create_policy(args) + except Exception as e: + logger.error(f"Failed to create policy: {e}") + logger.error(traceback.format_exc()) + sys.exit(1) + + # Get policy metadata + policy_metadata = policy.metadata + policy_metadata["env"] = args.env.value + policy_metadata["rtc"] = { + "enabled": True, + "execution_horizon": args.rtc_execution_horizon, + "max_guidance_weight": args.rtc_max_guidance_weight, + "prefix_attention_schedule": args.rtc_prefix_attention_schedule, + } + + # Get network info + hostname = socket.gethostname() + try: + local_ip = socket.gethostbyname(hostname) + except Exception: + local_ip = "unknown" + + logger.info(f"Server hostname: {hostname}") + logger.info(f"Server IP: {local_ip}") + logger.info(f"Server will be available at: ws://{args.host}:{args.port}") + logger.info(f"Health check endpoint: http://{args.host}:{args.port}/healthz") + + batching_str = ( + f"batch_size={args.max_batch_size}, wait={args.max_wait_time_ms}ms" + if args.max_batch_size + else "disabled" + ) + logger.info(f"Batching: {batching_str}") + + # Create and start server + server = WebsocketPolicyServer( + policy=policy, + host=args.host, + port=args.port, + metadata=policy_metadata, + max_batch_size=args.max_batch_size, + max_wait_time_ms=args.max_wait_time_ms, + max_queue_size=args.max_queue_size, + timeout_ms=args.timeout_ms, + ) + + logger.info("Starting server...") + try: + server.serve_forever() + except KeyboardInterrupt: + logger.info("Server stopped by user") + except Exception as e: + logger.error(f"Server error: {e}") + sys.exit(1) + + +if __name__ == "__main__": + main(tyro.cli(Args)) diff --git a/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py b/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py index 155d602..2306616 100644 --- a/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py +++ b/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py @@ -128,6 +128,18 @@ class WallXPolicy(BasePolicy): response: Dict[str, Any] = { "predict_action": np.asarray(predict_action, dtype=np.float32), } + # Single-arm checkpoints normally return raw model output because the + # generic serializer expects both arms. Still expose the official + # right-arm reconstruction (absolute xyz/euler/gripper) so legacy + # clients do not have to guess relative-action semantics. + state_action = model_output.get("robot_state_action_data") + if state_action is not None and self._is_single_arm_right_only(): + try: + response["follow2_pos"] = self.robot_preprocessor._get_right_arm_action( + state_action + ).astype(np.float32) + except Exception as exc: + logger.warning("Could not serialize single-arm right action: %s", exc) if "subtask" in model_output: response["subtask"] = model_output["subtask"] return response diff --git a/wall_x/_vendor/harrix/serving/policy/wall_x_policy_rtc.py b/wall_x/_vendor/harrix/serving/policy/wall_x_policy_rtc.py new file mode 100644 index 0000000..6d21e4c --- /dev/null +++ b/wall_x/_vendor/harrix/serving/policy/wall_x_policy_rtc.py @@ -0,0 +1,472 @@ +import base64 +import logging +import threading +from typing import Dict, Any, List + +import cv2 +import numpy as np +import torch +from wall_x._vendor.harrix.serving.websocket_policy_server import BasePolicy +from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig +from wall_x._vendor.harrix.serving._wallx_infer.model_wrapper_rtc import WallxModelWrapper +from wall_x._vendor.harrix.serving._wallx_infer.robot import ( + DesktopRobotPreprocessor, + TurtleRobotPreprocessor, + EX001RobotPreprocessor, +) +from wall_x._vendor.harrix.serving.policy._smoothing import apply_smoothing +from wall_x.utils.timers import ScopeTimer +from wall_x._vendor.harrix.serving.rtc_wallx import WallXRTCConfig +from wall_x._vendor.x2robot_utils import geometry as geom + +logger = logging.getLogger(__name__) + +# Inference modes aligned with model_wrapper for vqa / batch extensions +INFER_MODE_FLOW = "flow" +INFER_MODE_AR = "ar" +INFER_MODE_FLOW_WITH_SUBTASK = "flow_with_subtask" +INFER_MODE_DLLM_FLOW = "dllm" +INFER_MODE_DLLM_DD = "discrete_diffusion" +ACTION_INFER_MODES = ( + INFER_MODE_FLOW, + INFER_MODE_AR, + INFER_MODE_FLOW_WITH_SUBTASK, + INFER_MODE_DLLM_FLOW, + INFER_MODE_DLLM_DD, +) + + +class WallXPolicy(BasePolicy): + """Policy wrapper for Wall-X model that implements the BasePolicy interface.""" + + def __init__( + self, + config: InferConfig, + image_passing_mode: str = "base64", + default_infer_mode: str = INFER_MODE_FLOW, + serialize_actions: bool = True, + rtc_config: WallXRTCConfig | None = None, + ): + """Initialize the Wall-X policy. + + Args: + config: Inference configuration dataclass. + image_passing_mode: How images are passed from client ('base64' or 'numpy'). + default_infer_mode: Default inference mode ('flow', 'ar', 'flow_with_subtask', etc.). + serialize_actions: If True, actions are serialized via robot_preprocessor; + if False, raw model output is returned directly. + """ + self.config = config + self.rtc_config = rtc_config or WallXRTCConfig() + self.model_wrapper = WallxModelWrapper(config, rtc_config=self.rtc_config) + self.robot_preprocessor = self._register_robot_preprocessor() + self._rtc_sessions: dict[str, dict[str, Any]] = {} + self._rtc_infer_lock = threading.Lock() + self.image_passing_mode = image_passing_mode + self.default_infer_mode = default_infer_mode + self.serialize_actions = serialize_actions + logger.info( + "Image passing mode: %s, robot_type: %s, default_infer_mode: %s, " + "serialize_actions: %s, smooth_action: %s, smooth_gripper: %s", + self.image_passing_mode, + config.robot_type, + self.default_infer_mode, + self.serialize_actions, + config.smooth_action, + config.smooth_gripper, + ) + + def _register_robot_preprocessor(self): + """Select preprocessor by config.robot_type; mirrors env._register_robot.""" + if self.config.robot_type == "desktop": + return DesktopRobotPreprocessor(self.config) + if self.config.robot_type == "turtle": + return TurtleRobotPreprocessor(self.config) + if self.config.robot_type == "ex001": + return EX001RobotPreprocessor(self.config) + raise ValueError(f"Invalid robot_type: {self.config.robot_type!r}") + + @staticmethod + def _decode_base64_image(image_b64: str) -> np.ndarray: + """Decode client base64 JPEG/PNG payloads to RGB images.""" + img_bytes = base64.b64decode(image_b64) + img_array = np.frombuffer(img_bytes, dtype=np.uint8) + decoded_img = cv2.imdecode(img_array, cv2.IMREAD_COLOR) + if decoded_img is None: + raise ValueError("Failed to decode base64 image payload") + return cv2.cvtColor(decoded_img, cv2.COLOR_BGR2RGB) + + def reset(self) -> None: + """Reset the policy state.""" + self.action_buffer = [] + self.buffer_index = 0 + logger.debug("Policy reset") + + def _get_dof_config(self) -> dict: + train_config = self.config.train_config or {} + return ( + train_config.get("dof_config") + or train_config.get("task", {}).get("dof_config", {}) + or {} + ) + + def _is_single_arm_right_only(self) -> bool: + """True for single-arm LIBERO-style configs (right arm only, no left arm).""" + train_config = self.config.train_config or {} + agent_pos = ( + train_config.get("agent_pos_config") + or train_config.get("task", {}).get("agent_pos_config") + or {} + ) + skip = {"action_padding"} + has_left = any(k.startswith("follow_left_") for k in agent_pos if k not in skip) + has_right = any( + k.startswith("follow_right_") for k in agent_pos if k not in skip + ) + return has_right and not has_left + + def _pack_action_chunk_response( + self, model_output: Dict[str, Any] + ) -> Dict[str, Any]: + """Return a msgpack-safe action chunk for websocket clients.""" + predict_action = model_output["predict_action"] + if isinstance(predict_action, torch.Tensor): + predict_action = predict_action.detach().cpu().numpy() + response: Dict[str, Any] = { + "predict_action": np.asarray(predict_action, dtype=np.float32), + } + # Single-arm checkpoints normally return raw model output because the + # generic serializer expects both arms. Still expose the official + # right-arm reconstruction (absolute xyz/euler/gripper) so legacy + # clients do not have to guess relative-action semantics. + state_action = model_output.get("robot_state_action_data") + if state_action is not None and self._is_single_arm_right_only(): + try: + response["follow2_pos"] = self.robot_preprocessor._get_right_arm_action( + state_action + ).astype(np.float32) + except Exception as exc: + logger.warning("Could not serialize single-arm right action: %s", exc) + if "subtask" in model_output: + response["subtask"] = model_output["subtask"] + return response + + def _get_predict_action_keys(self) -> List[str]: + """Resolve flat action key order for smoothing / serialization.""" + data_cfg = self.config.data_config + keys = None + if isinstance(data_cfg, dict): + keys = data_cfg.get("predict_action_keys") + else: + keys = getattr(data_cfg, "predict_action_keys", None) + if keys: + return list(keys) + return list(self._get_dof_config().keys()) + + def _apply_smoothing(self, model_output: Dict[str, Any]) -> None: + """Apply Laplacian smoothing on predict_action before 6D->euler conversion.""" + cfg = self.config + if not getattr(cfg, "smooth_action", False): + return + dof_config = self._get_dof_config() + apply_smoothing( + model_output, + smooth_action=True, + smooth_gripper=cfg.smooth_gripper, + predict_action_keys=self._get_predict_action_keys(), + action_padding_dof=dof_config.get("action_padding"), + action_dim=cfg.action_dim, + ) + + def _rtc_build_normalized_prefix(self, state: Dict, rtc: Dict): + session_id = str(rtc.get("session_id", "default")) + if rtc.get("reset"): + self._rtc_sessions.pop(session_id, None) + return None + cached = self._rtc_sessions.get(session_id) + if not cached: + return None + + consumed = max(0, int(rtc.get("consumed_model_steps", 0))) + absolute = np.asarray(cached["follow2_pos"], dtype=np.float64) + if consumed >= len(absolute): + return None + absolute = absolute[consumed:] + current = np.asarray(state["follow2_pos"], dtype=np.float64).reshape(7) + + state_rotation = geom.euler_to_matrix_zyx_batch_nb( + current[None, 3:6] + )[0] + absolute_rotation = geom.euler_to_matrix_zyx_batch_nb(absolute[:, 3:6]) + delta_rotation = absolute_rotation @ state_rotation.T + delta_rotation_6d = delta_rotation[:, :2, :].reshape(len(absolute), 6) + + dof_config = self._get_dof_config() + columns = [] + for key, width in dof_config.items(): + width = int(width) + lowered = key.lower() + if key == "action_padding": + values = np.zeros((len(absolute), width), dtype=np.float64) + elif "follow_right" in lowered and "cartesian_pos" in lowered: + values = absolute[:, :3] - current[None, :3] + elif "follow_right" in lowered and "rotation_6d" in lowered: + values = delta_rotation_6d + elif "follow_right" in lowered and "rotation" in lowered: + raise ValueError("RTC currently requires the checkpoint's 6D rotation layout") + elif "follow_right" in lowered and "gripper" in lowered: + values = absolute[:, 6:7] + else: + values = np.zeros((len(absolute), width), dtype=np.float64) + if values.shape[1] != width: + raise ValueError( + f"RTC field {key!r} expected width {width}, got {values.shape[1]}" + ) + columns.append(values) + + raw = np.concatenate(columns, axis=1) + tensor = torch.from_numpy(raw).to( + device=self.config.model_device, dtype=torch.float32 + ).unsqueeze(0) + return self.model_wrapper.normalizer_action.normalize_data( + tensor, [self.model_wrapper.norm_key] + ) + + def _rtc_update_session( + self, + *, + state: Dict, + rtc: Dict, + model_output: Dict[str, Any], + ) -> None: + session_id = str(rtc.get("session_id", "default")) + state_action = model_output.get("robot_state_action_data") + if state_action is None: + return + absolute = self.robot_preprocessor._get_right_arm_action(state_action) + absolute = np.asarray(absolute, dtype=np.float64) + if len(absolute) == self.config.action_horizon + 1: + absolute = absolute[1:] + self._rtc_sessions[session_id] = { + "follow2_pos": absolute.copy(), + "chunk_id": int(rtc.get("request_id", 0)), + } + + def _run_action_infer( + self, + observation: Dict, + instruction: str, + mode: str, + *, + rtc_context: Dict | None = None, + rtc_prefix=None, + ) -> Dict[str, Any]: + """Run model_wrapper action inference by mode; returns model_output with robot_state_action_data. + + Supports flow | ar | flow_with_subtask. VQA can be added here later (different return format). + """ + if mode == INFER_MODE_FLOW: + rtc_context = rtc_context or {} + return self.model_wrapper.infer_flow_action( + observation, + instruction, + prev_chunk_left_over=rtc_prefix, + inference_delay=int(rtc_context.get("inference_delay_steps", 0)), + execution_horizon=int( + rtc_context.get( + "execution_horizon", + self.rtc_config.execution_horizon, + ) + ), + ) + if mode == INFER_MODE_AR: + return self.model_wrapper.infer_ar_action(observation, instruction) + if mode == INFER_MODE_FLOW_WITH_SUBTASK: + with ScopeTimer("infer_subtask"): + subtask = self.model_wrapper.infer_subtask(observation, instruction) + with ScopeTimer("infer_flow_action"): + model_output = self.model_wrapper.infer_flow_action( + observation, subtask + ) + model_output["subtask"] = subtask + return model_output + if mode == INFER_MODE_DLLM_FLOW: + return self.model_wrapper.infer_dllm_action( + observation, instruction, use_ar_action=False + ) + if mode == INFER_MODE_DLLM_DD: + return self.model_wrapper.infer_dllm_action( + observation, instruction, use_ar_action=True + ) + raise ValueError( + f"Unsupported infer_mode={mode!r}, expected one of {ACTION_INFER_MODES}" + ) + + def infer(self, obs: Dict) -> Dict: + with self._rtc_infer_lock: + return self._infer_locked(obs) + + def _infer_locked(self, obs: Dict) -> Dict: + """Infer action from observation. + + Args: + obs: Dictionary containing: + - 'state': Robot state + - 'views': Camera views (keyed by camera name) + - 'instruction': Task instruction + - Optional: 'infer_mode' - one of 'flow' | 'ar' | 'flow_with_subtask' + - Optional: 'robot_action_start_ratio' / 'robot_action_end_ratio' / + 'robot_action_interpolate_multiplier' - override config for action trim/interpolate + + Returns: + When serialize_actions=True (default): Serialized action dict + (e.g. follow1_pos/follow2_pos or follow1_joints/follow2_joints). + When serialize_actions=False: Raw model_output dict. + If infer_mode is flow_with_subtask, also includes 'subtask'. + """ + state = obs["state"] + views = obs["views"] + instruction = obs.get("instruction") or self.config.default_instruction or "" + if self.image_passing_mode == "base64": + for k, v in views.items(): + views[k] = np.expand_dims(self._decode_base64_image(v), axis=0) + + with ScopeTimer("get_observation"): + observation = self.robot_preprocessor.get_observation(state, views) + + mode = obs.get("infer_mode", self.default_infer_mode) + rtc_context = dict(obs.get("rtc") or {}) + rtc_prefix = self._rtc_build_normalized_prefix(state, rtc_context) + with ScopeTimer(f"infer_{mode}"): + model_output = self._run_action_infer( + observation, + instruction, + mode, + rtc_context=rtc_context, + rtc_prefix=rtc_prefix, + ) + + self._apply_smoothing(model_output) + if ( + "robot_state_action_data" in model_output + and "predict_action" in model_output + ): + model_output["robot_state_action_data"].save_action_data( + model_output["predict_action"] + ) + + self._rtc_update_session( + state=state, + rtc=rtc_context, + model_output=model_output, + ) + + # Single-arm LIBERO has no left-arm EE to serialize; dual-arm follow1/follow2 + # layout would fail in get_serialized_actions. Return raw action chunks instead. + if not self.serialize_actions or self._is_single_arm_right_only(): + response = self._pack_action_chunk_response(model_output) + else: + response = self.robot_preprocessor.get_serialized_actions( + model_output, robot_action_interpolate_multiplier=1 + ) + response["_rtc"] = { + "session_id": str(rtc_context.get("session_id", "default")), + "request_id": int(rtc_context.get("request_id", 0)), + "guided": rtc_prefix is not None, + "inference_delay_steps": int( + rtc_context.get("inference_delay_steps", 0) + ), + "execution_horizon": int( + rtc_context.get( + "execution_horizon", + self.rtc_config.execution_horizon, + ) + ), + } + return response + + # -- Batch inference ---------------------------------------------- + + def _preprocess_obs(self, obs: Dict): + """Preprocess a single observation dict into (observation, instruction). + + Handles both base64 and raw image modes. + """ + state = obs["state"] + views = obs["views"] + instruction = obs.get("instruction") or self.config.default_instruction or "" + + if self.image_passing_mode == "base64": + for k, v in views.items(): + views[k] = np.expand_dims(self._decode_base64_image(v), axis=0) + else: + # raw numpy: ensure (1, H, W, C) + for k, img in views.items(): + if isinstance(img, np.ndarray) and img.ndim == 3: + views[k] = np.expand_dims(img, axis=0) + + observation = self.robot_preprocessor.get_observation(state, views) + return observation, instruction + + def infer_batch( + self, obs_list: List[Dict[str, Any]], skip_serialize: bool = False + ) -> List[Dict[str, Any]]: + """Perform batch inference. + + Args: + obs_list: List of observations, each containing: + - "views": dict of camera images + - "state": robot state dict or array + - "instruction": text instruction + + Returns: + List of action dicts. + """ + batch_size = len(obs_list) + logger.info(f"WallXPolicy.infer_batch: processing {batch_size} observations") + + try: + observations = [] + instructions = [] + for obs in obs_list: + observation, instruction = self._preprocess_obs(obs) + observations.append(observation) + instructions.append(instruction) + + mode = obs_list[0].get("infer_mode", self.default_infer_mode) + + with torch.no_grad(): + if mode == INFER_MODE_FLOW: + model_outputs = self.model_wrapper.infer_flow_action_batch( + observations, instructions + ) + else: + # AR / flow_with_subtask: fall back to per-sample inference + model_outputs = [] + for obs_dict, instruction in zip(observations, instructions): + output = self._run_action_infer(obs_dict, instruction, mode) + model_outputs.append(output) + + if skip_serialize: + return [{}] * len(model_outputs) + + results = [] + for model_output in model_outputs: + self._apply_smoothing(model_output) + if self.serialize_actions and not self._is_single_arm_right_only(): + action = self.robot_preprocessor.get_serialized_actions( + model_output, robot_action_interpolate_multiplier=1 + ) + results.append(action) + else: + results.append(self._pack_action_chunk_response(model_output)) + + return results + + except Exception as e: + logger.error(f"Batch inference failed: {e}", exc_info=True) + return [{"action": {}, "error": str(e)} for _ in obs_list] + + @property + def metadata(self) -> Dict[str, Any]: + return {"batch_enabled": True, "model": "wall-x"} diff --git a/wall_x/_vendor/harrix/serving/rtc_wallx.py b/wall_x/_vendor/harrix/serving/rtc_wallx.py new file mode 100644 index 0000000..32550d7 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/rtc_wallx.py @@ -0,0 +1,119 @@ +"""Wall-X Real-Time Chunking helpers. + +The guided update follows LeRobot's RTC implementation, adapted to Wall-X's +flow convention: Wall-X integrates normalized actions from noise at t=0 to a +clean action at t=1 with velocity ``dx/dt``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import math + +import torch + + +@dataclass(frozen=True) +class WallXRTCConfig: + enabled: bool = True + execution_horizon: int = 6 + max_guidance_weight: float = 10.0 + prefix_attention_schedule: str = "linear" + + def __post_init__(self): + if self.execution_horizon <= 0: + raise ValueError("execution_horizon must be positive") + if self.max_guidance_weight <= 0: + raise ValueError("max_guidance_weight must be positive") + if self.prefix_attention_schedule not in {"zeros", "ones", "linear", "exp"}: + raise ValueError( + "prefix_attention_schedule must be zeros, ones, linear, or exp" + ) + + +class WallXRTCProcessor: + """Inference-time RTC guidance in normalized Wall-X action space.""" + + def __init__(self, config: WallXRTCConfig): + self.config = config + + def get_prefix_weights(self, start: int, end: int, total: int) -> torch.Tensor: + start = max(0, min(int(start), int(end), int(total))) + end = max(start, min(int(end), int(total))) + schedule = self.config.prefix_attention_schedule + if schedule == "zeros": + weights = torch.zeros(total) + weights[:start] = 1.0 + return weights + if schedule == "ones": + weights = torch.zeros(total) + weights[:end] = 1.0 + return weights + + middle_len = end - start + if middle_len: + middle = torch.linspace(1.0, 0.0, middle_len + 2)[1:-1] + if schedule == "exp": + middle = middle * torch.expm1(middle) / (math.e - 1.0) + else: + middle = torch.empty(0) + return torch.cat( + [torch.ones(start), middle, torch.zeros(total - end)], dim=0 + ) + + def guide_increasing_flow( + self, + *, + x_t: torch.Tensor, + time: torch.Tensor | float, + predict_velocity, + prev_chunk_left_over: torch.Tensor | None, + inference_delay: int = 0, + execution_horizon: int | None = None, + ) -> torch.Tensor: + """Return RTC-guided velocity for a flow integrated from t=0 to t=1.""" + if prev_chunk_left_over is None or not self.config.enabled: + return predict_velocity(x_t) + + x = x_t.detach().clone().requires_grad_(True) + prefix = prev_chunk_left_over.to(device=x.device, dtype=x.dtype) + if prefix.ndim == 2: + prefix = prefix.unsqueeze(0) + if prefix.shape[0] == 1 and x.shape[0] > 1: + prefix = prefix.expand(x.shape[0], -1, -1) + + padded = torch.zeros_like(x) + steps = min(prefix.shape[1], x.shape[1]) + dims = min(prefix.shape[2], x.shape[2]) + padded[:, :steps, :dims] = prefix[:, :steps, :dims] + + horizon = execution_horizon or self.config.execution_horizon + horizon = min(int(horizon), steps, x.shape[1]) + weights = self.get_prefix_weights(inference_delay, horizon, x.shape[1]) + weights = weights.to(device=x.device, dtype=x.dtype).view(1, -1, 1) + + with torch.enable_grad(): + velocity = predict_velocity(x) + t = torch.as_tensor(time, device=x.device, dtype=x.dtype) + remaining = torch.clamp(1.0 - t, min=1e-6) + clean_estimate = x + remaining * velocity + error = (padded - clean_estimate) * weights + correction = torch.autograd.grad( + clean_estimate, + x, + grad_outputs=error.detach(), + retain_graph=False, + )[0] + + # Same guidance schedule as LeRobot RTC after mapping its 1->0 time + # convention to Wall-X's 0->1 convention. + t_safe = torch.clamp(t, min=1e-6) + inv_r2 = (remaining.square() + t.square()) / remaining.square() + weight = (remaining / t_safe) * inv_r2 + weight = torch.nan_to_num( + weight, + nan=self.config.max_guidance_weight, + posinf=self.config.max_guidance_weight, + ).clamp(max=self.config.max_guidance_weight) + return (velocity + weight * correction).detach() + diff --git a/wall_x/_vendor/harrix/utils/ckpt_load.py b/wall_x/_vendor/harrix/utils/ckpt_load.py index 4e00350..a127c23 100644 --- a/wall_x/_vendor/harrix/utils/ckpt_load.py +++ b/wall_x/_vendor/harrix/utils/ckpt_load.py @@ -7,6 +7,8 @@ are intentionally left to the adapter. from __future__ import annotations import os +import json +import math from typing import Callable, Optional import torch @@ -17,6 +19,58 @@ def _noop_log(_msg: str, **_kw) -> None: pass +def resolve_lora_scale( + checkpoint_path: str, + train_config: dict, + state_dict: dict, + log_fn: Optional[Callable] = None, +) -> float: + """Read PEFT alpha/r from an exported config and validate it against weights.""" + log_fn = log_fn or _noop_log + ranks = { + int(tensor.shape[0]) + for name, tensor in state_dict.items() + if name.endswith(".lora_A.default.weight") + } + if not ranks: + return 1.0 + if len(ranks) != 1: + raise ValueError(f"LoRA weights contain multiple ranks: {sorted(ranks)}") + rank = ranks.pop() + + candidates = [ + os.path.join(checkpoint_path, "lora_config.json"), + os.path.join(checkpoint_path, "adapter_config.json"), + (train_config.get("model") or {}).get("lora_config_path"), + ] + source = next((path for path in candidates if path and os.path.isfile(path)), None) + if source is None: + log_fn( + "LoRA scale metadata is missing; using legacy alpha/r=2.0. " + "Export the training LoRA JSON as checkpoint/lora_config.json " + "before deploying a checkpoint with changed LoRA settings." + ) + return 2.0 + + with open(source, encoding="utf-8") as stream: + config = json.load(stream) + configured_rank = int(config.get("lora_r", config.get("r", rank))) + if configured_rank != rank: + raise ValueError( + f"LoRA rank mismatch: {source} has {configured_rank}, weights have {rank}" + ) + if config.get("rank_pattern") or config.get("alpha_pattern"): + raise ValueError(f"Per-layer LoRA scaling in {source} is unsupported") + alpha = config.get("lora_alpha") + if alpha is None: + raise ValueError(f"LoRA config {source} lacks lora_alpha") + scale = float(alpha) / (math.sqrt(rank) if config.get("use_rslora") else rank) + if not math.isfinite(scale) or scale <= 0: + raise ValueError(f"Invalid LoRA scale {scale} from {source}") + log_fn(f"LoRA scale={scale:g} from {source} (alpha={alpha}, rank={rank})") + return scale + + def _align_checkpoint_tensor( param: torch.Tensor, target: torch.Tensor, @@ -64,22 +118,58 @@ def _align_checkpoint_tensor( def reshape_compatible_state_dict( - state_dict: dict, model_sd: dict, log_fn: Optional[Callable] = None + state_dict: dict, + model_sd: dict, + log_fn: Optional[Callable] = None, + lora_scale: float = 2.0, ) -> dict: """Align checkpoint tensors to the target model shapes via crop / pad.""" log_fn = log_fn or _noop_log + # Training checkpoints keep PEFT modules (base_layer/lora_A/lora_B), + # while inference uses the plain fused projection tensors. Materialize + # those tensors here so serving does not silently drop the LoRA update. + merged = dict(state_dict) + for name, base in list(state_dict.items()): + suffix = ".base_layer.weight" + if not name.endswith(suffix): + continue + stem = name[: -len(suffix)] + a_name = stem + ".lora_A.default.weight" + b_name = stem + ".lora_B.default.weight" + if a_name not in state_dict or b_name not in state_dict: + continue + a = state_dict[a_name].to(dtype=torch.float32) + b = state_dict[b_name].to(dtype=torch.float32) + merged[name.replace(".base_layer.weight", ".weight")] = ( + base.to(dtype=torch.float32) + lora_scale * (b @ a) + ).to(dtype=base.dtype) + del merged[name] + del merged[a_name] + del merged[b_name] + log_fn(f"Merged LoRA weights for {stem}") + state_dict = merged out = {} for name, param in state_dict.items(): - if name not in model_sd: + # Training with PEFT wraps the backbone under ``base_model.model``; + # serving instantiates the unwrapped model. Normalize that prefix so + # fine-tuned backbone weights are applied instead of being ignored. + target_name = name + if ".base_layer." in target_name: + target_name = target_name.replace(".base_layer.", ".") + if target_name not in model_sd and target_name.startswith( + "model.base_model.model." + ): + target_name = "model." + target_name[len("model.base_model.model.") :] + if target_name not in model_sd: log_fn(f"Not used parameter: {name}") continue - target = model_sd[name] + target = model_sd[target_name] if param.shape == target.shape: - out[name] = param + out[target_name] = param continue - aligned = _align_checkpoint_tensor(param, target, name, log_fn) + aligned = _align_checkpoint_tensor(param, target, target_name, log_fn) if aligned is not None: - out[name] = aligned + out[target_name] = aligned return out diff --git a/wall_x/_vendor/harrix/utils/normalizer.py b/wall_x/_vendor/harrix/utils/normalizer.py index 2494b83..bb2eaad 100644 --- a/wall_x/_vendor/harrix/utils/normalizer.py +++ b/wall_x/_vendor/harrix/utils/normalizer.py @@ -1,8 +1,10 @@ """Build action/proprio normalizers and resolve the effective norm key. Public inference artifacts must carry their own normalization data. This module -uses checkpoint-local ``norm_stats.json`` first, then checkpoint-side normalizer -state dicts, and finally an explicit ``customized_action_statistic_dof`` path. +uses checkpoint-local ``norm_stats.json`` for the 7D Euler layout, and uses +checkpoint-side normalizer state dicts for checkpoints whose train config contains +6D rotation fields. It finally supports an explicit +``customized_action_statistic_dof`` path. It does not fall back to internal default action statistics. """ @@ -59,6 +61,19 @@ def _missing_normalizer_error(checkpoint_path: str, train_config: dict) -> FileN ) +def _uses_rotation_6d(train_config: dict) -> bool: + """Return whether the checkpoint was trained with 6D rotation fields.""" + for layout_name in ("dof_config", "agent_pos_config", "ar_dof_config"): + layout = train_config.get(layout_name) + if layout is None and isinstance(train_config.get("task"), dict): + layout = train_config["task"].get(layout_name) + if isinstance(layout, dict) and any( + "rotation_6d" in str(key).lower() for key in layout + ): + return True + return False + + def build_normalizers( checkpoint_path: str, train_config: dict, @@ -66,14 +81,24 @@ def build_normalizers( ) -> tuple[Normalizer, Normalizer, str]: """Return action/proprio normalizers and the resolved norm key.""" norm_stats_path = os.path.join(checkpoint_path, "norm_stats.json") - if os.path.exists(norm_stats_path): + action_pth = os.path.join(checkpoint_path, "normalizer_action.pth") + propri_pth = os.path.join(checkpoint_path, "normalizer_propri.pth") + prefer_checkpoint_pth = ( + os.path.exists(action_pth) + and os.path.exists(propri_pth) + and _uses_rotation_6d(train_config) + ) + if os.path.exists(norm_stats_path) and not prefer_checkpoint_pth: propri_stats = _load_norm_stats(norm_stats_path, "observation.state") action_stats = _load_norm_stats(norm_stats_path, "action") normalizer_propri = Normalizer.from_lerobot_norm_stats(propri_stats, norm_key) normalizer_action = Normalizer.from_lerobot_norm_stats(action_stats, norm_key) else: - action_pth = os.path.join(checkpoint_path, "normalizer_action.pth") - propri_pth = os.path.join(checkpoint_path, "normalizer_propri.pth") + if prefer_checkpoint_pth: + logger.info( + "Using checkpoint normalizer .pth files because train config " + "contains rotation_6D; ignoring 7D norm_stats.json" + ) custom_stats = _load_custom_action_stats(train_config) if custom_stats is None and (not os.path.exists(action_pth) or not os.path.exists(propri_pth)): raise _missing_normalizer_error(checkpoint_path, train_config) diff --git a/wall_x/config/model_config.py b/wall_x/config/model_config.py index 4edaf44..3a34ad6 100644 --- a/wall_x/config/model_config.py +++ b/wall_x/config/model_config.py @@ -26,6 +26,7 @@ class QActModelConfig(ModelConfig): """QAct model config for the Qwen2.5 VLA path.""" config_path: str = "" + lora_config_path: Optional[str] = None processor_path: str = "" pretrained_path: Optional[str] = None backbone: str = "qwen2_5" diff --git a/wall_x/data/backends/lerobot/loader.py b/wall_x/data/backends/lerobot/loader.py index c000ec8..2af5043 100644 --- a/wall_x/data/backends/lerobot/loader.py +++ b/wall_x/data/backends/lerobot/loader.py @@ -9,6 +9,7 @@ from typing import Protocol, SupportsIndex, TypeVar import numpy as np import torch from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from lerobot.datasets.utils import check_delta_timestamps, get_delta_indices from qwen_vl_utils.vision_process import smart_resize from torch.utils.data import DistributedSampler, random_split from transformers import AutoProcessor @@ -788,6 +789,27 @@ def load_lerobot_data( video_backend="pyav", ) + # Some locally patched LeRobot releases return from timestamp + # validation before initializing delta_indices. Restore the normal + # upstream behavior so action chunks contain the requested horizon. + if ( + train_dataset.delta_timestamps is not None + and train_dataset.delta_indices is None + ): + check_delta_timestamps( + train_dataset.delta_timestamps, + train_dataset.fps, + train_dataset.tolerance_s, + ) + train_dataset.delta_indices = get_delta_indices( + train_dataset.delta_timestamps, train_dataset.fps + ) + logger.warning( + "[Global rank %s] LeRobot did not initialize delta_indices; " + "restored them in Wall-X for action chunk loading", + global_rank, + ) + logger.info( "[Global rank %s] Finished loading on local_rank=%s", global_rank, diff --git a/wall_x/data/backends/lerobot/utils.py b/wall_x/data/backends/lerobot/utils.py index e002840..606ba31 100644 --- a/wall_x/data/backends/lerobot/utils.py +++ b/wall_x/data/backends/lerobot/utils.py @@ -10,6 +10,33 @@ import numpy as np import torch from transformers import BatchFeature +KEY_MAPPINGS = { + "lerobot/aloha_mobile_cabinet": { + "camera": { + "observation.images.cam_high": "face_view", + "observation.images.cam_left_wrist": "left_wrist_view", + "observation.images.cam_right_wrist": "right_wrist_view", + }, + "state": "observation.state", + "action": "action", + }, + "physical-intelligence/libero": { + "camera": { + "image": "face_view", + "wrist_image": "left_wrist_view", + }, + "state": "state", + "action": "actions", + }, + "/home/xiehaolv/huanghuagui/dataset/aloha_sim_transfer_cube_human": { + "camera": { + "observation.images.top": "face_view", + }, + "state": "observation.state", + "action": "action", + }, ##根据meta查询出来的结果，这一数据集只有"observation.images.top"，所以其余部分实际上无法得到引用结果 +} ##真实进入训练时需要对这一部分进行修改 + logger = logging.getLogger(__name__) diff --git a/wall_x/model/core/ops/moe.py b/wall_x/model/core/ops/moe.py index c50d946..314f98a 100644 --- a/wall_x/model/core/ops/moe.py +++ b/wall_x/model/core/ops/moe.py @@ -5,6 +5,7 @@ Mixture-of-Experts processing. """ import logging +import os import torch @@ -13,6 +14,14 @@ from wall_x.model.core.ops.base import OpsProxy logger = logging.getLogger(__name__) +def _force_pytorch_backend() -> bool: + return os.environ.get("WALL_X_FORCE_TORCH_MOE_OPS", "").lower() in { + "1", + "true", + "yes", + } + + class PermuteOp(OpsProxy): """Reorder tokens by expert assignment for MoE processing. @@ -21,9 +30,11 @@ class PermuteOp(OpsProxy): @property def _external_accel_name(self): - return "permute" + return None if _force_pytorch_backend() else "permute" def _get_cuda_kernel(self): + if _force_pytorch_backend(): + return None try: from wall_x.model.core.ops._cuda_wrappers import permute_kernel @@ -64,9 +75,11 @@ class UnpermuteOp(OpsProxy): @property def _external_accel_name(self): - return "unpermute" + return None if _force_pytorch_backend() else "unpermute" def _get_cuda_kernel(self): + if _force_pytorch_backend(): + return None try: from wall_x.model.core.ops._cuda_wrappers import unpermute_kernel diff --git a/wall_x/model/core/vla_mixin.py b/wall_x/model/core/vla_mixin.py index b438589..9d9bf64 100644 --- a/wall_x/model/core/vla_mixin.py +++ b/wall_x/model/core/vla_mixin.py @@ -321,6 +321,20 @@ class ActionModelMixMin: noisy_action_emb = noisy_action_emb.to( inputs_embeds.device, inputs_embeds.dtype ) + action_token_count = int(mask.sum().item()) + if noisy_action_emb.ndim != 3 or ( + action_token_count * inputs_embeds.shape[-1] + != noisy_action_emb.numel() + ): + raise ValueError( + "Flow action placeholder/embedding mismatch: " + f"input_ids={tuple(input_ids.shape)}, " + f"action_tokens={action_token_count}, " + f"action_chunk={tuple(action_chunk.shape)}, " + f"inputs_embeds={tuple(inputs_embeds.shape)}, " + f"noisy_action_emb={tuple(noisy_action_emb.shape)}, " + f"configured_horizon={self.config.action_horizon_flow}" + ) inputs_embeds = inputs_embeds.masked_scatter(action_mask, noisy_action_emb) return inputs_embeds, flow, adarms_cond @@ -530,7 +544,10 @@ class ActionGenerationMixin(GenerationMixin): target_modules=target_modules, lora_dropout=lora_dropout, bias="none", - task_type="CAUSAL_LM", + # ``self.model`` is Wall-X's decoder, not the outer generation model. + # Leave task_type unset so PEFT uses its generic wrapper and forwards + # Wall-X's custom MoE arguments unchanged. + task_type=None, ) self.model = get_peft_model(self.model, config) # Log trainable parameter information diff --git a/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py index 9eff1fe..1f76b13 100644 --- a/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py +++ b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py @@ -65,7 +65,7 @@ from transformers.utils import ( from .configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig from wall_x.model.core.attention.selector import AttentionsSelectorMixin from wall_x.model.core.ops import rot_pos_emb, get_window_index, m_rope - +##这一部分代码是对flash_attn的import的处理，在确认可使用的情况下再将flash_attn相关库做引入 if is_flash_attn_2_available(): from flash_attn import flash_attn_func from flash_attn import flash_attn_varlen_func @@ -91,8 +91,8 @@ except ImportError: logger = logging.get_logger(__name__) _CONFIG_FOR_DOC = "Qwen2_5_VLConfig" - - +#GEMM对齐实际上是将模型参数按照128一组对齐，实际上是tile的处理匹配，那么训练时没有对应的处理，所以下载之后需要做pad +#针对VL模型中的MLP部分做实现与优化，做了GEMM对齐，gate_up合并Linear，做了checkpoint的pad / strip，以及activate的选择性重算 class Qwen2_5_VLMLP(nn.Module): # Align intermediate_size to this value for cuBLAS GEMM tile efficiency. # E.g. 3420 -> 3456 (= 128 * 27), improving MFU from 34.8% to ~60%+. @@ -1000,9 +1000,13 @@ class Qwen2_5_VLRotaryEmbedding(nn.Module): else "cpu" ) with torch.autocast(device_type=device_type, enabled=False): - freqs = ( - inv_freq_expanded.float() @ position_ids_expanded.float() - ).transpose(2, 3) + # This is an outer product: (..., rotary_dim, 1) x + # (..., 1, sequence_length). Broadcasting is mathematically + # identical to batched matmul and avoids a cuBLAS failure seen + # intermittently under multi-GPU FSDP. + freqs = (inv_freq_expanded.float() * position_ids_expanded.float()).transpose( + 2, 3 + ) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() sin = emb.sin() diff --git a/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py index 4a0f991..1bda139 100644 --- a/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py +++ b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py @@ -1,3 +1,4 @@ +import os import re import torch import torch.nn as nn @@ -809,6 +810,11 @@ class Qwen2_5_VLMoEForAction( use_selective_recompute=False, ): super().__init__(config) + if not getattr(config, "use_cuda_moe_ops", True): + os.environ["WALL_X_FORCE_TORCH_MOE_OPS"] = "1" + logger.warning_once( + "Using PyTorch MoE routing ops because use_cuda_moe_ops=false" + ) self.visual = self._build_visual(config, use_selective_recompute) self.model = Qwen2_5_VLMoEModel( config, use_selective_recompute=use_selective_recompute @@ -850,6 +856,8 @@ class Qwen2_5_VLMoEForAction( target_modules=config.lora_target_modules, lora_dropout=config.lora_dropout, ) + if getattr(config, "lora_train_action_expert", False): + self._set_lora_action_expert_trainable() # Initialize weights and apply final processing self.post_init() self._post_init_engine(config) @@ -861,6 +869,34 @@ class Qwen2_5_VLMoEForAction( use_selective_recompute=use_selective_recompute, ) + def _set_lora_action_expert_trainable(self) -> None: + """Train LoRA adapters and the action expert while freezing the VLM base.""" + action_expert_keywords = ( + "action_preprocessor", + "moe.experts.1", + "qkv_proj_experts.1", + "o_proj_experts.1", + "input_layernorms.1", + "post_attention_layernorms.1", + "model.norms.1", + ) + trainable_count = 0 + trainable_parameters = 0 + for name, param in self.named_parameters(): + is_lora = "lora_" in name + is_action_expert = any(keyword in name for keyword in action_expert_keywords) + # Dataset normalizer statistics are model Parameters but must remain fixed. + is_normalizer_stat = "normalizer" in name and not is_lora + param.requires_grad = (is_lora or is_action_expert) and not is_normalizer_stat + if param.requires_grad: + trainable_count += 1 + trainable_parameters += param.numel() + logger.info( + "LoRA action-expert mode: %d trainable tensors (%d parameters)", + trainable_count, + trainable_parameters, + ) + def _post_init_engine(self, config): """Hook called at the end of __init__. Subclasses may override to add engine-specific state.""" pass diff --git a/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act_rtc.py b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act_rtc.py new file mode 100644 index 0000000..c0c9f5f --- /dev/null +++ b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act_rtc.py @@ -0,0 +1,3139 @@ +import os +import re +import torch +import torch.nn as nn +from dataclasses import dataclass +from torch.nn import CrossEntropyLoss +from typing import Optional, List, Tuple +from torchdiffeq import odeint +import numpy as np +from typing import Any, Dict, Union +from wall_x.model.qact.qwen2_5.configuration_qwen2_5_vl import Qwen2_5_VLConfig +from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl import ( + QWEN2_5_VL_ATTENTION_CLASSES, + Qwen2_5_VLMLP, + Qwen2_5_VLRotaryEmbedding, + Qwen2_5_VLPreTrainedModel, + Qwen2RMSNorm, + Qwen2_5_VLPatchMerger, + Qwen2_5_VLVisionBlock, + Qwen2_5_VisionTransformerPretrainedModel, +) + +from transformers import AutoConfig + +from transformers.cache_utils import ( + Cache, + DynamicCache, + StaticCache, +) + +try: + from transformers.cache_utils import SlidingWindowCache +except ImportError: + # Compatibility with newer transformers versions + SlidingWindowCache = StaticCache + +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + ModelOutput, +) +from transformers.utils import ( + logging, + is_torchdynamo_compiling, +) +from transformers.modeling_attn_mask_utils import ( + AttentionMaskConverter, +) +from wall_x.model.core.vla_mixin import ActionGenerationMixin, ActionModelMixMin +from wall_x.model.core.action.normalizer import ( + normalize_data_with_virtual_tail, + unnormalize_data_with_virtual_tail, +) +from wall_x.model.core.action.moe import TokenTypeRouter, SparseMoeBlock +from wall_x.model.core.attention.selector import ( + ATTENTION_TYPES_WITH_2D_MASK, + ATTENTION_TYPES_WITH_FLASH_MASK, +) +from wall_x.model.core.attention.joint import JOINT_QWEN_ATTENTION_CLASSES +from wall_x.model.qact.qwen2_5.inference_mixin import VLAInferenceMixin + +from wall_x.model.core.ops import unpermute, permute, get_rope_index + +logger = logging.get_logger(__name__) + +_QWEN25_DMUON_BLOCKED_NAME_PARTS = ( + "embed_tokens", + "pos_embed", + "lm_head", + "norm", +) + + +def _is_qwen25_dmuon_target_param(name: str, param: nn.Parameter) -> bool: + """Select trainable Qwen2.5 VLA matrices for DMuon.""" + if not param.requires_grad or param.ndim != 2 or not name.endswith(".weight"): + return False + return not any(part in name for part in _QWEN25_DMUON_BLOCKED_NAME_PARTS) + + +@dataclass +class Qwen2_5_VLACausalLMOutputWithPast(ModelOutput): + loss: Optional[torch.FloatTensor] = None + flow_loss: Optional[torch.FloatTensor] = None + cross_entropy_loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[List[torch.FloatTensor]] = None + hidden_states: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[torch.FloatTensor]] = None + rope_deltas: Optional[torch.LongTensor] = None + + channel_loss_dict: Optional[dict[torch.FloatTensor]] = None + channel_loss_count_dict: Optional[dict[torch.FloatTensor]] = None + + +class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module, ActionModelMixMin): + def __init__( + self, + config: Qwen2_5_VLConfig, + layer_idx: int, + num_experts: int, + use_selective_recompute: bool = True, + ): + super().__init__() + self.hidden_size = config.hidden_size + self.use_selective_recompute = use_selective_recompute + + if ( + config.use_sliding_window + and config._attn_implementation != "flash_attention_2" + ): + logger.warning_once( + f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " + "unexpected results may be encountered." + ) + if config.attention_moe: + self.self_attn = JOINT_QWEN_ATTENTION_CLASSES[config._attn_implementation]( + config, layer_idx + ) + else: + self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation]( + config, layer_idx + ) + + if config.use_adarms: + adarms_cond_dims = [None, config.adarms_cond_dim] + else: + adarms_cond_dims = [None, None] + + if config.norm_moe: + self.input_layernorms = nn.ModuleList( + [ + Qwen2RMSNorm( + config.dim_inputs[i], + eps=config.rms_norm_eps, + cond_dim=adarms_cond_dims[i], + ) + for i in range(num_experts) + ] + ) + self.post_attention_layernorms = nn.ModuleList( + [ + Qwen2RMSNorm( + config.dim_inputs[i], + eps=config.rms_norm_eps, + cond_dim=adarms_cond_dims[i], + ) + for i in range(num_experts) + ] + ) + self.input_layernorm, self.post_attention_layernorm = None, None + else: + self.input_layernorm = Qwen2RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = Qwen2RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.input_layernorms, self.post_attention_layernorms = None, None + + if config.mlp_moe: + self.router = TokenTypeRouter(num_experts=num_experts) + self.moe = SparseMoeBlock( + config, + num_experts=num_experts, + use_selective_recompute=use_selective_recompute, + ) + self.mlp = None + else: + self.mlp = Qwen2_5_VLMLP(config) + self.moe, self.router = None, None + + self.config = config + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC + # for vla + token_types: Optional[torch.LongTensor] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + probs: Optional[torch.Tensor] = None, + row_id_map: Optional[torch.Tensor] = None, + orig_shape: Optional[Tuple[int, int, int]] = None, + adarms_conds: Optional[List[torch.Tensor]] = [None, None], + **kwargs, + ) -> Tuple[ + torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] + ]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. + position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + + residual = hidden_states + hidden_states, gate, gate_mask = self._apply_norm_moe( + hidden_states, + token_types, + adarms_conds, + self.input_layernorms, + self.input_layernorm, + start_indices, + end_indices, + self.use_selective_recompute, + ) + + # Self Attention + if self.config.attention_moe: + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + token_types=token_types, + start_indices=start_indices, + end_indices=end_indices, + probs=probs, + row_id_map=row_id_map, + orig_shape=orig_shape, + position_embeddings=position_embeddings, + ) + else: + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + + hidden_states = self._gated_residual( + residual, hidden_states, gate, start_indices, end_indices + ) + + # Fully Connected + residual = hidden_states + + hidden_states, gate, gate_mask = self._apply_norm_moe( + hidden_states, + token_types, + adarms_conds, + self.post_attention_layernorms, + self.post_attention_layernorm, + start_indices, + end_indices, + self.use_selective_recompute, + ) + + hidden_states = self._apply_mlp_moe( + hidden_states, token_types, start_indices, end_indices + ) + + hidden_states = self._gated_residual( + residual, hidden_states, gate, start_indices, end_indices + ) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + if use_cache: + outputs += (present_key_value,) + return outputs + + +class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin): + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path, num_experts=None, *args, **kwargs + ): + # If num_experts is provided, make sure it is added to config + config = kwargs.get("config", None) + if config is None: + config = AutoConfig.from_pretrained(pretrained_model_name_or_path) + + if num_experts is not None: + config.num_experts = num_experts + + kwargs["config"] = config + return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) + + def __init__(self, config: Qwen2_5_VLConfig, use_selective_recompute=False): + super().__init__(config) + self.config = config + self.use_selective_recompute = use_selective_recompute + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx + ) + self.layers = nn.ModuleList( + [ + Qwen2_5_VLDecoderLayer_with_MoE( + config, + layer_idx, + config.num_experts, + use_selective_recompute=use_selective_recompute, + ) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self._attn_implementation = config._attn_implementation + + if config.use_adarms: + adarms_cond_dims = [None, config.adarms_cond_dim] + else: + adarms_cond_dims = [None, None] + + if config.norm_moe: + self.norms = nn.ModuleList( + [ + Qwen2RMSNorm( + config.dim_inputs[i], + eps=config.rms_norm_eps, + cond_dim=adarms_cond_dims[i], + ) + for i in range(config.num_experts) + ] + ) + self.norm = None + else: + self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norms = None + + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, # new parameter + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + positional_masks: Optional[ + dict + ] = None, # stores token position masks needed by each category + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + adarms_conds: Optional[List[torch.Tensor]] = [None, None], + **kwargs, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) + if moe_token_types is None: + raise ValueError("moe_token_types must be provided for MoE routing") + if start_indices is None or end_indices is None: + raise ValueError( + "start_indices and end_indices must be provided for MoE routing" + ) + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # torch.jit.trace() doesn't support cache objects in the output + if use_cache and past_key_values is None and not torch.jit.is_tracing(): + past_key_values = DynamicCache() + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if cache_position is None: + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + cache_position = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ) + + # the hard coded `3` is for temporal, height and width. + if position_ids is None: + position_ids = cache_position.view(1, 1, -1).expand( + 3, inputs_embeds.shape[0], -1 + ) + elif position_ids.dim() == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + + if not self.config.attention_moe: + causal_mask = self._update_causal_mask( + attention_mask, + inputs_embeds, + cache_position, + past_key_values, + output_attentions, + moe_token_types, + ) + else: + causal_mask = attention_mask + + hidden_states = inputs_embeds + + if ( + self.config._attn_implementation != "flash_attention_2" + and self.config.attention_moe is True + ): + position_ids = self._update_position_ids( + position_ids, moe_token_types, positional_masks + ) + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # When mot_opt is enabled, permute tokens from different experts first; shape becomes [Tokens, HiddenSize] + orig_shape = hidden_states.shape + moe_token_types = moe_token_types.contiguous() + if self.config.mot_opt: + hidden_states = hidden_states.view(-1, hidden_states.size(-1)) + hidden_states, row_id_map = permute(hidden_states, moe_token_types.view(-1)) + else: + row_id_map = None + # Use reshape instead of view to support non-contiguous tensors + probs = torch.ones_like(moe_token_types, dtype=torch.float32).reshape(-1, 1) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + # generate 2d attention mask if needed + if ( + self.config._attn_implementation in ATTENTION_TYPES_WITH_2D_MASK + and self.config.attention_moe is True + ): + if causal_mask is not None and inputs_embeds.shape[1] > 1: + causal_mask = self._update_joint_attention_mask_2d( + attention_mask=causal_mask, + moe_token_types=moe_token_types, + positional_masks=positional_masks, + ) + + if ( + self.config._attn_implementation in ATTENTION_TYPES_WITH_FLASH_MASK + and self.config.attention_moe is True + ): + causal_mask = self._update_joint_attention_flash_mask( + attention_mask=causal_mask, + moe_token_types=moe_token_types, + positional_masks=positional_masks, + ) + + for decoder_layer in self.layers: + if output_hidden_states: + assert ( + self.config.mot_opt is False + ), "When using mot_opt, output_hidden_states is not supported yet." + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + # for vla + moe_token_types, + start_indices, + end_indices, + probs, + row_id_map, + orig_shape, + adarms_conds, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + # for vla + token_types=moe_token_types, + start_indices=start_indices, + end_indices=end_indices, + probs=probs, + row_id_map=row_id_map, + orig_shape=orig_shape, + adarms_conds=adarms_conds, + ) + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states, _, _ = self._apply_norm_moe( + hidden_states, + moe_token_types, + adarms_conds, + self.norms, + self.norm, + start_indices, + end_indices, + self.use_selective_recompute, + ) + + # add hidden states from the last decoder layer + if output_hidden_states: + assert ( + self.config.mot_opt is False + ), "When using mot_opt, output_hidden_states is not supported yet." + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + + if self.config.mot_opt: + hidden_states = unpermute(hidden_states, row_id_map, probs) + hidden_states = hidden_states.view(orig_shape) + + if not return_dict: + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] + if v is not None + ) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + moe_token_types: Optional[torch.LongTensor] = None, + ): + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and past_key_values is not None: + is_padding_right = ( + attention_mask[:, -1].sum().item() != input_tensor.size()[0] + ) + if is_padding_right: + raise ValueError( + "You are attempting to perform batched generation with padding_side='right'" + " this may lead to unexpected behaviour for Flash Attention version of Qwen2_5_VL. Make sure to " + " call `tokenizer.padding_side = 'left'` before tokenizing the input. " + ) + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + using_static_cache = isinstance(past_key_values, StaticCache) + using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if ( + self.config._attn_implementation == "sdpa" + and not (using_static_cache or using_sliding_window_cache) + and not output_attentions + ): + if attention_mask.ndim == 2: + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + sliding_window=self.config.sliding_window, + is_training=self.training, + ): + return None + elif attention_mask.ndim == 3: + return attention_mask + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + # SlidingWindowCache or StaticCache + if using_sliding_window_cache or using_static_cache: + target_length = past_key_values.get_max_cache_shape() + # DynamicCache or no cache + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). + causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=target_length, + dtype=dtype, + device=device, + cache_position=cache_position, + batch_size=input_tensor.shape[0], + config=self.config, + past_key_values=past_key_values, + ) + # Modify the mask to support bidirectional attention + if moe_token_types is not None: + # Find all token positions with type 1 + type1_tokens = ( + (moe_token_types == 1).unsqueeze(1).unsqueeze(2) + ) # [B, 1, 1, S] + + # Create a square mask for the type-1 region + type1_mask = torch.zeros_like(causal_mask) # [B, num_heads, S, S] + type1_region = type1_tokens & type1_tokens.transpose(-1, -2) # [B, 1, S, S] + type1_mask = type1_mask.masked_fill(type1_region, 1.0).to(torch.bool) + # Zero the original causal_mask in the type-1 region, then add type1_mask + causal_mask = torch.where( + type1_mask, # Expand dimensions to match causal_mask + torch.zeros_like(causal_mask), # zero the type-1 region + causal_mask, # keep other regions unchanged + ) + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type in ["cuda", "xpu"] + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended( + causal_mask, min_dtype + ) + + return causal_mask + + @staticmethod + def _prepare_4d_causal_attention_mask_with_cache_position( + attention_mask: torch.Tensor, + sequence_length: int, + target_length: int, + dtype: torch.dtype, + device: torch.device, + cache_position: torch.Tensor, + batch_size: int, + config: Qwen2_5_VLConfig, + past_key_values: Cache, + ): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. + + Args: + attention_mask (`torch.Tensor`): + A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. + sequence_length (`int`): + The sequence length being processed. + target_length (`int`): + The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. + dtype (`torch.dtype`): + The dtype to use for the 4D attention mask. + device (`torch.device`): + The device to plcae the 4D attention mask on. + cache_position (`torch.Tensor`): + Indices depicting the position of the input sequence tokens in the sequence. + batch_size (`torch.Tensor`): + Batch size. + config (`Qwen2_5_VLConfig`): + The model's configuration class + past_key_values (`Cache`): + The cache class that is being used currently to generate + """ + if attention_mask is not None and attention_mask.dim() == 4: + # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. + causal_mask = attention_mask + else: + min_dtype = torch.finfo(dtype).min + causal_mask = torch.full( + (sequence_length, target_length), + fill_value=min_dtype, + dtype=dtype, + device=device, + ) + diagonal_attend_mask = torch.arange( + target_length, device=device + ) > cache_position.reshape(-1, 1) + if config.sliding_window is not None: + # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also + # the check is needed to verify is current checkpoint was trained with sliding window or not + if ( + not isinstance(past_key_values, SlidingWindowCache) + or sequence_length > target_length + ): + sliding_attend_mask = torch.arange( + target_length, device=device + ) <= (cache_position.reshape(-1, 1) - config.sliding_window) + diagonal_attend_mask.bitwise_or_(sliding_attend_mask) + causal_mask *= diagonal_attend_mask + causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) + if attention_mask is not None: + causal_mask = ( + causal_mask.clone() + ) # copy to contiguous memory for in-place edit + if attention_mask.shape[-1] > target_length: + attention_mask = attention_mask[:, :target_length] + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[ + :, None, None, : + ].to(causal_mask.device) + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[ + :, :, :, :mask_length + ].masked_fill(padding_mask, min_dtype) + return causal_mask + + +class Qwen2_5_VLMoEForAction( + Qwen2_5_VLPreTrainedModel, + ActionGenerationMixin, + ActionModelMixMin, + VLAInferenceMixin, +): + # Compatibility with newer transformers versions (5.x):_tied_weights_keys changed from list[str] to dict[str, str] (target -> source mapping). + # Older versions (4.x) iterate over dict keys, so the behavior is equivalent and compatible. + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + config_class = Qwen2_5_VLConfig + _no_split_modules = ["Qwen2_5_VLDecoderLayer_with_MoE", "Qwen2_5_VLVisionBlock"] + + def __init__( + self, + config: Qwen2_5_VLConfig, + processor=None, + tokenizer_mixin=None, + use_selective_recompute=False, + ): + super().__init__(config) + if not getattr(config, "use_cuda_moe_ops", True): + os.environ["WALL_X_FORCE_TORCH_MOE_OPS"] = "1" + logger.warning_once( + "Using PyTorch MoE routing ops because use_cuda_moe_ops=false" + ) + self.visual = self._build_visual(config, use_selective_recompute) + self.model = Qwen2_5_VLMoEModel( + config, use_selective_recompute=use_selective_recompute + ) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + self.loss_fct = CrossEntropyLoss( + reduction="none" + ) # do not do reduction to compute channel loss + + # Manage through tokenizer_mixin; in pure flow mode the mixin is None + self.processor = processor + self.tokenizer_mixin = tokenizer_mixin + if self.tokenizer_mixin is not None: + self.action_tokenizer_type = self.tokenizer_mixin.tokenizer_type + self.action_tokenizer = self.tokenizer_mixin.tokenizer + self.action_mapper = self.tokenizer_mixin.build_action_mapper(processor) + else: + self.action_tokenizer_type = None + self.action_tokenizer = None + self.action_mapper = None + self.define_action_token_id() + + self.times_cache = {} # cache times linspace for each num_inference_timesteps + self._infer_stable_cache: dict = {} # cache stable-state inference structures + + self.rope_deltas = None # cache rope_deltas here + + from wall_x.model.core.action.processor import ActionProcessor + + self.action_preprocessor = ActionProcessor(config) # action processing + + # Apply LoRA if the configuration contains LoRA settings + if hasattr(config, "use_lora") and config.use_lora: + self.add_lora( + r=config.lora_r, + lora_alpha=config.lora_alpha, + target_modules=config.lora_target_modules, + lora_dropout=config.lora_dropout, + ) + if getattr(config, "lora_train_action_expert", False): + self._set_lora_action_expert_trainable() + # Initialize weights and apply final processing + self.post_init() + self._post_init_engine(config) + + def _build_visual(self, config, use_selective_recompute): + """Factory method for visual encoder. Subclasses may override to swap implementation.""" + return Qwen2_5_VisionTransformerPretrainedModel( + config=config.vision_config, + use_selective_recompute=use_selective_recompute, + ) + + def _set_lora_action_expert_trainable(self) -> None: + """Train LoRA adapters and the action expert while freezing the VLM base.""" + action_expert_keywords = ( + "action_preprocessor", + "moe.experts.1", + "qkv_proj_experts.1", + "o_proj_experts.1", + "input_layernorms.1", + "post_attention_layernorms.1", + "model.norms.1", + ) + trainable_count = 0 + trainable_parameters = 0 + for name, param in self.named_parameters(): + is_lora = "lora_" in name + is_action_expert = any(keyword in name for keyword in action_expert_keywords) + # Dataset normalizer statistics are model Parameters but must remain fixed. + is_normalizer_stat = "normalizer" in name and not is_lora + param.requires_grad = (is_lora or is_action_expert) and not is_normalizer_stat + if param.requires_grad: + trainable_count += 1 + trainable_parameters += param.numel() + logger.info( + "LoRA action-expert mode: %d trainable tensors (%d parameters)", + trainable_count, + trainable_parameters, + ) + + def _post_init_engine(self, config): + """Hook called at the end of __init__. Subclasses may override to add engine-specific state.""" + pass + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def get_rope_index( + self, + input_ids: Optional[torch.LongTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Calculate the 3D rope index based on image and video's temporal, height and width in LLM. + + Explanation: + Each embedding sequence contains vision embedding and text embedding or just contains text embedding. + + For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs. + Examples: + input_ids: [T T T T T], here T is for text. + temporal position_ids: [0, 1, 2, 3, 4] + height position_ids: [0, 1, 2, 3, 4] + width position_ids: [0, 1, 2, 3, 4] + + For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part + and 1D rotary position embeddin for text part. + Examples: + Temporal (Time): 3 patches, representing different segments of the video in time. + Height: 2 patches, dividing each frame vertically. + Width: 2 patches, dividing each frame horizontally. + We also have some important parameters: + fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second. + tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity. + temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames. + interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs. + input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision. + vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100] + vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] + vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1] + text temporal position_ids: [101, 102, 103, 104, 105] + text height position_ids: [101, 102, 103, 104, 105] + text width position_ids: [101, 102, 103, 104, 105] + Here we calculate the text start position_ids as the max vision position_ids plus 1. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*): + The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + Returns: + position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`) + mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`) + """ + spatial_merge_size = self.config.vision_config.spatial_merge_size + image_token_id = self.config.image_token_id + video_token_id = self.config.video_token_id + vision_start_token_id = self.config.vision_start_token_id + mrope_position_deltas = [] + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + for i, input_ids in enumerate(total_input_ids): + input_ids = input_ids[attention_mask[i] == 1] + image_nums, video_nums = 0, 0 + vision_start_indices = torch.argwhere( + input_ids == vision_start_token_id + ).squeeze(1) + vision_tokens = input_ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + input_tokens = input_ids.tolist() + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + second_per_grid_t = 0 + image_index += 1 + remain_images -= 1 + ed = ed_image + + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + if second_per_grid_ts is not None: + second_per_grid_t = second_per_grid_ts[video_index] + else: + second_per_grid_t = 1.0 + video_index += 1 + remain_videos -= 1 + ed = ed_video + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + range_tensor = torch.arange(llm_grid_t).view(-1, 1) + expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) + + time_tensor = ( + expanded_range + * second_per_grid_t + * self.config.vision_config.tokens_per_second + ) + + time_tensor_long = time_tensor.long() + t_index = time_tensor_long.flatten() + + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx + ) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = ( + llm_pos_ids_list[-1].max() + 1 + if len(llm_pos_ids_list) > 0 + else 0 + ) + text_len = len(input_tokens) - st + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( + position_ids.device + ) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]) + ) + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=input_ids.device + ).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = ( + position_ids.unsqueeze(0) + .expand(3, -1, -1) + .to(attention_mask.device) + ) + max_position_ids = position_ids.max(0, keepdim=False)[0].max( + -1, keepdim=True + )[0] + mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] + else: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + + return position_ids, mrope_position_deltas + + def get_attention_maps( + self, + input_ids, + action_horizon, + action_dim, + num_inference_timesteps: int = 10, + prefix_length: Optional[int] = None, # prefix length + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + positional_masks: Optional[torch.LongTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + action_chunk: Optional[torch.FloatTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, + unnorm_proprioception: Optional[torch.FloatTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + dataset_names: Optional[str] = None, + dof_mask: Optional[torch.FloatTensor] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + **kwargs, + ): + + if input_ids is not None: + batch_size, seq_length = input_ids.shape + + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if start_indices is None or end_indices is None: + # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2 + ): + # calculate RoPE index once per generation in the pre-fill stage only + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): + position_ids, rope_deltas = get_rope_index( + input_ids=input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + attention_mask=attention_mask, + spatial_merge_size=self.config.vision_config.spatial_merge_size, + image_token_id=self.config.image_token_id, + video_token_id=self.config.video_token_id, + vision_start_token_id=self.config.vision_start_token_id, + tokens_per_second=self.config.vision_config.tokens_per_second, + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + # batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(cache_position.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=cache_position.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + img_mask = input_ids == self.config.image_token_id + mask_unsqueezed = img_mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + image_mask = mask_expanded.to(inputs_embeds.device) + + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + + mask = input_ids == self.config.video_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + video_mask = mask_expanded.to(inputs_embeds.device) + + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + inputs_embeds = self.scatter_proprioception_embeddings( + input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask + ) + inputs_embeds, flow, adarms_cond = self.scatter_flow_action_embeddings( + input_ids, inputs_embeds, action_chunk, dataset_names, dof_mask + ) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + outputs = self.model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + moe_token_types=moe_token_types, # pass token_types + start_indices=start_indices, + end_indices=end_indices, + positional_masks=positional_masks, # pass masks for each category + use_cache=False, + output_attentions=True, + output_hidden_states=False, + return_dict=return_dict, + adarms_conds=[None, adarms_cond], + # cache_position=cache_position, + ) + attention_maps = torch.stack(outputs.attentions, dim=0)[ + None, :, -action_horizon: + ] # [flow steps, layer depths, all tokens, all tokens] + attention_maps = attention_maps.mean(2) + image_mask_indices = img_mask[0].nonzero(as_tuple=True)[0] + attention_maps = attention_maps[:, :, image_mask_indices] + return attention_maps + + def train_step_forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + # for vla + sample_time: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + positional_masks: Optional[dict] = None, + action_chunk: Optional[torch.FloatTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, + dataset_names: Optional[str] = None, + dof_mask: Optional[torch.FloatTensor] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + flow_loss_mask: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> Union[Tuple, Qwen2_5_VLACausalLMOutputWithPast]: + if input_ids is not None: + batch_size, seq_length = input_ids.shape + + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if start_indices is None or end_indices is None: + # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2 + ): + # calculate RoPE index once per generation in the pre-fill stage only + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): + position_ids, rope_deltas = get_rope_index( + input_ids=input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + attention_mask=attention_mask, + spatial_merge_size=self.config.vision_config.spatial_merge_size, + image_token_id=self.config.image_token_id, + video_token_id=self.config.video_token_id, + vision_start_token_id=self.config.vision_start_token_id, + tokens_per_second=self.config.vision_config.tokens_per_second, + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + # batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(cache_position.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=cache_position.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + mask = input_ids == self.config.image_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + image_mask = mask_expanded.to(inputs_embeds.device) + + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + + mask = input_ids == self.config.video_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + video_mask = mask_expanded.to(inputs_embeds.device) + + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + inputs_embeds = self.scatter_proprioception_embeddings( + input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask + ) + inputs_embeds, flow, adarms_cond = self.scatter_flow_action_embeddings( + input_ids, + inputs_embeds, + action_chunk, + dataset_names, + sample_time, + dof_mask, + ) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + outputs = self.model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + moe_token_types=moe_token_types, # pass token_types + start_indices=start_indices, + end_indices=end_indices, + positional_masks=positional_masks, # pass masks for each category + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + adarms_conds=[None, adarms_cond], + # cache_position=cache_position, + ) + + hidden_states = outputs[0] + + # --- optimization: loss-token-only LM Head --- + # During training, only compute lm_head for tokens that contribute to CE loss, + # reducing GEMM size from [B*S, V] to [N_loss, V]. + # FSDP fp32 layernorms emit fp32 hidden_states; lm_head runs in bf16. + _lm_dtype = self.lm_head.weight.dtype + _lm_loss_mask = None + if labels is not None and self.training: + shift_labels = labels[..., 1:].contiguous() + _lm_loss_mask = shift_labels != -100 # [B, S-1] + if _lm_loss_mask.any(): + lm_hidden = hidden_states[:, :-1, :][_lm_loss_mask] # [N_loss, H] + logits = self.lm_head(lm_hidden.to(_lm_dtype)) # [N_loss, V] + else: + logits = None + else: + logits = self.lm_head(hidden_states.to(_lm_dtype)) + + ( + loss, + cross_entropy_loss, + flow_loss, + channel_loss_dict, + channel_loss_count_dict, + ) = self.compute_loss( + hidden_states=hidden_states, + logits=logits, + input_ids=input_ids, + dataset_names=dataset_names, + labels=labels, + action_chunk=action_chunk, + dof_mask=dof_mask, + flow=flow, + flow_loss_mask=flow_loss_mask, + _lm_loss_mask=_lm_loss_mask, + ) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return Qwen2_5_VLACausalLMOutputWithPast( + loss=loss, + cross_entropy_loss=( + cross_entropy_loss.clone() if cross_entropy_loss is not None else None + ), + flow_loss=flow_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=self.rope_deltas, + channel_loss_dict=channel_loss_dict, + channel_loss_count_dict=channel_loss_count_dict, + ) + + def predict_action(self, predict_mode: str, **kwargs): + assert predict_mode in ["fast", "spatial_token", "flow", "diffusion", "dllm"] + + if predict_mode in ["fast", "spatial_token"]: + output = self.generate_ar_action(**kwargs) + elif predict_mode in ["diffusion", "flow"]: + output = self.generate_flow_action(**kwargs) + elif predict_mode == "dllm": + output = self.generate_dllm_action(**kwargs) + + return output["predict_action"], output.get("gt_action", None) + + @torch.no_grad() + def generate_text( + self, + input_ids: torch.LongTensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, # new parameter + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_side_image: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + side_image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + action_chunk: Optional[torch.FloatTensor] = None, # action chunk + proprioception: Optional[torch.FloatTensor] = None, # joint positions + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + dataset_names: Optional[str] = None, + dof_mask: Optional[torch.FloatTensor] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + gt_output_ids: Optional[torch.LongTensor] = None, + prefix_length: Optional[int] = None, # prefix length + positional_masks: Optional[dict] = None, # position masks for each modality + re_generate: bool = False, + **kwargs, + ): + # assert self.config._attn_implementation == "flash_attention_2", "generate_text only support flash_attention_2 attn implementation" + + # generate text only need prefix part + if prefix_length is not None: + input_ids, _gt_output_ids = ( + input_ids[:, :prefix_length], + input_ids[:, prefix_length:], + ) + gt_output_ids = ( + gt_output_ids if gt_output_ids is not None else _gt_output_ids + ) + attention_mask = ( + attention_mask[:, :prefix_length] + if attention_mask is not None + else None + ) + moe_token_types = ( + moe_token_types[:, :prefix_length] + if moe_token_types is not None + else None + ) + else: + prefix_length = input_ids.shape[1] + + if start_indices is None or end_indices is None: + # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + + batch = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "moe_token_types": moe_token_types, + "start_indices": start_indices, + "end_indices": end_indices, + "image_grid_thw": image_grid_thw, + "dof_mask": dof_mask, + "agent_pos_mask": agent_pos_mask, + "proprioception": proprioception, + "dataset_names": dataset_names, + # "prefix_length": prefix_length, + } + predict_output_ids = self.generate( + **batch, + max_new_tokens=100, + eos_token_id=[ + self.processor.tokenizer.eos_token_id + ], # set multiple end markers + use_cache=True, + pad_token_id=self.processor.tokenizer.pad_token_id, # explicitly set pad_token_id + temperature=( + 1.0 if not re_generate else 0.7 + ), # use a higher temperature when regenerating + do_sample=( + False if not re_generate else True + ), # use sampling when regenerating + ) + input_text = self.processor.batch_decode( + input_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True + ) + if gt_output_ids is not None and gt_output_ids.shape[-1] != 0: + gt_output_text = self.processor.batch_decode( + gt_output_ids, + skip_special_tokens=False, + clean_up_tokenization_spaces=True, + ) + else: + gt_output_text = None + predict_output_text = self.processor.batch_decode( + predict_output_ids[:, prefix_length:], + skip_special_tokens=False, + clean_up_tokenization_spaces=True, + ) + output = { + "input_text": input_text, + "gt_output_text": gt_output_text, + "predict_output_text": predict_output_text, + "predict_output_ids": predict_output_ids, + } + return output + + @torch.no_grad() + def generate_ar_action( + self, + input_ids, + action_horizon, + action_dim, + state=None, + dataset_names=None, + action_chunk=None, + dof_mask=None, + unnorm=True, + max_retries=3, # maximum retry count + proprioception=None, # proprioception for v3.1_delta decoding + robot_type_id=None, # robot type ID for v3.1_delta decoding + **kwargs, + ): + # Regeneration mechanism + predict_action = None + output = None + decode_success = False + if dataset_names is None and (unnorm or action_chunk is not None): + raise KeyError( + "generate_ar_action requires dataset_names when normalizing or " + "unnormalizing actions." + ) + + for retry_attempt in range(max_retries): + use_re_generate = retry_attempt > 0 + + if retry_attempt > 0: + logger.info( + "Retrying action generation (attempt %s/%s)...", + retry_attempt + 1, + max_retries, + ) + + output = self.generate_text( + input_ids=input_ids, re_generate=use_re_generate, **kwargs + ) + predict_output_ids = output["predict_output_ids"] + + # Decode uniformly through tokenizer_mixin + predict_action, decode_success = self.tokenizer_mixin.decode_action( + output_ids=predict_output_ids, + action_mapper=self.action_mapper, + action_horizon=action_horizon, + action_dim=action_dim, + device=input_ids.device, + proprioception=proprioception, + dof_mask=dof_mask, + robot_type_id=robot_type_id, + state=state, + ) + + if decode_success: + if retry_attempt > 0: + logger.info( + "Action generation succeeded on attempt %s", + retry_attempt + 1, + ) + break + else: + if retry_attempt < max_retries - 1: + logger.warning( + "Error in decoding action (attempt %s/%s), retrying with re_generate mode...", + retry_attempt + 1, + max_retries, + ) + else: + logger.warning( + "Error in decoding action after %s attempts, returning None", + max_retries, + ) + + # Check whether dof_mask is required through mixin attributes + uses_dof_mask = ( + self.tokenizer_mixin.uses_dof_mask_for_unnorm + if self.tokenizer_mixin is not None + else True + ) + + if action_chunk is not None: + action_chunk = action_chunk.to(input_ids.device).to(torch.bfloat16) + if uses_dof_mask: + action_chunk = action_chunk[:, :, dof_mask[0, 0, :].bool()] + output["gt_action"] = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + action_chunk, dataset_names, dof_mask + ) + ) + else: + # Output full dimensions without mask filtering + output["gt_action"] = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + action_chunk, dataset_names, None + ) + ) + + if not decode_success or predict_action is None: + logger.warning("Error in decoding action, predict_action is None") + output["predict_action"] = None + else: + if unnorm: + if isinstance(predict_action, np.ndarray): + predict_action = torch.tensor( + predict_action, device=input_ids.device + ) + elif predict_action.device != input_ids.device: + predict_action = predict_action.to(input_ids.device) + + # Add the batch dimension + if predict_action.dim() == 2: + predict_action = predict_action.unsqueeze(0) + + if uses_dof_mask: + predict_action = unnormalize_data_with_virtual_tail( + self.action_preprocessor.normalizer_action, + predict_action, + dataset_names, + dof_mask, + ) + else: + # Output full dimensions without mask filtering + predict_action = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names, None + ) + ) + output["predict_action"] = predict_action + + return output + + @torch.no_grad() + def generate_flow_action( + self, + input_ids, + action_horizon, + action_dim, + num_inference_timesteps: int = 10, + prefix_length: Optional[int] = None, # prefix length + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + positional_masks: Optional[torch.LongTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + action_chunk: Optional[torch.FloatTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, + unnorm_proprioception: Optional[torch.FloatTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + dataset_names: Optional[str] = None, + dof_mask: Optional[torch.FloatTensor] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + rtc_processor=None, + prev_chunk_left_over: Optional[torch.FloatTensor] = None, + inference_delay: int = 0, + execution_horizon: Optional[int] = None, + **kwargs, + ): + + # assert self.config._attn_implementation == "sdpa", "generate_flow_action only support sdpa attn implementation" + ctx = self._prepare_flow_action_inputs( + input_ids=input_ids, + action_horizon=action_horizon, + action_dim=action_dim, + num_inference_timesteps=num_inference_timesteps, + prefix_length=prefix_length, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + moe_token_types=moe_token_types, + start_indices=start_indices, + end_indices=end_indices, + positional_masks=positional_masks, + labels=labels, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + action_chunk=action_chunk, + proprioception=proprioception, + unnorm_proprioception=unnorm_proprioception, + rope_deltas=rope_deltas, + cache_position=cache_position, + second_per_grid_ts=second_per_grid_ts, + dataset_names=dataset_names, + dof_mask=dof_mask, + agent_pos_mask=agent_pos_mask, + ) + ctx["rtc_processor"] = rtc_processor + ctx["prev_chunk_left_over"] = prev_chunk_left_over + ctx["inference_delay"] = max(0, int(inference_delay or 0)) + ctx["execution_horizon"] = execution_horizon + + action_trajectory, attention_maps = self._execute_flow_action(ctx) + + output = self._finalize_flow_action_output(action_trajectory, ctx) + if output_attentions: + output["attention_maps"] = attention_maps + return output + + def _execute_flow_action(self, ctx): + """Strategy hook: selects which flow-action backend to use. + + Returns: + action_trajectory: list of action tensors + attention_maps: attention maps or None + """ + return self._generate_flow_action_vanilla(ctx) + + def _prepare_flow_action_inputs( + self, + input_ids, + action_horizon, + action_dim, + num_inference_timesteps, + prefix_length, + attention_mask, + position_ids, + past_key_values, + inputs_embeds, + moe_token_types, + start_indices, + end_indices, + positional_masks, + labels, + use_cache, + output_attentions, + output_hidden_states, + return_dict, + pixel_values, + pixel_values_videos, + image_grid_thw, + video_grid_thw, + action_chunk, + proprioception, + unnorm_proprioception, + rope_deltas, + cache_position, + second_per_grid_ts, + dataset_names, + dof_mask, + agent_pos_mask, + ): + batch_size = ( + input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] + ) + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # Timing: input embedding processing + img_mask = None + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + # from wall_x.utils.timers import ScopeTimer + # with ScopeTimer("pixel_values.visual"): + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_features = image_embeds.shape[0] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + + img_mask = input_ids == self.config.image_token_id + mask_unsqueezed = img_mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + image_mask = mask_expanded.to(inputs_embeds.device) + + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + + mask = input_ids == self.config.video_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + video_mask = mask_expanded.to(inputs_embeds.device) + + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + if ( + proprioception is not None + and not self.config.use_state_string_representation + ): + proprioception = proprioception.to(inputs_embeds.device) + agent_pos_mask = agent_pos_mask.to(inputs_embeds.device) + proprio_embed = self.action_preprocessor.proprioception_proj( + proprioception, + dataset_names, + agent_pos_mask, + use_history=proprioception.shape[1] > 1, + ) + proprioception_mask = ( + input_ids == self.action_token_id_set["propri_token_id"] + ) + inputs_embeds[proprioception_mask] = proprio_embed.reshape( + -1, inputs_embeds.shape[-1] + ).to(inputs_embeds.dtype) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2 + ): + # calculate RoPE index once per generation in the pre-fill stage only + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): + position_ids, rope_deltas = get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + attention_mask, + self.config.vision_config.spatial_merge_size, + self.config.image_token_id, + self.config.video_token_id, + self.config.vision_start_token_id, + self.config.vision_config.tokens_per_second, + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + flow_action_mask = input_ids == self.action_token_id_set["action_token_id"] + + if prefix_length is None: + has_true = flow_action_mask.any(dim=1) + prefix_length = torch.argmax(flow_action_mask.float(), dim=1, keepdim=True) + prefix_length[~has_true] = flow_action_mask.shape[1] + # check if prefix_length is the same for all batch + if not torch.all(prefix_length == prefix_length[0]): + raise ValueError( + "prefix_length differs across batch; batch prompts must align" + ) + prefix_length = int(prefix_length[0].item()) + + if start_indices is None or end_indices is None: + # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here + # Cache key: shape + sum is sufficient for stable-state (moe_token_types structure doesn't change) + _moe_key = (moe_token_types.shape, int(moe_token_types.sum().item())) + _moe_cache = self._infer_stable_cache.get("moe_indices") + if _moe_cache is None or _moe_cache[0] != _moe_key: + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + self._infer_stable_cache["moe_indices"] = ( + _moe_key, + start_indices, + end_indices, + ) + else: + _, start_indices, end_indices = _moe_cache + + # Timing: action initialization + if action_chunk is not None: + action_chunk = action_chunk.to(inputs_embeds.device).to(torch.float32) + + # Reproduce + # torch.manual_seed(0) + noise = torch.randn( + size=(batch_size, action_horizon, action_dim), + dtype=torch.float32, + device=inputs_embeds.device, + ) + noisy_action = noise.clone() + + dof_mask = dof_mask.to(inputs_embeds.device).to(torch.float32) + + if num_inference_timesteps not in self.times_cache: + self.times_cache[num_inference_timesteps] = ( + self.action_preprocessor.get_inference_times( + num_inference_timesteps, inputs_embeds.device, torch.float32 + ) + ) + times = self.times_cache[num_inference_timesteps] + time_0 = times[0].unsqueeze(0).repeat(noisy_action.shape[0]) + action_embed, adarms_cond = self.action_preprocessor.step( + timestep=time_0, noisy_action=noisy_action, dof_mask=dof_mask + ) + action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]).to( + inputs_embeds.dtype + ) + inputs_embeds[flow_action_mask] = action_embed + + # Automatically generate padding actions from dof_mask + if not dof_mask.all(): + padding_action = ( + torch.zeros((1, dof_mask.shape[-1])) + .to(dof_mask.device) + .to(torch.float32) + ) + padding_action = normalize_data_with_virtual_tail( + self.action_preprocessor.normalizer_action, + padding_action, + dataset_names, + dof_mask, + ) + else: + padding_action = None + + return { + "batch_size": batch_size, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "position_ids": position_ids, + "prefix_length": prefix_length, + "flow_action_mask": flow_action_mask, + "start_indices": start_indices, + "end_indices": end_indices, + "dof_mask": dof_mask, + "noise": noise, + "noisy_action": noisy_action, + "times": times, + "adarms_cond": adarms_cond, + "action_chunk": action_chunk, + "padding_action": padding_action, + "action_horizon": action_horizon, + "action_dim": action_dim, + "input_ids": input_ids, + "moe_token_types": moe_token_types, + "positional_masks": positional_masks, + "dataset_names": dataset_names, + "output_attentions": output_attentions, + "img_mask": img_mask, + } + + def _generate_flow_action_vanilla(self, ctx): + # Timing: prefill forward pass + all_attention_maps = [] + output_attentions = True if ctx["output_attentions"] else False + prefetch_output = self.model( + input_ids=None, + attention_mask=ctx["attention_mask"], + position_ids=ctx["position_ids"], + past_key_values=None, + inputs_embeds=ctx["inputs_embeds"], + moe_token_types=ctx["moe_token_types"], + start_indices=ctx["start_indices"], + end_indices=ctx["end_indices"], + positional_masks=ctx["positional_masks"], + use_cache=True, + output_attentions=output_attentions, + output_hidden_states=False, + return_dict=True, + adarms_conds=[None, ctx["adarms_cond"]], + ) + hidden_states = prefetch_output.last_hidden_state + prefix_kv_cache = prefetch_output.past_key_values + if output_attentions: + all_attention_maps.append(prefetch_output.attentions) + action_hidden_states = hidden_states[ctx["flow_action_mask"]].to(torch.float32) + action_pred = self.action_preprocessor.action_proj_back( + action_hidden_states[:, : self.action_preprocessor.action_hidden_size] + ) + if getattr(self.config, "use_x_pred", False): + v_0 = action_pred - ctx["noise"].reshape(-1, ctx["noise"].shape[-1]) + else: + v_0 = action_pred + v_0 = v_0.reshape(ctx["batch_size"], ctx["action_horizon"], ctx["action_dim"]) + + v_padding = None + if (not ctx["dof_mask"].all()) and (ctx["padding_action"] is not None): + v_padding = ctx["padding_action"] - ctx["noisy_action"] + v_0 = (v_padding) * (1 - ctx["dof_mask"]) + v_0 * ctx["dof_mask"] + dt = ctx["times"][1] - ctx["times"][0] + ctx["noisy_action"] = ctx["noisy_action"] + dt * v_0 + + ( + postfix_position_ids, + postfix_inputs_embeds, + _postfix_attention_mask, + postfix_moe_token_types, + postfix_input_ids, + postfix_start_indices, + postfix_end_indices, + padding_mask, + ) = self._prepare_flow_postfix(ctx, prefix_kv_cache) + + def step_with_kvcache(timestep, noisy_action): + action_mask = ( + postfix_input_ids == self.action_token_id_set["action_token_id"] + ) + assert action_mask.any(), "No action token found in input_ids" + timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0]) + action_embed, adarms_cond = self.action_preprocessor.step( + timestep=timestep, noisy_action=noisy_action, dof_mask=ctx["dof_mask"] + ) + action_embed = action_embed.reshape(-1, postfix_inputs_embeds.shape[-1]) + + # Create a temporary inputs_embeds copy for thread safety + temp_inputs_embeds = postfix_inputs_embeds.clone() + temp_inputs_embeds[action_mask] = action_embed.to(temp_inputs_embeds.dtype) + transformer_outputs = self.model( + input_ids=None, + attention_mask=_postfix_attention_mask, + position_ids=postfix_position_ids, + past_key_values=prefix_kv_cache, + inputs_embeds=temp_inputs_embeds, + moe_token_types=postfix_moe_token_types, + start_indices=postfix_start_indices, + end_indices=postfix_end_indices, + use_cache=False, + output_attentions=output_attentions, + output_hidden_states=False, + return_dict=True, + adarms_conds=[None, adarms_cond], + ) + if output_attentions: + all_attention_maps.append(transformer_outputs.attentions) + hidden_states = transformer_outputs.last_hidden_state + action_hidden_states = hidden_states[action_mask].to(torch.float32) + action_pred = self.action_preprocessor.action_proj_back( + action_hidden_states[:, : self.action_preprocessor.action_hidden_size] + ) + if getattr(self.config, "use_x_pred", False): + # Align noisy_action (and timestep) to action_pred's shape + B, action_horizon, action_dim = noisy_action.shape + noisy_action_flat = noisy_action.reshape( + -1, action_dim + ) # [B * action_horizon, action_dim] + timestep_expand = timestep.view(-1, 1).repeat_interleave( + action_horizon, dim=0 + ) # [B * action_horizon, 1] + v_t = (action_pred - noisy_action_flat) / torch.clamp( + 1 - timestep_expand, min=0.05 + ) + else: + v_t = action_pred + v_t = v_t.reshape( + ctx["batch_size"], ctx["action_horizon"], ctx["action_dim"] + ) + + if (not ctx["dof_mask"].all()) and (ctx["padding_action"] is not None): + v_t = (v_padding) * (1 - ctx["dof_mask"]) + v_t * ctx["dof_mask"] + + return v_t + + # The first Euler step above bootstraps the reusable prefix KV cache. + # RTC guidance is applied to every subsequent denoising step. + current_action = ctx["noisy_action"] + action_states = [current_action] + rtc_processor = ctx.get("rtc_processor") + rtc_prefix = ctx.get("prev_chunk_left_over") + use_rtc = rtc_processor is not None and rtc_prefix is not None + for step_index in range(1, len(ctx["times"]) - 1): + timestep = ctx["times"][step_index] + next_timestep = ctx["times"][step_index + 1] + + def predict_velocity(action, current_timestep=timestep): + return step_with_kvcache(current_timestep, action) + + if use_rtc: + velocity = rtc_processor.guide_increasing_flow( + x_t=current_action, + time=timestep, + predict_velocity=predict_velocity, + prev_chunk_left_over=rtc_prefix, + inference_delay=ctx["inference_delay"], + execution_horizon=ctx["execution_horizon"], + ) + else: + velocity = predict_velocity(current_action) + current_action = current_action + (next_timestep - timestep) * velocity + action_states.append(current_action) + + action_trajectory = torch.stack(action_states, dim=0) + + attention_maps = None + if output_attentions and ctx["img_mask"] is not None: + attention_maps = [torch.stack(map, dim=0) for map in all_attention_maps] + attention_maps = torch.stack( + attention_maps, dim=0 + ) # [flow steps, layer depths, action tokens, all tokens] + attention_maps = attention_maps.mean(2) + image_mask_indices = ctx["img_mask"][0].nonzero(as_tuple=True)[0] + attention_maps = attention_maps[:, :, image_mask_indices] + return action_trajectory, attention_maps + + def _prepare_flow_postfix(self, ctx, prefix_kv_cache=None): + postfix_position_ids = ctx["position_ids"][:, :, ctx["prefix_length"] :] + postfix_inputs_embeds = ctx["inputs_embeds"][:, ctx["prefix_length"] :, :] + postfix_attention_mask = ctx["attention_mask"][:, ctx["prefix_length"] :] + postfix_moe_token_types = ctx["moe_token_types"][:, ctx["prefix_length"] :] + postfix_input_ids = ctx["input_ids"][:, ctx["prefix_length"] :] + + postfix_start_indices = None + postfix_end_indices = None + padding_mask = None + + if prefix_kv_cache is not None: + if ctx["prefix_length"] is None: + has_true = ctx["flow_action_mask"].any(dim=1) + prefix_length = torch.argmax( + ctx["flow_action_mask"].float(), dim=1, keepdim=True + ) + prefix_length[~has_true] = ctx["flow_action_mask"].shape[1] + # check if prefix_length is the same for all batch + if not torch.all(prefix_length == prefix_length[0]): + raise ValueError( + "prefix_length differs across batch; batch prompts must align" + ) + prefix_length = int(prefix_length[0].item()) + ctx["prefix_length"] = prefix_length + + if hasattr(prefix_kv_cache, "key_cache"): + for layer_i in range(len(prefix_kv_cache.key_cache)): + prefix_kv_cache.key_cache[layer_i] = prefix_kv_cache.key_cache[ + layer_i + ][:, :, : ctx["prefix_length"], :] + prefix_kv_cache.value_cache[layer_i] = prefix_kv_cache.value_cache[ + layer_i + ][:, :, : ctx["prefix_length"], :] + else: + for layer_i in range(len(prefix_kv_cache.layers)): + prefix_kv_cache.layers[layer_i].keys = prefix_kv_cache.layers[ + layer_i + ].keys[:, :, : ctx["prefix_length"], :] + prefix_kv_cache.layers[layer_i].values = prefix_kv_cache.layers[ + layer_i + ].values[:, :, : ctx["prefix_length"], :] + + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (postfix_moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + postfix_start_indices = torch.cumsum(group_size, dim=0) - group_size + postfix_end_indices = torch.cumsum(group_size, dim=0) + + pad_token_id = self.processor.tokenizer.pad_token_id + padding_mask = ctx["input_ids"] == pad_token_id + else: + pad_token_id = self.processor.tokenizer.pad_token_id + padding_mask = ctx["input_ids"] == pad_token_id + + postfix_length = ctx["input_ids"].shape[-1] - ctx["prefix_length"] + _postfix_attention_mask = torch.ones( + (ctx["batch_size"], postfix_length, ctx["prefix_length"] + postfix_length), + dtype=torch.bool, + device=postfix_attention_mask.device, + ) + + postfix_padding_mask = padding_mask[:, ctx["prefix_length"] :] + full_padding_mask = padding_mask + + if self.config.causal_action_attention_mask: + _postfix_attention_mask[:, :, ctx["prefix_length"] :] = torch.tril( + torch.ones( + (postfix_length, postfix_length), + dtype=torch.bool, + device=postfix_attention_mask.device, + ) + ) + + # postfix_padding_mask: [B, postfix_length], True where padding (query rows) + # full_padding_mask: [B, prefix_length+postfix_length], True where padding (key cols) + _postfix_attention_mask.masked_fill_(postfix_padding_mask.unsqueeze(2), False) + _postfix_attention_mask.masked_fill_(full_padding_mask.unsqueeze(1), False) + + return ( + postfix_position_ids, + postfix_inputs_embeds, + _postfix_attention_mask, + postfix_moe_token_types, + postfix_input_ids, + postfix_start_indices, + postfix_end_indices, + padding_mask, + ) + + def _finalize_flow_action_output(self, action_trajectory, ctx): + predict_action = action_trajectory[-1] + normalized_predict_action = predict_action.detach().clone() + predict_action = unnormalize_data_with_virtual_tail( + self.action_preprocessor.normalizer_action, + predict_action, + ctx["dataset_names"], + ctx.get("dof_mask"), + ) + output = { + "predict_action": predict_action, + "_rtc_normalized_predict_action": normalized_predict_action, + } + if ctx["action_chunk"] is not None: + output["gt_action"] = unnormalize_data_with_virtual_tail( + self.action_preprocessor.normalizer_action, + ctx["action_chunk"], + ctx["dataset_names"], + ctx.get("dof_mask"), + ) + + return output + + def forward(self, mode: Optional[str] = None, **kwargs): + if not mode: + with torch.no_grad(): + return self.train_step_forward(**kwargs) + elif mode == "train": + return self.train_step_forward(use_cache=False, **kwargs) + elif mode == "validate": + with torch.no_grad(): + return self.train_step_forward(use_cache=False, **kwargs) + else: + raise NotImplementedError("invalid key") + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + moe_token_types=None, + start_indices=None, + end_indices=None, + cache_position=None, + position_ids=None, + use_cache=True, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + second_per_grid_ts=None, + dataset_names=None, + proprioception=None, + dof_mask=None, + agent_pos_mask=None, + **kwargs, + ): + # Overwritten -- in specific circumstances we don't want to forward image inputs to the model + + # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens + # Exception 1: when passing input_embeds, input_ids may be missing entries + # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here + # Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case. + # (we can't check exception 3 while compiling) + # Exception 4: If input_embeds are passed then slice it through `cache_position`, to keep only the unprocessed tokens and + # generate the first token for each sequence. Later use the generated Input ids for continuation. + + if moe_token_types is None: + moe_token_types = torch.zeros_like( + input_ids + ) # FIXME input_ids or input_embeds + else: + # Ensure moe_token_types has the same length as input_ids + if moe_token_types.shape[1] < input_ids.shape[1]: + # Compute the required padding length + pad_length = input_ids.shape[1] - moe_token_types.shape[1] + # Create the padding tensor using 0 as the default token type + pad_tensor = torch.zeros( + (moe_token_types.shape[0], pad_length), + dtype=moe_token_types.dtype, + device=moe_token_types.device, + ) + # Append padding after the existing moe_token_types + moe_token_types = torch.cat([moe_token_types, pad_tensor], dim=1) + + if past_key_values is not None: + if inputs_embeds is not None and input_ids.shape[1] == 0: # Exception 4 + inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :] + moe_token_types = moe_token_types[:, -cache_position.shape[0] :] + elif inputs_embeds is not None or ( # Exception 1 + is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1] + ): # Exception 3 + input_ids = input_ids[:, -cache_position.shape[0] :] + moe_token_types = moe_token_types[:, -cache_position.shape[0] :] + elif ( + input_ids.shape[1] != cache_position.shape[0] + ): # Default case (the "else", a no op, is Exception 2) + cache_pos = cache_position.clone() + input_ids = input_ids[:, cache_pos] + moe_token_types = moe_token_types[:, cache_pos] + + if cache_position[0] != 0: + pixel_values = None + pixel_values_videos = None + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]: + model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None} + else: + model_inputs = {"input_ids": input_ids, "inputs_embeds": None} + + if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2: + if model_inputs["inputs_embeds"] is not None: + batch_size, sequence_length, _ = inputs_embeds.shape + device = inputs_embeds.device + else: + batch_size, sequence_length = input_ids.shape + device = input_ids.device + + attention_mask = ( + self.model._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=past_key_values.get_max_cache_shape(), + dtype=self.lm_head.weight.dtype, + device=device, + cache_position=cache_position, + batch_size=batch_size, + config=self.config, + past_key_values=past_key_values, + ) + ) + + if ( + self.config._attn_implementation == "flash_attention_2_ki" + and self.config.attention_moe is True + ): + attention_mask = None # Pass None as the mask to use the official flash attention for autoregressive prediction + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "moe_token_types": moe_token_types, + "start_indices": start_indices, + "end_indices": end_indices, + "use_cache": use_cache, + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "pixel_values_videos": pixel_values_videos, + "image_grid_thw": image_grid_thw, + "video_grid_thw": video_grid_thw, + "cache_position": cache_position, + "second_per_grid_ts": second_per_grid_ts, + "proprioception": proprioception, + "dataset_names": dataset_names, + "dof_mask": dof_mask, + "agent_pos_mask": agent_pos_mask, + } + ) + return model_inputs + + def _get_image_nums_and_video_nums( + self, + input_ids: Optional[torch.LongTensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Get the number of images and videos for each sample to calculate the separation length of the sample tensor. + These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Returns: + image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`) + video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`) + """ + image_token_id = self.config.image_token_id + video_token_id = self.config.video_token_id + vision_start_token_id = self.config.vision_start_token_id + + vision_start_mask = input_ids == vision_start_token_id + vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1) + image_mask = input_ids == image_token_id + video_mask = input_ids == video_token_id + image_nums = torch.sum(vision_first_mask & image_mask, dim=1) + video_nums = torch.sum(vision_first_mask & video_mask, dim=1) + + return image_nums, video_nums + + def _expand_inputs_for_generation( + self, + expand_size: int = 1, + is_encoder_decoder: bool = False, + input_ids: Optional[torch.LongTensor] = None, + **model_kwargs, + ) -> Tuple[torch.LongTensor, Dict[str, Any]]: + # Overwritten -- Support for expanding tensors without a batch size dimension + # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t + # pixel_values.shape[0] is sum(seqlen_images for samples) + # image_grid_thw.shape[0] is sum(num_images for samples) + + if expand_size == 1: + return input_ids, model_kwargs + + visual_keys = [ + "pixel_values", + "image_grid_thw", + "pixel_values_videos", + "video_grid_thw", + "second_per_grid_ts", + ] + + def _expand_dict_for_generation_visual(dict_to_expand): + image_grid_thw = model_kwargs.get("image_grid_thw", None) + video_grid_thw = model_kwargs.get("video_grid_thw", None) + image_nums, video_nums = self._get_image_nums_and_video_nums(input_ids) + + def _repeat_interleave_samples(x, lengths, repeat_times): + samples = torch.split(x, lengths) + repeat_args = [repeat_times] + [1] * (x.dim() - 1) + result = torch.cat( + [sample.repeat(*repeat_args) for sample in samples], dim=0 + ) + return result + + for key in dict_to_expand: + if key == "pixel_values": + # split images into samples + samples = torch.split(image_grid_thw, list(image_nums)) + # compute the sequence length of images for each sample + lengths = [torch.prod(sample, dim=1).sum() for sample in samples] + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "image_grid_thw": + # get the num of images for each sample + lengths = list(image_nums) + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "pixel_values_videos": + samples = torch.split(video_grid_thw, list(video_nums)) + lengths = [torch.prod(sample, dim=1).sum() for sample in samples] + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "video_grid_thw": + lengths = list(video_nums) + dict_to_expand[key] = _repeat_interleave_samples( + dict_to_expand[key], lengths=lengths, repeat_times=expand_size + ) + elif key == "second_per_grid_ts": + if not isinstance(dict_to_expand[key], list): + raise TypeError( + f"Expected value for key '{key}' to be a list, but got {type(dict_to_expand[key])} instead." + ) + tensor = torch.tensor(dict_to_expand[key]) + lengths = list(video_nums) + tensor = _repeat_interleave_samples( + tensor, lengths=lengths, repeat_times=expand_size + ) + dict_to_expand[key] = tensor.tolist() + return dict_to_expand + + def _expand_dict_for_generation(dict_to_expand): + for key in dict_to_expand: + if ( + key != "cache_position" + and dict_to_expand[key] is not None + and isinstance(dict_to_expand[key], torch.Tensor) + and key not in visual_keys + ): + dict_to_expand[key] = dict_to_expand[key].repeat_interleave( + expand_size, dim=0 + ) + return dict_to_expand + + # input_ids is required for expanding visual inputs + # If input_ids is unavailable, visual inputs will not be used; therefore, there is no need to expand visual inputs. + if input_ids is not None and input_ids.numel() != 0: + model_kwargs = _expand_dict_for_generation_visual(model_kwargs) + + if input_ids is not None: + input_ids = input_ids.repeat_interleave(expand_size, dim=0) + + model_kwargs = _expand_dict_for_generation(model_kwargs) + + if is_encoder_decoder: + if model_kwargs.get("encoder_outputs") is None: + raise ValueError( + "If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined." + ) + model_kwargs["encoder_outputs"] = _expand_dict_for_generation( + model_kwargs["encoder_outputs"] + ) + + return input_ids, model_kwargs + + def rename_vlm_weights_for_vla(self, merged_weights): + renamed = {} + + for key, value in merged_weights.items(): + + if key.startswith("model.layers") and "mlp." in key and self.config.mlp_moe: + layer_num = key.split(".layers.")[1].split(".mlp")[0] + new_key = key.replace( + f"layers.{layer_num}.mlp.", f"layers.{layer_num}.moe.experts.0." + ) + renamed[new_key] = value + continue + + if ( + key.startswith("model.layers") + and "self_attn." in key + and self.config.attention_moe + ): + layer_num = key.split(".layers.")[1].split(".self_attn")[0] + proj_types = ["q_proj", "k_proj", "v_proj", "o_proj"] + for proj in proj_types: + if proj in key: + new_key = key.replace( + f"layers.{layer_num}.self_attn.{proj}", + f"layers.{layer_num}.self_attn.{proj}_experts.0", + ) + renamed[new_key] = value + break + continue + + if self.config.norm_moe and ".input_layernorm." in key: + renamed[key.replace("input_layernorm", "input_layernorms.0")] = value + continue + if self.config.norm_moe and ".post_attention_layernorm." in key: + renamed[ + key.replace( + "post_attention_layernorm", "post_attention_layernorms.0" + ) + ] = value + continue + if self.config.norm_moe and ".norm." in key: + renamed[key.replace("norm", "norms.0")] = value + continue + + renamed[key] = value + + fused = Qwen2_5_VLMoEForAction.convert_to_fused(renamed) + + return fused + + @staticmethod + def fuse_gate_up( + fused, prefix, suffix_gate="gate_proj", suffix_up="up_proj", out="gate_up_proj" + ): + gate_w = fused.get(prefix + f"{suffix_gate}.weight") + up_w = fused.get(prefix + f"{suffix_up}.weight") + gate_b = fused.get(prefix + f"{suffix_gate}.bias") + up_b = fused.get(prefix + f"{suffix_up}.bias") + + # Skip fusion if gate_up_proj already exists + if prefix + f"{out}.weight" in fused: + return + + if gate_w is not None and up_w is not None: + fused[prefix + f"{out}.weight"] = torch.cat([gate_w, up_w], dim=0) + if gate_b is not None and up_b is not None: + fused[prefix + f"{out}.bias"] = torch.cat([gate_b, up_b], dim=0) + + # Remove the old modules + for n in [ + f"{suffix_gate}.weight", + f"{suffix_up}.weight", + f"{suffix_gate}.bias", + f"{suffix_up}.bias", + ]: + full = prefix + n + if full in fused: + del fused[full] + + @staticmethod + def fuse_qkv(fused, prefix, *, experts=False, expert_id=None): + """ + prefix: + - non-MoE: model.layers.X.self_attn. + - MoE: model.layers.X.self_attn. + + When experts=True, expert_id must be provided + """ + + if experts: + assert expert_id is not None, "expert_id must be provided for MoE attention" + + eid = expert_id + + q_w = fused.get(prefix + f"q_proj_experts.{eid}.weight") + k_w = fused.get(prefix + f"k_proj_experts.{eid}.weight") + v_w = fused.get(prefix + f"v_proj_experts.{eid}.weight") + + q_b = fused.get(prefix + f"q_proj_experts.{eid}.bias") + k_b = fused.get(prefix + f"k_proj_experts.{eid}.bias") + v_b = fused.get(prefix + f"v_proj_experts.{eid}.bias") + + out_w = f"qkv_proj_experts.{eid}.weight" + out_b = f"qkv_proj_experts.{eid}.bias" + + remove_list = [ + f"q_proj_experts.{eid}.weight", + f"k_proj_experts.{eid}.weight", + f"v_proj_experts.{eid}.weight", + f"q_proj_experts.{eid}.bias", + f"k_proj_experts.{eid}.bias", + f"v_proj_experts.{eid}.bias", + ] + + else: + q_w = fused.get(prefix + "q_proj.weight") + k_w = fused.get(prefix + "k_proj.weight") + v_w = fused.get(prefix + "v_proj.weight") + + q_b = fused.get(prefix + "q_proj.bias") + k_b = fused.get(prefix + "k_proj.bias") + v_b = fused.get(prefix + "v_proj.bias") + + out_w = "qkv_proj.weight" + out_b = "qkv_proj.bias" + + remove_list = [ + "q_proj.weight", + "k_proj.weight", + "v_proj.weight", + "q_proj.bias", + "k_proj.bias", + "v_proj.bias", + ] + + if q_w is None or k_w is None or v_w is None: + return + + fused[prefix + out_w] = torch.cat([q_w, k_w, v_w], dim=0) + + if q_b is not None and k_b is not None and v_b is not None: + fused[prefix + out_b] = torch.cat([q_b, k_b, v_b], dim=0) + + for n in remove_list: + full = prefix + n + if full in fused: + del fused[full] + + @staticmethod + def is_fused(state_dict): + # Check whether the model is already fused by looking for fused weight markers + return any(".moe.experts.0.gate_up_proj" in key for key in state_dict.keys()) + + @staticmethod + def convert_to_fused(state_dict): + """ + Convert an unfused checkpoint to a fused checkpoint + """ + fused = {} + + # ========================= + # Phase 1: copy everything first + # ========================= + for k, v in state_dict.items(): + fused[k] = v + + # ========================= + # Phase 2: fuse gate + up (MoE aware) + # ========================= + for key in list(fused.keys()): + # language MoE mlp (ANY expert) + m = re.match( + r"(model\.layers\.\d+\.moe\.experts\.\d+\.)gate_proj\.weight", + key, + ) + if m: + prefix = m.group(1) + Qwen2_5_VLMoEForAction.fuse_gate_up(fused, prefix) + continue + + # language non-MoE mlp + if re.match(r"(model\.layers\.\d+\.mlp\.)gate_proj\.weight", key): + prefix = key.replace("gate_proj.weight", "") + Qwen2_5_VLMoEForAction.fuse_gate_up(fused, prefix) + continue + + # visual mlp + if re.match(r"(visual\.blocks\.\d+\.mlp\.)gate_proj\.weight", key): + prefix = key.replace("gate_proj.weight", "") + Qwen2_5_VLMoEForAction.fuse_gate_up(fused, prefix) + + # ========================= + # Phase 3: fuse attention qkv (MoE aware) + # ========================= + for key in list(fused.keys()): + # language MoE attention: ANY expert + m = re.match( + r"(model\.layers\.\d+\.self_attn\.)q_proj_experts\.(\d+)\.weight", + key, + ) + if m: + prefix = m.group(1) + expert_id = int(m.group(2)) + Qwen2_5_VLMoEForAction.fuse_qkv( + fused, + prefix, + experts=True, + expert_id=expert_id, + ) + continue + + # language non-MoE attention + if re.match(r"(model\.layers\.\d+\.self_attn\.)q_proj\.weight", key): + prefix = key.replace("q_proj.weight", "") + Qwen2_5_VLMoEForAction.fuse_qkv( + fused, + prefix, + experts=False, + ) + continue + + # visual attention + if re.match(r"(visual\.blocks\.\d+\.attn\.)q_proj\.weight", key): + prefix = key.replace("q_proj.weight", "") + Qwen2_5_VLMoEForAction.fuse_qkv( + fused, + prefix, + experts=False, + ) + + # ========================= + # Phase 4: sanity check(optional) + # ========================= + for k in fused.keys(): + assert not any( + x in k + for x in [ + ".gate_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ] + ), f"Unfused key still exists: {k}" + + return fused + + def convert_to_mix_precision(self): + # Mix Precision + params_to_keep_float32 = [] + for name, _ in self.named_parameters(): + if any( + k in name + for k in [ + "input_layernorm", + "post_attention_layernorm", + "model.norm", + "action_preprocessor", + ] + ): + params_to_keep_float32.append(name) + for name, param in self.named_parameters(): + if name not in params_to_keep_float32: + param.data = param.data.to(torch.bfloat16) + if name in params_to_keep_float32: + param.data = param.data.to(torch.float32) + + def convert_to_fsdp( + self, + *, + mesh, + mp_policy, + offload_policy=None, + reshard_after_forward: bool = True, + use_dmuon: bool = False, + ): + """Wrap with FSDP2: per-decoder-layer + per-vision-block fully_shard + with bf16 ``mp_policy``; layernorms and ``action_preprocessor`` leaves + get their own nested ``fully_shard`` with an fp32 policy so their + forward stays in fp32 (matching the old FSDP1 NO_SHARD + fp32_policy + layout). + + ``mesh`` and ``mp_policy`` are required and produced by + ``FSDPStrategy._build_fsdp2_layout``. The pre-migration FSDP1 wrap + used ``ShardingStrategy.SHARD_GRAD_OP`` regardless of the yaml + config; the new path honors ``fsdp_sharding_strategy`` (default + ``full_shard`` -> ``reshard_after_forward=True``). Active runs that + relied on the implicit SHARD_GRAD_OP throughput should set + ``distributed.fsdp_sharding_strategy: shard_grad_op`` explicitly. + + When ``use_dmuon=True``, DMuon owns selected trainable + matrix weights before the remaining parameters are wrapped by FSDP2. + """ + from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy + + if use_dmuon: + import dmuon + + # DMuon keeps fp32 master weights for dedicated params; symmetric + # FSDP2 params still use MixedPrecisionPolicy for bf16 compute. + for param in self.parameters(): + if param.requires_grad and param.dtype != torch.float32: + param.data = param.data.float() + + action_preprocessor_linear_ids = { + id(module) + for module in self.action_preprocessor.modules() + if isinstance(module, nn.Linear) + } + + def hook_boundary(module: nn.Module) -> bool: + # ActionProcessor helper methods call internal Linear layers + # directly, bypassing ActionProcessor.forward(). Those internal + # Linear layers need their own hooks. + if id(module) in action_preprocessor_linear_ids: + return True + return isinstance( + module, + ( + Qwen2_5_VLDecoderLayer_with_MoE, + Qwen2_5_VLVisionBlock, + Qwen2_5_VLPatchMerger, + type(self.action_preprocessor), + ), + ) + + assignment = dmuon.dedicate_params( + self, + mesh, + predicate=_is_qwen25_dmuon_target_param, + hook_boundary_predicate=hook_boundary, + hook_boundary_strict=True, + compute_dtype=torch.bfloat16, + reshard_after_forward=reshard_after_forward, + ) + logger.info( + "[Qwen2.5 FSDP2 + DMuon] dedicate_params: %d params assigned", + len(assignment), + ) + + fp32_mp = MixedPrecisionPolicy( + param_dtype=torch.float32, + reduce_dtype=torch.float32, + ) + shard_kwargs = dict( + mesh=mesh, + mp_policy=mp_policy, + reshard_after_forward=reshard_after_forward, + ) + if offload_policy is not None: + shard_kwargs["offload_policy"] = offload_policy + fp32_shard_kwargs = dict(shard_kwargs, mp_policy=fp32_mp) + + def fully_shard_fp32_leaf_or_container(name: str, module: nn.Module) -> None: + # FSDP2 cannot shard containers without forward(), e.g. ModuleList. + # When a norm match resolves to a container, shard only norm leaves. + def is_norm_leaf(leaf_name: str, leaf: nn.Module) -> bool: + return ( + "norm" in leaf_name.lower() or "norm" in type(leaf).__name__.lower() + ) + + if isinstance(module, (nn.ModuleList, nn.ModuleDict)): + for leaf_name, leaf in module.named_modules(): + if not leaf_name: + continue + if any(True for _ in leaf.children()): + continue + if not is_norm_leaf(leaf_name, leaf): + continue + logger.info( + "[FSDP2] fully_shard fp32 layernorm: %s.%s", + name, + leaf_name, + ) + fully_shard(leaf, **fp32_shard_kwargs) + return + + logger.info("[FSDP2] fully_shard fp32 layernorm: %s", name) + fully_shard(module, **fp32_shard_kwargs) + + # Layernorm leaves inside decoder layers / vision blocks / model.norm + # - wrap with fp32 mp_policy BEFORE wrapping their parents so the + # nested policy survives. Match-by-FQN mirrors the FSDP1 logic. + for module_name, module in list(self.named_modules()): + for child_name, child in list(module.named_children()): + if any( + k in child_name.lower() + for k in ("input_layernorm", "post_attention_layernorm") + ) or ("norm" in child_name.lower() and module_name.endswith("model")): + fully_shard_fp32_leaf_or_container( + f"{module_name}.{child_name}", child + ) + + # action_preprocessor leaves: same fp32 treatment as layernorms. + for name, module in list(self.named_modules()): + if "action_preprocessor" not in name.lower(): + continue + if any(True for _ in module.children()): + continue # only leaves + logger.info( + "[FSDP2] fully_shard fp32 action_preprocessor leaf: %s", + name, + ) + fully_shard(module, **fp32_shard_kwargs) + + # Each transformer decoder block: bf16 mp_policy from strategy. + for idx, layer in enumerate(self.model.layers): + fully_shard(layer, **shard_kwargs) + if idx == 0: + logger.info( + "[FSDP2] fully_shard model.layers.* (bf16, " + "reshard_after_forward=%s)", + reshard_after_forward, + ) + + # Each vision tower block: same. + if hasattr(self, "visual") and hasattr(self.visual, "blocks"): + for idx, block in enumerate(self.visual.blocks): + fully_shard(block, **shard_kwargs) + if idx == 0: + logger.info("[FSDP2] fully_shard visual.blocks.* (bf16)") + + # Root: covers lm_head, embed_tokens, and any params not in a + # nested fully_shard. Inner fp32 wraps maintain their own policy. + logger.info("[FSDP2] fully_shard root (bf16)") + fully_shard(self, **shard_kwargs) + return self diff --git a/wall_x/trainer/adapters/base_adapter.py b/wall_x/trainer/adapters/base_adapter.py index 99f0803..ad615d8 100644 --- a/wall_x/trainer/adapters/base_adapter.py +++ b/wall_x/trainer/adapters/base_adapter.py @@ -94,6 +94,10 @@ class ModelAdapter(ABC): """Load pretrained weights, resize embeddings, set normalizers, etc.""" ... + def finalize_model_after_weight_load(self, model, model_config): + """Apply wrappers that must be created after loading base model weights.""" + return model + # ---- FSDP wrapping ---- @abstractmethod def get_transformer_layer_cls(self): diff --git a/wall_x/trainer/adapters/vla_model_adapter.py b/wall_x/trainer/adapters/vla_model_adapter.py index 955c9dc..89a7aff 100644 --- a/wall_x/trainer/adapters/vla_model_adapter.py +++ b/wall_x/trainer/adapters/vla_model_adapter.py @@ -1,3 +1,4 @@ +import json import os import torch @@ -107,6 +108,8 @@ class VLAdapter(ModelAdapter): flat = dataclasses.asdict(self.cfg.model) flat["model_type"] = self.cfg.model_type flat["data"] = dict(self.cfg._raw_data or {}) + flat["data"]["action_horizon"] = self.cfg.task.action_horizon + flat["data"]["action_horizon_flow"] = self.cfg.task.action_horizon_flow flat["dof_config"] = self.cfg.task.dof_config flat["agent_pos_config"] = self.cfg.task.agent_pos_config if self.cfg.task.ar_dof_config is not None: @@ -145,6 +148,16 @@ class VLAdapter(ModelAdapter): else: model_config = ConfigClass.from_pretrained(qwen_vl_act_config_path) + lora_config_path = getattr(self.cfg.model, "lora_config_path", None) + if lora_config_path: + with open(lora_config_path, encoding="utf-8") as f: + lora_overrides = json.load(f) + if not isinstance(lora_overrides, dict): + raise ValueError("model.lora_config_path must contain a JSON object") + for key, value in lora_overrides.items(): + setattr(model_config, key, value) + self.logger.info("Loaded LoRA overrides from %s", lora_config_path) + assert self.model_type in model_config.model_type, ( f"Mismatch of model type: model type in config file is " f"{model_config.model_type}, but the model type is {self.model_type}." @@ -156,12 +169,36 @@ class VLAdapter(ModelAdapter): def create_model(self, processor, tokenizer_mixin, model_config): ModelClass, _ = self._get_model_and_config_class() use_selective_recompute = self.cfg.distributed.use_selective_recompute - return ModelClass( - model_config, - processor, - tokenizer_mixin=tokenizer_mixin, - use_selective_recompute=use_selective_recompute, + # PEFT changes parameter names (base_model/base_layer). Build the plain + # model first so Qwen and Wall-OSS checkpoints load against native keys. + use_lora = bool(getattr(model_config, "use_lora", False)) + if use_lora: + model_config.use_lora = False + try: + return ModelClass( + model_config, + processor, + tokenizer_mixin=tokenizer_mixin, + use_selective_recompute=use_selective_recompute, + ) + finally: + if use_lora: + model_config.use_lora = True + + def finalize_model_after_weight_load(self, model, model_config): + """Inject LoRA only after all unwrapped checkpoint weights are loaded.""" + if not getattr(model_config, "use_lora", False): + return model + model.add_lora( + r=model_config.lora_r, + lora_alpha=model_config.lora_alpha, + target_modules=model_config.lora_target_modules, + lora_dropout=model_config.lora_dropout, ) + if getattr(model_config, "lora_train_action_expert", False): + model._set_lora_action_expert_trainable() + self.logger.info("Applied LoRA after loading base and Wall-OSS weights") + return model def load_weights(self, model, normalizer_action, normalizer_propri, **kwargs): import copy diff --git a/wall_x/trainer/fsdp_trainer/fsdp_trainer.py b/wall_x/trainer/fsdp_trainer/fsdp_trainer.py index 19ffd77..12b4f54 100644 --- a/wall_x/trainer/fsdp_trainer/fsdp_trainer.py +++ b/wall_x/trainer/fsdp_trainer/fsdp_trainer.py @@ -95,6 +95,18 @@ class FSDPTrainer(DistributedTrainer): self.load_processor() self.load_model() + # Single-file pretrained checkpoints use native model parameter names. + # Load them before PEFT or other parameter-name-changing wrappers. + if self._resume_from_single_file(): + self.load_state_dict( + self.model, + {"ckpt": self.cfg.checkpoint.resume_from}, + ) + + self.model = self.adapter.finalize_model_after_weight_load( + self.model, self.model_config + ) + # DDP needs to see correct requires_grad at wrap time, so freeze first. self._freeze_params_if_needed(self.model) @@ -107,12 +119,6 @@ class FSDPTrainer(DistributedTrainer): name: p.shape for name, p in self.model.named_parameters() } - if self._resume_from_single_file(): - self.load_state_dict( - self.model, - {"ckpt": self.cfg.checkpoint.resume_from}, - ) - self._wrap_model(self.model) self._create_optimizer() self._create_scheduler() diff --git a/wall_x/trainer/trainer_utils.py b/wall_x/trainer/trainer_utils.py index eab59f0..36f3b1d 100644 --- a/wall_x/trainer/trainer_utils.py +++ b/wall_x/trainer/trainer_utils.py @@ -161,6 +161,8 @@ def load_wallx_processors_from_cfg( flat = dataclasses.asdict(cfg.model) flat["model_type"] = cfg.model_type flat["data"] = dict(cfg._raw_data or {}) + flat["data"]["action_horizon"] = cfg.task.action_horizon + flat["data"]["action_horizon_flow"] = cfg.task.action_horizon_flow flat["dof_config"] = cfg.task.dof_config flat["agent_pos_config"] = cfg.task.agent_pos_config if cfg.task.ar_dof_config is not None: