Add Wall-X serving and Turtle2 TCP WebSocket bridge
Pre-commit / pre-commit (push) Canceled after 0s
Pre-commit / pre-commit (push) Canceled after 0s
This commit is contained in:
@@ -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/
|
||||
|
||||
@@ -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 仍需要现场与训练产物确认。
|
||||
@@ -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`,建议先停止机器人端,再停止桥接端和推流端。
|
||||
@@ -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
|
||||
@@ -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
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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}"
|
||||
|
||||
|
||||
Executable
+342
@@ -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[@]}"
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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])
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
@@ -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"}
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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__)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user