Add Wall-X serving and Turtle2 TCP WebSocket bridge
Pre-commit / pre-commit (push) Canceled after 0s

This commit is contained in:
2026-09-23 21:04:17 +08:00
parent 6764e8f12f
commit d1cc7d96ad
40 changed files with 8591 additions and 40 deletions
+5
View File
@@ -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/
+84
View File
@@ -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: <All keys matched successfully>`。当前模型可能打印“词嵌入表被裁剪”和“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 合并和 `<All keys matched successfully>`。三路 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 仍需要现场与训练产物确认。
+159
View File
@@ -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: <All keys matched successfully>`;桥接日志出现 `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`,建议先停止机器人端,再停止桥接端和推流端。
+260
View File
@@ -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
+200
View File
@@ -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
```
+103
View File
@@ -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)
+41
View File
@@ -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
+36
View File
@@ -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}"
+342
View File
@@ -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[@]}"
+833
View File
@@ -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("<L", size_b)[0]
return await _recvall(reader, size)
async def _recv_state(reader):
raw = await _recv_frame(reader)
if raw is None:
return None
state = json.loads(raw.decode("utf-8"))
if not isinstance(state, dict):
raise ValueError("robot state is not a JSON object")
for key in ("follow1_pos", "follow2_pos"):
if key not in state:
raise ValueError(f"robot state missing required key {key!r}")
return state
async def _recv_image(reader, index):
raw = await _recv_frame(reader)
if raw is None:
raise ConnectionError("robot closed during image stream")
arr = np.frombuffer(raw, np.uint8)
img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
if img is None:
raise ValueError(f"failed to decode image #{index}")
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
def _img_to_b64(rgb):
ok, buf = cv2.imencode(".jpg", cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR))
if not ok:
raise ValueError("failed to JPEG-encode camera image")
return base64.b64encode(buf.tobytes()).decode("ascii")
def _as_action_chunk(value, key):
try:
arr = np.asarray(value, dtype=np.float64)
except (TypeError, ValueError) as exc:
raise ValueError(f"{key} must be numeric") from exc
if arr.ndim != 2 or arr.shape[1] != 7:
raise ValueError(f"{key} must have shape (T,7), got {arr.shape}")
if not np.isfinite(arr).all():
raise ValueError(f"{key} contains NaN/Inf")
return arr
def _single_arm_action(ws_action, current_right=None):
"""Decode a right-arm-only raw model response into a (T, 7) chunk."""
if "predict_action" not in ws_action:
raise ValueError("WebSocket response missing predict_action")
raw = np.asarray(ws_action["predict_action"], dtype=np.float64)
if raw.ndim == 3:
if raw.shape[0] != 1:
raise ValueError(f"predict_action batch must have size 1, got {raw.shape}")
raw = raw[0]
if raw.ndim != 2 or not np.isfinite(raw).all():
raise ValueError(f"predict_action must be finite with shape (T,D), got {raw.shape}")
# Training layout: 16 padding values, then position(3), rotation-6D(6),
# and gripper(1). Convert the rotation representation without importing
# the LIBERO/MuJoCo package, which is unavailable on headless deployments.
if raw.shape[1] < 26:
raise ValueError(f"predict_action has {raw.shape[1]} dims; expected at least 26")
pos = raw[:, 16:19]
rot6d = raw[:, 19:25].reshape(-1, 2, 3)
first = rot6d[:, 0]
second = rot6d[:, 1]
first = first / np.maximum(np.linalg.norm(first, axis=1, keepdims=True), 1e-12)
second = second - np.sum(first * second, axis=1, keepdims=True) * first
second = second / np.maximum(np.linalg.norm(second, axis=1, keepdims=True), 1e-12)
third = np.cross(first, second)
matrices = np.stack((first, second, third), axis=-1)
rot = Rotation.from_matrix(matrices).as_euler("zyx")
if current_right is not None:
current = np.asarray(current_right, dtype=np.float64).reshape(7)
if not np.isfinite(current).all():
raise ValueError("state.follow2_pos contains NaN/Inf")
# Training keys are explicitly *_relative: compose the predicted
# translation and rotation with the robot's current EE pose.
pos = current[None, :3] + pos
from wall_x._vendor.x2robot_utils import geometry as geom
rot = geom.compose_state_and_delta_to_abs_rpy(
raw[:, 19:25], current[3:6]
)
grip = raw[:, 25:26]
right = np.concatenate((pos, rot, grip), axis=1)
return _as_action_chunk(right, "predict_action.right_arm")
def _validate_optional_series(ws_action, key, length, width=None):
"""Validate an optional serialized trajectory and return JSON-safe lists."""
if key not in ws_action:
return None
value = np.asarray(ws_action[key], dtype=np.float64)
if value.ndim == 1:
value = value[:, None]
if value.ndim != 2 or (width is not None and value.shape[1] != width):
raise ValueError(f"{key} has invalid shape {value.shape}")
if value.shape[0] != length or not np.isfinite(value).all():
raise ValueError(f"{key} has invalid length or non-finite values")
return value.tolist()
def _legacy_resample(values, key, end_ratio, interpolate_multiplier):
"""Apply the legacy infer.py trim + linear interpolation policy."""
arr = np.asarray(values, dtype=np.float64)
if arr.ndim == 1:
arr = arr[:, None]
if arr.ndim != 2 or arr.shape[0] == 0:
raise ValueError(f"{key} has invalid shape {arr.shape}")
end = int(end_ratio * arr.shape[0])
if end <= 0:
raise ValueError(
f"{key} trim is empty: end_ratio={end_ratio} length={arr.shape[0]}"
)
trimmed = arr[:end]
target_length = int(interpolate_multiplier * len(trimmed))
if target_length <= 0:
raise ValueError(f"{key} interpolation produced no frames")
if target_length == len(trimmed):
return trimmed.tolist()
source_idx = np.linspace(0, len(trimmed) - 1, len(trimmed))
target_idx = np.linspace(0, len(trimmed) - 1, target_length)
out = np.empty((target_length, arr.shape[1]), dtype=np.float64)
for col in range(arr.shape[1]):
out[:, col] = np.interp(target_idx, source_idx, trimmed[:, col])
return out.tolist()
def _limit_action_packet(traj, current, position_limit, rotation_limit):
"""Scale one absolute action packet without changing its path shape.
The limits apply to both the packet's excursion from the observed pose and
its largest inter-sample step. A single scale factor per position/rotation
group keeps the model's direction and timing intact; independently clipping
each sample would turn a large target into an artificial diagonal ramp.
"""
limited = np.asarray(traj, dtype=np.float64).copy()
current = np.asarray(current, dtype=np.float64).reshape(7)
if limited.ndim != 2 or limited.shape[1] != 7:
raise ValueError(f"right-arm trajectory must have shape (T,7), got {limited.shape}")
if not np.isfinite(limited).all() or not np.isfinite(current).all():
raise ValueError("right-arm trajectory/current contains NaN/Inf")
if position_limit <= 0 or rotation_limit <= 0:
raise ValueError("action delta limits must be positive")
scales = {}
for name, columns, limit in (
("position", slice(0, 3), float(position_limit)),
("rotation", slice(3, 6), float(rotation_limit)),
):
delta = limited[:, columns] - current[columns]
peak = float(np.max(np.abs(delta))) if delta.size else 0.0
if len(delta) > 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("<L", len(payload)))
writer.write(payload)
await writer.drain()
previous_right_plan = {
"initial": np.asarray(state["follow2_pos"], dtype=np.float64).reshape(7),
"target": np.asarray(actions["follow2_pos"][-1], dtype=np.float64).reshape(7),
}
logger.info("sent action back to robot T=%d", len(actions["follow1_pos"]))
cycles += 1
if max_action_cycles > 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)
+872
View File
@@ -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("<L", size_b)[0]
return await _recvall(reader, size)
async def _recv_state(reader):
raw = await _recv_frame(reader)
if raw is None:
return None
state = json.loads(raw.decode("utf-8"))
if not isinstance(state, dict):
raise ValueError("robot state is not a JSON object")
for key in ("follow1_pos", "follow2_pos"):
if key not in state:
raise ValueError(f"robot state missing required key {key!r}")
return state
async def _recv_image(reader, index):
raw = await _recv_frame(reader)
if raw is None:
raise ConnectionError("robot closed during image stream")
arr = np.frombuffer(raw, np.uint8)
img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
if img is None:
raise ValueError(f"failed to decode image #{index}")
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
def _img_to_b64(rgb):
ok, buf = cv2.imencode(".jpg", cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR))
if not ok:
raise ValueError("failed to JPEG-encode camera image")
return base64.b64encode(buf.tobytes()).decode("ascii")
def _as_action_chunk(value, key):
try:
arr = np.asarray(value, dtype=np.float64)
except (TypeError, ValueError) as exc:
raise ValueError(f"{key} must be numeric") from exc
if arr.ndim != 2 or arr.shape[1] != 7:
raise ValueError(f"{key} must have shape (T,7), got {arr.shape}")
if not np.isfinite(arr).all():
raise ValueError(f"{key} contains NaN/Inf")
return arr
def _single_arm_action(ws_action, current_right=None):
"""Decode a right-arm-only raw model response into a (T, 7) chunk."""
if "predict_action" not in ws_action:
raise ValueError("WebSocket response missing predict_action")
raw = np.asarray(ws_action["predict_action"], dtype=np.float64)
if raw.ndim == 3:
if raw.shape[0] != 1:
raise ValueError(f"predict_action batch must have size 1, got {raw.shape}")
raw = raw[0]
if raw.ndim != 2 or not np.isfinite(raw).all():
raise ValueError(f"predict_action must be finite with shape (T,D), got {raw.shape}")
# Training layout: 16 padding values, then position(3), rotation-6D(6),
# and gripper(1). Convert the rotation representation without importing
# the LIBERO/MuJoCo package, which is unavailable on headless deployments.
if raw.shape[1] < 26:
raise ValueError(f"predict_action has {raw.shape[1]} dims; expected at least 26")
pos = raw[:, 16:19]
rot6d = raw[:, 19:25].reshape(-1, 2, 3)
first = rot6d[:, 0]
second = rot6d[:, 1]
first = first / np.maximum(np.linalg.norm(first, axis=1, keepdims=True), 1e-12)
second = second - np.sum(first * second, axis=1, keepdims=True) * first
second = second / np.maximum(np.linalg.norm(second, axis=1, keepdims=True), 1e-12)
third = np.cross(first, second)
matrices = np.stack((first, second, third), axis=-1)
rot = Rotation.from_matrix(matrices).as_euler("zyx")
if current_right is not None:
current = np.asarray(current_right, dtype=np.float64).reshape(7)
if not np.isfinite(current).all():
raise ValueError("state.follow2_pos contains NaN/Inf")
# Training keys are explicitly *_relative: compose the predicted
# translation and rotation with the robot's current EE pose.
pos = current[None, :3] + pos
from wall_x._vendor.x2robot_utils import geometry as geom
rot = geom.compose_state_and_delta_to_abs_rpy(
raw[:, 19:25], current[3:6]
)
grip = raw[:, 25:26]
right = np.concatenate((pos, rot, grip), axis=1)
return _as_action_chunk(right, "predict_action.right_arm")
def _validate_optional_series(ws_action, key, length, width=None):
"""Validate an optional serialized trajectory and return JSON-safe lists."""
if key not in ws_action:
return None
value = np.asarray(ws_action[key], dtype=np.float64)
if value.ndim == 1:
value = value[:, None]
if value.ndim != 2 or (width is not None and value.shape[1] != width):
raise ValueError(f"{key} has invalid shape {value.shape}")
if value.shape[0] != length or not np.isfinite(value).all():
raise ValueError(f"{key} has invalid length or non-finite values")
return value.tolist()
def _legacy_resample(values, key, end_ratio, interpolate_multiplier):
"""Apply the legacy infer.py trim + linear interpolation policy."""
arr = np.asarray(values, dtype=np.float64)
if arr.ndim == 1:
arr = arr[:, None]
if arr.ndim != 2 or arr.shape[0] == 0:
raise ValueError(f"{key} has invalid shape {arr.shape}")
end = int(end_ratio * arr.shape[0])
if end <= 0:
raise ValueError(
f"{key} trim is empty: end_ratio={end_ratio} length={arr.shape[0]}"
)
trimmed = arr[:end]
target_length = int(interpolate_multiplier * len(trimmed))
if target_length <= 0:
raise ValueError(f"{key} interpolation produced no frames")
if target_length == len(trimmed):
return trimmed.tolist()
source_idx = np.linspace(0, len(trimmed) - 1, len(trimmed))
target_idx = np.linspace(0, len(trimmed) - 1, target_length)
out = np.empty((target_length, arr.shape[1]), dtype=np.float64)
for col in range(arr.shape[1]):
out[:, col] = np.interp(target_idx, source_idx, trimmed[:, col])
return out.tolist()
def _limit_action_packet(traj, current, position_limit, rotation_limit):
"""Scale one absolute action packet without changing its path shape.
The limits apply to both the packet's excursion from the observed pose and
its largest inter-sample step. A single scale factor per position/rotation
group keeps the model's direction and timing intact; independently clipping
each sample would turn a large target into an artificial diagonal ramp.
"""
limited = np.asarray(traj, dtype=np.float64).copy()
current = np.asarray(current, dtype=np.float64).reshape(7)
if limited.ndim != 2 or limited.shape[1] != 7:
raise ValueError(f"right-arm trajectory must have shape (T,7), got {limited.shape}")
if not np.isfinite(limited).all() or not np.isfinite(current).all():
raise ValueError("right-arm trajectory/current contains NaN/Inf")
if position_limit <= 0 or rotation_limit <= 0:
raise ValueError("action delta limits must be positive")
scales = {}
for name, columns, limit in (
("position", slice(0, 3), float(position_limit)),
("rotation", slice(3, 6), float(rotation_limit)),
):
delta = limited[:, columns] - current[columns]
peak = float(np.max(np.abs(delta))) if delta.size else 0.0
if len(delta) > 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("<L", len(payload)))
writer.write(payload)
await writer.drain()
previous_right_plan = {
"initial": np.asarray(state["follow2_pos"], dtype=np.float64).reshape(7),
"target": np.asarray(actions["follow2_pos"][-1], dtype=np.float64).reshape(7),
}
logger.info("sent action back to robot T=%d", len(actions["follow1_pos"]))
cycles += 1
if max_action_cycles > 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)
+2
View File
@@ -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",
+45
View File
@@ -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)
+29
View File
@@ -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)
+149
View File
@@ -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])
+40
View File
@@ -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
+137
View File
@@ -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,
)
+47
View File
@@ -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
+3 -3
View File
@@ -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",
}
@@ -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=(
@@ -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
@@ -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
@@ -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))
@@ -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
@@ -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"}
+119
View File
@@ -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()
+96 -6
View File
@@ -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
+30 -5
View File
@@ -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)
+1
View File
@@ -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"
+22
View File
@@ -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,
+27
View File
@@ -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__)
+15 -2
View File
@@ -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
+18 -1
View File
@@ -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
@@ -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()
@@ -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
File diff suppressed because it is too large Load Diff
+4
View File
@@ -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):
+42 -5
View File
@@ -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
+12 -6
View File
@@ -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()
+2
View File
@@ -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: