From 9a10812fd94acb80b67b13c3adf0094b90f7d4e4 Mon Sep 17 00:00:00 2001 From: yangping Date: Wed, 10 Jun 2026 11:02:52 +0800 Subject: [PATCH 1/3] refactor: restructure codebase with harrix vendor, config system, and FSDP trainer Integrate harrix for serving/evaluation, introduce modular config and data backends, reorganize model/trainer layers, and update workspace examples. Co-authored-by: Cursor --- scripts/README.md | 70 + scripts/compute_norm_stats.py | 698 ++++- scripts/draw_openloop_plot.py | 1141 +++++++- scripts/fake_inference.py | 188 +- scripts/infer_libero.py | 413 +-- scripts/infer_robochallenge.py | 1164 -------- scripts/merge_sharded_weights.py | 69 +- scripts/normalize.py | 4 +- scripts/run_libero.sh | 204 ++ scripts/run_serving.sh | 16 + scripts/vqa_inference.py | 2 +- wall_x/_vendor/__init__.py | 4 + wall_x/_vendor/harrix/__init__.py | 5 + wall_x/_vendor/harrix/adapters/__init__.py | 6 + wall_x/_vendor/harrix/adapters/base.py | 51 + wall_x/_vendor/harrix/adapters/qwen_vlact.py | 465 ++++ wall_x/_vendor/harrix/adapters/registry.py | 39 + .../harrix/adapters/variants/__init__.py | 15 + .../harrix/adapters/variants/qwen2_5.py | 17 + .../harrix/drivers}/__init__.py | 0 .../_vendor/harrix/drivers/inproc/__init__.py | 197 ++ .../harrix/drivers/inproc/env_handle.py | 152 ++ .../harrix/drivers/inproc/model_handle.py | 25 + wall_x/_vendor/harrix/drivers/job_state.py | 343 +++ wall_x/_vendor/harrix/envs/__init__.py | 8 + wall_x/_vendor/harrix/envs/base.py | 119 + wall_x/_vendor/harrix/envs/libero.py | 342 +++ wall_x/_vendor/harrix/envs/libero_common.py | 359 +++ wall_x/_vendor/harrix/envs/libero_sim.py | 181 ++ wall_x/_vendor/harrix/envs/registry.py | 53 + wall_x/_vendor/harrix/eval_config.py | 217 ++ .../harrix/serving-old}/__init__.py | 0 .../serving-old/_wallx_infer/__init__.py | 1 + .../_wallx_infer}/base_dataclass.py | 183 +- .../serving-old/_wallx_infer/infer_config.py | 343 +++ .../serving-old/_wallx_infer}/logger.py | 108 +- .../serving-old/_wallx_infer/model_wrapper.py | 827 ++++++ .../harrix/serving-old/_wallx_infer/robot.py | 964 +++++++ .../_wallx_infer/socket_controller.py | 423 +++ .../harrix/serving-old/_wallx_infer}/utils.py | 83 +- wall_x/_vendor/harrix/serving-old/client.py | 260 ++ .../harrix/serving-old/launch_serving.py | 268 ++ .../harrix/serving-old}/policy/__init__.py | 0 .../serving-old/policy/wall_x_policy.py | 250 ++ .../_vendor/harrix/serving-old/scheduler.py | 160 ++ .../serving-old/websocket_policy_server.py | 184 ++ wall_x/_vendor/harrix/serving/__init__.py | 3 + .../harrix/serving/_wallx_infer/__init__.py | 1 + .../serving/_wallx_infer/base_dataclass.py | 476 ++++ .../serving/_wallx_infer/infer_config.py | 434 +++ .../harrix/serving/_wallx_infer/logger.py | 277 ++ .../serving/_wallx_infer/model_wrapper.py | 764 ++++++ .../harrix/serving/_wallx_infer/robot.py | 1183 ++++++++ .../serving/_wallx_infer/socket_controller.py | 423 +++ .../harrix/serving/_wallx_infer/utils.py | 280 ++ wall_x/_vendor/harrix/serving/client.py | 262 ++ .../_vendor/harrix/serving/launch_serving.py | 281 ++ .../_vendor/harrix/serving/policy/__init__.py | 5 + .../harrix/serving/policy/_smoothing.py | 104 + .../harrix/serving/policy/wall_x_policy.py | 328 +++ wall_x/_vendor/harrix/serving/scheduler.py | 156 ++ .../serving/websocket_policy_server.py | 82 +- wall_x/_vendor/harrix/utils/__init__.py | 0 wall_x/_vendor/harrix/utils/ckpt_load.py | 171 ++ wall_x/_vendor/harrix/utils/normalizer.py | 135 + wall_x/_vendor/harrix/utils/seed.py | 23 + wall_x/_vendor/harrix/utils/train_config.py | 371 +++ wall_x/_vendor/x2robot_utils/__init__.py | 5 + wall_x/_vendor/x2robot_utils/geometry.py | 309 +++ wall_x/_vendor/x2robot_utils/grounding.py | 172 ++ .../_vendor/x2robot_utils/text_templates.py | 235 ++ wall_x/config/__init__.py | 62 + wall_x/config/data_config.py | 63 + wall_x/config/hyperparams_config.py | 129 + wall_x/config/infra_config.py | 98 + wall_x/config/loader.py | 285 ++ wall_x/config/model_config.py | 42 + wall_x/config/registry.py | 100 + wall_x/config/task_config.py | 17 + wall_x/config/train_config.py | 81 + wall_x/data/__init__.py | 31 + wall_x/data/_bundle.py | 39 + wall_x/data/_protocol.py | 37 + wall_x/data/_registry.py | 263 ++ wall_x/data/backends/__init__.py | 16 + wall_x/data/backends/lerobot/__init__.py | 26 + wall_x/data/backends/lerobot/build.py | 337 +++ wall_x/data/backends/lerobot/config.py | 133 + wall_x/data/backends/lerobot/loader.py | 886 ++++++ .../data/backends/lerobot/rotation_layout.py | 95 + wall_x/data/{ => backends/lerobot}/utils.py | 937 +++---- wall_x/data/backends/mm | 188 ++ wall_x/data/config.py | 126 - wall_x/data/load_lerobot_dataset.py | 715 ----- wall_x/fusions/backend.py | 433 --- wall_x/fusions/ops.py | 742 ----- wall_x/infer/env.py | 143 - wall_x/infer/env_libero.py | 744 ----- wall_x/infer/infer_config.py | 587 ---- wall_x/infer/utils_libero.py | 323 --- wall_x/model/action_head.py | 808 ------ wall_x/model/core/__init__.py | 1 + wall_x/model/core/action/__init__.py | 10 + wall_x/model/core/action/head.py | 27 + wall_x/model/core/action/moe.py | 127 + wall_x/model/core/action/normalizer.py | 447 +++ wall_x/model/core/action/processor.py | 317 +++ wall_x/model/core/attention/__init__.py | 9 + .../attention/joint.py} | 505 +++- wall_x/model/core/attention/mask.py | 259 ++ wall_x/model/core/attention/selector.py | 138 + wall_x/model/core/ops/__init__.py | 17 + wall_x/model/core/ops/base.py | 99 + wall_x/model/core/ops/index.py | 339 +++ wall_x/model/core/ops/moe.py | 96 + wall_x/model/core/ops/norm.py | 41 + wall_x/model/core/ops/rope.py | 233 ++ wall_x/model/core/vla_mixin.py | 751 ++++++ wall_x/model/model_utils.py | 319 --- wall_x/model/qact/__init__.py | 1 + .../qwen2_5}/__init__.py | 6 - wall_x/model/qact/qwen2_5/adapter.py | 56 + .../qwen2_5}/configuration_qwen2_5_vl.py | 80 +- wall_x/model/qact/qwen2_5/inference_mixin.py | 1039 +++++++ .../qwen2_5}/modeling_qwen2_5_vl.py | 628 ++++- .../qwen2_5}/modeling_qwen2_5_vl_act.py | 2391 ++++++++++------- wall_x/model/qact/tokenizer_mixin.py | 963 +++++++ wall_x/model/registry.py | 64 + wall_x/model/vla_mixin.py | 987 ------- wall_x/serving/README.md | 263 -- wall_x/serving/client.py | 397 --- wall_x/serving/launch_serving.py | 220 -- wall_x/serving/policy/utils.py | 265 -- wall_x/serving/policy/wall_x_policy.py | 182 -- wall_x/trainer/adapters/__init__.py | 129 + wall_x/trainer/adapters/base_adapter.py | 495 ++++ wall_x/trainer/adapters/vla_model_adapter.py | 990 +++++++ wall_x/trainer/fsdp_trainer/__init__.py | 36 + wall_x/trainer/fsdp_trainer/base_trainer.py | 429 +++ wall_x/trainer/fsdp_trainer/checkpoint_io.py | 1027 +++++++ .../fsdp_trainer/distribution_strategy.py | 400 +++ wall_x/trainer/fsdp_trainer/fsdp_trainer.py | 963 +++++++ .../trainer/fsdp_trainer/metrics/__init__.py | 3 + wall_x/trainer/fsdp_trainer/metrics/norms.py | 123 + wall_x/trainer/fsdp_trainer/metrics_logger.py | 111 + wall_x/trainer/fsdp_trainer/train_fsdp.py | 179 ++ wall_x/trainer/optimizer/__init__.py | 9 + wall_x/trainer/optimizer/dmuon/__init__.py | 16 + wall_x/trainer/optimizer/dmuon/utils.py | 156 ++ wall_x/trainer/optimizer/utils.py | 230 ++ wall_x/trainer/qwen_vl_act_trainer.py | 1144 -------- wall_x/trainer/scheduler/scheduler.py | 70 + wall_x/trainer/trainer_utils.py | 574 ++++ wall_x/trainer/utils/__init__.py | 6 + wall_x/trainer/utils/data.py | 46 + wall_x/trainer/utils/diagnostics.py | 40 + wall_x/utils/constant.py | 421 +-- wall_x/utils/cudagraph_wrapper.py | 277 ++ wall_x/utils/logger.py | 95 + wall_x/utils/metrics.py | 94 + wall_x/utils/timers.py | 209 +- wall_x/utils/weight_convert/__init__.py | 0 wall_x/utils/weight_convert/convert.py | 215 ++ wall_x/utils/weight_convert/weight_loader.py | 377 +++ workspace/README.md | 421 ++- workspace/example/libero.yml | 125 + workspace/example/maniparena_example.yml | 134 + workspace/lerobot_example/config_qact.yml | 151 -- .../lerobot_example/config_qact_from_vlm.yml | 152 -- .../evaluation/lerobot_openloop.png | 3 - .../libero/config_qact_libero_from_vlm.yml | 100 - workspace/lerobot_example/qwen25_config.json | 106 - workspace/lerobot_example/run.sh | 23 - 173 files changed, 33676 insertions(+), 13267 deletions(-) create mode 100644 scripts/README.md delete mode 100644 scripts/infer_robochallenge.py create mode 100755 scripts/run_libero.sh create mode 100755 scripts/run_serving.sh create mode 100644 wall_x/_vendor/__init__.py create mode 100644 wall_x/_vendor/harrix/__init__.py create mode 100644 wall_x/_vendor/harrix/adapters/__init__.py create mode 100644 wall_x/_vendor/harrix/adapters/base.py create mode 100644 wall_x/_vendor/harrix/adapters/qwen_vlact.py create mode 100644 wall_x/_vendor/harrix/adapters/registry.py create mode 100644 wall_x/_vendor/harrix/adapters/variants/__init__.py create mode 100644 wall_x/_vendor/harrix/adapters/variants/qwen2_5.py rename wall_x/{fusions => _vendor/harrix/drivers}/__init__.py (100%) create mode 100644 wall_x/_vendor/harrix/drivers/inproc/__init__.py create mode 100644 wall_x/_vendor/harrix/drivers/inproc/env_handle.py create mode 100644 wall_x/_vendor/harrix/drivers/inproc/model_handle.py create mode 100644 wall_x/_vendor/harrix/drivers/job_state.py create mode 100644 wall_x/_vendor/harrix/envs/__init__.py create mode 100644 wall_x/_vendor/harrix/envs/base.py create mode 100644 wall_x/_vendor/harrix/envs/libero.py create mode 100644 wall_x/_vendor/harrix/envs/libero_common.py create mode 100644 wall_x/_vendor/harrix/envs/libero_sim.py create mode 100644 wall_x/_vendor/harrix/envs/registry.py create mode 100644 wall_x/_vendor/harrix/eval_config.py rename wall_x/{serving => _vendor/harrix/serving-old}/__init__.py (100%) create mode 100644 wall_x/_vendor/harrix/serving-old/_wallx_infer/__init__.py rename wall_x/{infer => _vendor/harrix/serving-old/_wallx_infer}/base_dataclass.py (63%) create mode 100644 wall_x/_vendor/harrix/serving-old/_wallx_infer/infer_config.py rename wall_x/{infer => _vendor/harrix/serving-old/_wallx_infer}/logger.py (71%) create mode 100644 wall_x/_vendor/harrix/serving-old/_wallx_infer/model_wrapper.py create mode 100644 wall_x/_vendor/harrix/serving-old/_wallx_infer/robot.py create mode 100644 wall_x/_vendor/harrix/serving-old/_wallx_infer/socket_controller.py rename wall_x/{infer => _vendor/harrix/serving-old/_wallx_infer}/utils.py (76%) create mode 100644 wall_x/_vendor/harrix/serving-old/client.py create mode 100644 wall_x/_vendor/harrix/serving-old/launch_serving.py rename wall_x/{serving => _vendor/harrix/serving-old}/policy/__init__.py (100%) create mode 100644 wall_x/_vendor/harrix/serving-old/policy/wall_x_policy.py create mode 100644 wall_x/_vendor/harrix/serving-old/scheduler.py create mode 100644 wall_x/_vendor/harrix/serving-old/websocket_policy_server.py create mode 100644 wall_x/_vendor/harrix/serving/__init__.py create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/__init__.py create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/base_dataclass.py create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/logger.py create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/robot.py create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py create mode 100644 wall_x/_vendor/harrix/serving/_wallx_infer/utils.py create mode 100644 wall_x/_vendor/harrix/serving/client.py create mode 100644 wall_x/_vendor/harrix/serving/launch_serving.py create mode 100644 wall_x/_vendor/harrix/serving/policy/__init__.py create mode 100644 wall_x/_vendor/harrix/serving/policy/_smoothing.py create mode 100644 wall_x/_vendor/harrix/serving/policy/wall_x_policy.py create mode 100644 wall_x/_vendor/harrix/serving/scheduler.py rename wall_x/{ => _vendor/harrix}/serving/websocket_policy_server.py (57%) create mode 100644 wall_x/_vendor/harrix/utils/__init__.py create mode 100644 wall_x/_vendor/harrix/utils/ckpt_load.py create mode 100644 wall_x/_vendor/harrix/utils/normalizer.py create mode 100644 wall_x/_vendor/harrix/utils/seed.py create mode 100644 wall_x/_vendor/harrix/utils/train_config.py create mode 100644 wall_x/_vendor/x2robot_utils/__init__.py create mode 100644 wall_x/_vendor/x2robot_utils/geometry.py create mode 100644 wall_x/_vendor/x2robot_utils/grounding.py create mode 100644 wall_x/_vendor/x2robot_utils/text_templates.py create mode 100644 wall_x/config/__init__.py create mode 100644 wall_x/config/data_config.py create mode 100644 wall_x/config/hyperparams_config.py create mode 100644 wall_x/config/infra_config.py create mode 100644 wall_x/config/loader.py create mode 100644 wall_x/config/model_config.py create mode 100644 wall_x/config/registry.py create mode 100644 wall_x/config/task_config.py create mode 100644 wall_x/config/train_config.py create mode 100644 wall_x/data/_bundle.py create mode 100644 wall_x/data/_protocol.py create mode 100644 wall_x/data/_registry.py create mode 100644 wall_x/data/backends/__init__.py create mode 100644 wall_x/data/backends/lerobot/__init__.py create mode 100644 wall_x/data/backends/lerobot/build.py create mode 100644 wall_x/data/backends/lerobot/config.py create mode 100644 wall_x/data/backends/lerobot/loader.py create mode 100644 wall_x/data/backends/lerobot/rotation_layout.py rename wall_x/data/{ => backends/lerobot}/utils.py (72%) create mode 100644 wall_x/data/backends/mm delete mode 100644 wall_x/data/config.py delete mode 100644 wall_x/data/load_lerobot_dataset.py delete mode 100644 wall_x/fusions/backend.py delete mode 100644 wall_x/fusions/ops.py delete mode 100644 wall_x/infer/env.py delete mode 100644 wall_x/infer/env_libero.py delete mode 100644 wall_x/infer/infer_config.py delete mode 100644 wall_x/infer/utils_libero.py delete mode 100644 wall_x/model/action_head.py create mode 100644 wall_x/model/core/__init__.py create mode 100644 wall_x/model/core/action/__init__.py create mode 100644 wall_x/model/core/action/head.py create mode 100644 wall_x/model/core/action/moe.py create mode 100644 wall_x/model/core/action/normalizer.py create mode 100644 wall_x/model/core/action/processor.py create mode 100644 wall_x/model/core/attention/__init__.py rename wall_x/model/{joint_attention.py => core/attention/joint.py} (56%) create mode 100644 wall_x/model/core/attention/mask.py create mode 100644 wall_x/model/core/attention/selector.py create mode 100644 wall_x/model/core/ops/__init__.py create mode 100644 wall_x/model/core/ops/base.py create mode 100644 wall_x/model/core/ops/index.py create mode 100644 wall_x/model/core/ops/moe.py create mode 100644 wall_x/model/core/ops/norm.py create mode 100644 wall_x/model/core/ops/rope.py create mode 100644 wall_x/model/core/vla_mixin.py delete mode 100644 wall_x/model/model_utils.py create mode 100644 wall_x/model/qact/__init__.py rename wall_x/model/{qwen2_5_based => qact/qwen2_5}/__init__.py (58%) create mode 100644 wall_x/model/qact/qwen2_5/adapter.py rename wall_x/model/{qwen2_5_based => qact/qwen2_5}/configuration_qwen2_5_vl.py (77%) create mode 100644 wall_x/model/qact/qwen2_5/inference_mixin.py rename wall_x/model/{qwen2_5_based => qact/qwen2_5}/modeling_qwen2_5_vl.py (82%) rename wall_x/model/{qwen2_5_based => qact/qwen2_5}/modeling_qwen2_5_vl_act.py (60%) create mode 100644 wall_x/model/qact/tokenizer_mixin.py create mode 100644 wall_x/model/registry.py delete mode 100644 wall_x/model/vla_mixin.py delete mode 100644 wall_x/serving/README.md delete mode 100644 wall_x/serving/client.py delete mode 100644 wall_x/serving/launch_serving.py delete mode 100644 wall_x/serving/policy/utils.py delete mode 100644 wall_x/serving/policy/wall_x_policy.py create mode 100644 wall_x/trainer/adapters/__init__.py create mode 100644 wall_x/trainer/adapters/base_adapter.py create mode 100644 wall_x/trainer/adapters/vla_model_adapter.py create mode 100644 wall_x/trainer/fsdp_trainer/__init__.py create mode 100644 wall_x/trainer/fsdp_trainer/base_trainer.py create mode 100644 wall_x/trainer/fsdp_trainer/checkpoint_io.py create mode 100644 wall_x/trainer/fsdp_trainer/distribution_strategy.py create mode 100644 wall_x/trainer/fsdp_trainer/fsdp_trainer.py create mode 100644 wall_x/trainer/fsdp_trainer/metrics/__init__.py create mode 100644 wall_x/trainer/fsdp_trainer/metrics/norms.py create mode 100644 wall_x/trainer/fsdp_trainer/metrics_logger.py create mode 100644 wall_x/trainer/fsdp_trainer/train_fsdp.py create mode 100644 wall_x/trainer/optimizer/__init__.py create mode 100644 wall_x/trainer/optimizer/dmuon/__init__.py create mode 100644 wall_x/trainer/optimizer/dmuon/utils.py create mode 100644 wall_x/trainer/optimizer/utils.py delete mode 100644 wall_x/trainer/qwen_vl_act_trainer.py create mode 100644 wall_x/trainer/scheduler/scheduler.py create mode 100644 wall_x/trainer/trainer_utils.py create mode 100644 wall_x/trainer/utils/__init__.py create mode 100644 wall_x/trainer/utils/data.py create mode 100644 wall_x/trainer/utils/diagnostics.py create mode 100644 wall_x/utils/cudagraph_wrapper.py create mode 100644 wall_x/utils/logger.py create mode 100644 wall_x/utils/metrics.py create mode 100644 wall_x/utils/weight_convert/__init__.py create mode 100644 wall_x/utils/weight_convert/convert.py create mode 100644 wall_x/utils/weight_convert/weight_loader.py create mode 100644 workspace/example/libero.yml create mode 100644 workspace/example/maniparena_example.yml delete mode 100644 workspace/lerobot_example/config_qact.yml delete mode 100644 workspace/lerobot_example/config_qact_from_vlm.yml delete mode 100644 workspace/lerobot_example/evaluation/lerobot_openloop.png delete mode 100644 workspace/lerobot_example/libero/config_qact_libero_from_vlm.yml delete mode 100644 workspace/lerobot_example/qwen25_config.json delete mode 100644 workspace/lerobot_example/run.sh diff --git a/scripts/README.md b/scripts/README.md new file mode 100644 index 00000000..13ab106c --- /dev/null +++ b/scripts/README.md @@ -0,0 +1,70 @@ +# Scripts + +This directory contains the public Wall-X command-line helpers. Run commands +from the repository root unless noted otherwise. + +## Inference smoke test + +Use `fake_inference.py` to verify that a checkpoint can be loaded and can +produce one action chunk from a synthetic LIBERO-style observation. + +```bash +python scripts/fake_inference.py \ + --checkpoint-path /path/to/checkpoint +``` + +If the training config is not stored next to the checkpoint as `config.yml` or +`config.yaml`, pass it explicitly: + +```bash +python scripts/fake_inference.py \ + --checkpoint-path /path/to/checkpoint \ + --train-config-path /path/to/config.yml +``` + +## LIBERO evaluation + +`run_libero.sh` is a small shell wrapper around `infer_libero.py`. + +```bash +CHECKPOINT_PATH=/path/to/checkpoint bash scripts/run_libero.sh +``` + +Useful environment variables: + +```bash +CHECKPOINT_PATH=/path/to/checkpoint +TRAIN_CONFIG_PATH=/path/to/config.yml +TASK_SUITE_NAME=libero_spatial +TASK_INDICES=0,1,2 +NUM_TRIALS_PER_TASK=50 +CUDA_ID=0 +SMOKE=1 +``` + +For full control, call the Python entry directly: + +```bash +python scripts/infer_libero.py \ + --checkpoint-path /path/to/checkpoint \ + --task-suite-name libero_spatial \ + --num-trials-per-task 50 \ + --driver-mode in_process +``` + +You can also pass a complete harrix eval config: + +```bash +python scripts/infer_libero.py --config /path/to/eval_config.yml +``` + +## Dataset and checkpoint utilities + +- `compute_norm_stats.py`: compute action normalization statistics for a + LeRobot dataset. +- `merge_sharded_weights.py`: merge FSDP sharded checkpoint files into a single + checkpoint directory. +- `normalize.py`: helper for generating normalization files from dataset + statistics. + +Most scripts support `--help` for their command-line options. diff --git a/scripts/compute_norm_stats.py b/scripts/compute_norm_stats.py index a820d959..f1ce6dd3 100644 --- a/scripts/compute_norm_stats.py +++ b/scripts/compute_norm_stats.py @@ -1,182 +1,594 @@ #!/usr/bin/env python3 +"""Compute LeRobot normalization stats (mean, std, q01, q99) for training. +Writes JSON in the format expected by wall-x training configs:: + + {"norm_stats": { + "observation.state": {"mean": [...], "std": [...], "q01": [...], "q99": [...]}, + "action": {"mean": [...], "std": [...], "q01": [...], "q99": [...]} + }} + +When ``--train_config`` is provided, the script reads ``data.lerobot_config.repo_id``, +``norm_stats_path``, ``task.dof_config``, ``task.agent_pos_config``, and +``task.action_horizon`` from the YAML. Per-DOF slices are aggregated separately; +keys ending with ``_relative`` use the same relative-pose logic as the LeRobot loader. + +Usage +----- +Recommended: pass a finetune YAML (paths in the config can be placeholders; override +with CLI flags if needed):: + + python scripts/compute_norm_stats.py \\ + --train_config workspace/example/fintune_bus2602/libero.yml + +CVPR / multi-task example:: + + python scripts/compute_norm_stats.py \\ + --train_config workspace/example/fintune_bus2602/cvpr_example.yml + +Override dataset or output path from the command line:: + + python scripts/compute_norm_stats.py \\ + --train_config workspace/example/fintune_bus2602/libero.yml \\ + --data_root /path/to/repo_id \\ + --output_path /path/to/norm_stats_path + +Without a train config (global stats only, no per-DOF relative slices):: + + python scripts/compute_norm_stats.py \\ + --data_root /path/to/lerobot_dataset \\ + --output_path /path/to/norm_stats.json + +Requirements +------------ +- Local LeRobot v3 dataset at ``--data_root`` (or ``data.lerobot_config.repo_id``) +- ``lerobot>=0.3``, ``numpy``, ``pyyaml``, ``tqdm`` +- Run from the wall-x repo root (or set ``PYTHONPATH`` so ``wall_x`` imports resolve) + +After running, set ``norm_stats_path`` in your training YAML to the generated JSON. +""" + +import argparse import json import logging -from collections import defaultdict +from dataclasses import dataclass from pathlib import Path -from typing import Dict, List -from tqdm import tqdm +from typing import Any import numpy as np - +import yaml from lerobot.datasets.lerobot_dataset import LeRobotDataset +from numba import jit, prange +from tqdm import tqdm +from wall_x._vendor.x2robot_utils.geometry import ( + canonicalize_euler_zyx_batch_nb, + euler_to_matrix_zyx_batch_nb, + matrix_to_euler_zyx_batch_nb, + so3_to_matrix_batch_nb, +) -def write_json(path: Path, data: Dict) -> None: - path.write_text( - json.dumps(data, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) +SKIP_DOF_KEYS = frozenset( + {"velocity_decomposed", "height", "head_actions", "action_padding"} +) -def compute_action_statistics( - action_data_by_robot: Dict[str, Dict[str, List]] -) -> Dict[str, Dict[str, Dict]]: - """ - Compute statistics (min, q01, q99, max) for each action type and dimension. +@dataclass +class TrainConfigContext: + data_root: Path + output_path: Path + state_key: str + action_key: str + propri_ranges: dict[str, list[int]] + action_ranges: dict[str, list[int]] + action_chunk: int - Args: - action_data_by_robot: Dict[robot_id][action_type] -> list of arrays/lists - Returns: - Dict[robot_id][action_type] -> { - "min": [min for each dim], - "q01": [quantile 1% for each dim], - "q99": [quantile 99% for each dim], - "max": [max for each dim], - "delta": [max - min for each dim] - "delta_q99_q01": [q99 - q01 for each dim] - } - """ - stats = {} - - for robot_id, action_data in action_data_by_robot.items(): - stats[robot_id] = {} - - for action_type, values_list in action_data.items(): - if not values_list: - continue - - # Convert to numpy array: shape (num_samples, num_dims) - try: - values_array = np.array(values_list) - if values_array.size == 0: - continue - - # Handle both 1D and 2D cases - if values_array.ndim == 1: - values_array = values_array.reshape(-1, 1) - elif values_array.ndim == 2: - pass - else: - logging.warning( - f"Unexpected shape for {robot_id}/{action_type}: {values_array.shape}" - ) - continue - - # Compute statistics for each dimension - min_vals = np.min(values_array, axis=0).tolist() - max_vals = np.max(values_array, axis=0).tolist() - q01_vals = np.quantile(values_array, 0.01, axis=0).tolist() - q99_vals = np.quantile(values_array, 0.99, axis=0).tolist() - delta_vals = (np.array(max_vals) - np.array(min_vals)).tolist() - delta_q99_q01_vals = (np.array(q99_vals) - np.array(q01_vals)).tolist() - - stats[robot_id][action_type] = { - "min": min_vals, - "q01": q01_vals, - "q99": q99_vals, - "max": max_vals, - "delta": delta_vals, - "delta_q99_q01": delta_q99_q01_vals, - } - - except Exception as e: - logging.warning( - f"Error computing statistics for {robot_id}/{action_type}: {e}" - ) - continue +def write_norm_stats(path: Path, norm_stats: dict[str, dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps({"norm_stats": norm_stats}, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + - return stats +def compute_vector_stats(values: np.ndarray) -> dict[str, list[float]]: + if values.ndim == 1: + values = values.reshape(-1, 1) + return { + "mean": np.mean(values, axis=0).tolist(), + "std": np.std(values, axis=0).tolist(), + "q01": np.quantile(values, 0.01, axis=0).tolist(), + "q99": np.quantile(values, 0.99, axis=0).tolist(), + } -def load_lerobot_dataset( - repo_id: str, - trajectory_keys: Dict, - base_dir: Path, +def _apply_slice_stats( + full_stats: dict[str, list[float]], + index_range: list[int], + slice_stats: dict[str, list[float]], ) -> None: + start, end = index_range + for field in ("mean", "std", "q01", "q99"): + full_stats[field][start:end] = slice_stats[field] + + +def config_to_index_ranges(config: dict[str, int]) -> dict[str, list[int]]: + ranges: dict[str, list[int]] = {} + cur = 0 + for key, dim in config.items(): + if key in SKIP_DOF_KEYS: + continue + ranges[key] = [cur, cur + int(dim)] + cur += int(dim) + return ranges + + +def load_train_config(path: Path) -> dict[str, Any]: + with path.open(encoding="utf-8") as f: + config = yaml.load(f, Loader=yaml.SafeLoader) + if not isinstance(config, dict): + raise ValueError(f"train config must be a YAML mapping, got {type(config)}") + return config + + +def parse_train_config(config: dict[str, Any]) -> TrainConfigContext: + task = config.get("task") + if not isinstance(task, dict): + raise ValueError("train config must contain a 'task' section") + + dof_config = task.get("dof_config") + agent_pos_config = task.get("agent_pos_config") + if not isinstance(dof_config, dict) or not dof_config: + raise ValueError("task.dof_config is required in train config") + if not isinstance(agent_pos_config, dict) or not agent_pos_config: + raise ValueError("task.agent_pos_config is required in train config") + + data_cfg = config.get("data") + if not isinstance(data_cfg, dict): + raise ValueError("train config must contain a 'data' section") + + lerobot_config = data_cfg.get("lerobot_config") + if not isinstance(lerobot_config, dict): + raise ValueError("data.lerobot_config is required in train config") + + repo_id = lerobot_config.get("repo_id") + if not repo_id: + raise ValueError("data.lerobot_config.repo_id is required in train config") + + norm_stats_path = config.get("norm_stats_path") or data_cfg.get("norm_stats_path") + if not norm_stats_path: + raise ValueError( + "norm_stats_path is required in train config " "(top-level or under data)" + ) + + key_mappings = data_cfg.get("key_mappings") + if not isinstance(key_mappings, dict): + raise ValueError("data.key_mappings is required in train config") + + state_key = key_mappings.get("state", "observation.state") + action_key = key_mappings.get("action", "action") + + action_chunk = int( + task.get("action_horizon") + or task.get("action_horizon_flow") + or data_cfg.get("action_horizon") + or 32 + ) - # Load local or remote dataset - dataset = LeRobotDataset(base_dir) + return TrainConfigContext( + data_root=Path(repo_id), + output_path=Path(norm_stats_path), + state_key=state_key, + action_key=action_key, + propri_ranges=config_to_index_ranges(agent_pos_config), + action_ranges=config_to_index_ranges(dof_config), + action_chunk=action_chunk, + ) - # Iterate through all data - frames: Dict[str, Dict[str, List]] = defaultdict(lambda: defaultdict(list)) - all_features = dataset.features - non_image_columns = [col for col in all_features if "image" not in col] +def _load_state_action_table( + data_root: Path, + state_key: str, + action_key: str, +): + dataset = LeRobotDataset(str(data_root), root=None, video_backend="pyav") + non_image_columns = [ + col for col in dataset.features if "image" not in col and col not in {"task"} + ] + if state_key not in non_image_columns or action_key not in non_image_columns: + raise ValueError( + f"Expected keys {state_key!r} and {action_key!r} in dataset columns, " + f"got {non_image_columns}" + ) + table = dataset.hf_dataset.select_columns([state_key, action_key]) + return table, state_key, action_key + + +def _table_to_arrays( + table, + state_key: str, + action_key: str, +) -> tuple[np.ndarray, np.ndarray]: + """Load full columns once; avoids O(N*chunk) random row access.""" + logging.info("Loading state/action columns into memory...") + try: + states = np.asarray(table[state_key], dtype=np.float32) + actions = np.asarray(table[action_key], dtype=np.float32) + except (KeyError, TypeError, ValueError) as exc: + logging.warning( + "Column-wise load failed (%s); falling back to per-row stack.", exc + ) + states = np.stack( + [ + np.asarray(table[i][state_key], dtype=np.float32) + for i in range(len(table)) + ] + ) + actions = np.stack( + [ + np.asarray(table[i][action_key], dtype=np.float32) + for i in range(len(table)) + ] + ) + if states.ndim == 1: + states = states.reshape(-1, 1) + if actions.ndim == 1: + actions = actions.reshape(-1, 1) + logging.info( + " frames=%d state_dim=%d action_dim=%d", + len(states), + states.shape[1], + actions.shape[1], + ) + return states, actions + + +def _collect_relative_cartesian( + actions: np.ndarray, + states: np.ndarray, + index_range: list[int], + action_chunk: int, +) -> np.ndarray: + start, end = index_range + max_start = max(0, len(actions) - action_chunk) + chunks = [] + anchor_states = states[:max_start, start:end] + for offset in range(action_chunk): + chunks.append(actions[offset : offset + max_start, start:end] - anchor_states) + return np.concatenate(chunks, axis=0) + + +def _compute_delta_from_state_and_abs_rot( + rotations: np.ndarray, state: np.ndarray +) -> np.ndarray: + """Relative rotation: R_rel = R_abs @ R_state^T (same convention as the loader).""" + if rotations.shape[-1] == 3: + rotations_matrix = euler_to_matrix_zyx_batch_nb(rotations) + out_is_euler = True + elif rotations.shape[-1] == 6: + rotations_matrix = so3_to_matrix_batch_nb(rotations) + out_is_euler = False + else: + raise ValueError( + f"Only 3D euler or 6D rotation supported, got {rotations.shape[-1]}D" + ) + + if state.shape[-1] == 3: + state_matrix = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0] + elif state.shape[-1] == 6: + state_matrix = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] + else: + raise ValueError( + f"Only 3D euler or 6D rotation supported, got {state.shape[-1]}D" + ) + + return _abs_rot_to_delta(rotations_matrix, state_matrix, out_is_euler) + + +@jit(nopython=True, parallel=True) +def _abs_rot_to_delta( + rotations_matrix: np.ndarray, + state_matrix: np.ndarray, + out_is_euler: bool, +) -> np.ndarray: + st = np.empty((3, 3), dtype=np.float64) + st[0, 0] = state_matrix[0, 0] + st[0, 1] = state_matrix[1, 0] + st[0, 2] = state_matrix[2, 0] + st[1, 0] = state_matrix[0, 1] + st[1, 1] = state_matrix[1, 1] + st[1, 2] = state_matrix[2, 1] + st[2, 0] = state_matrix[0, 2] + st[2, 1] = state_matrix[1, 2] + st[2, 2] = state_matrix[2, 2] + + n = rotations_matrix.shape[0] + r_rel = np.empty((n, 3, 3), dtype=np.float64) + for i in prange(n): + a00 = rotations_matrix[i, 0, 0] + a01 = rotations_matrix[i, 0, 1] + a02 = rotations_matrix[i, 0, 2] + a10 = rotations_matrix[i, 1, 0] + a11 = rotations_matrix[i, 1, 1] + a12 = rotations_matrix[i, 1, 2] + a20 = rotations_matrix[i, 2, 0] + a21 = rotations_matrix[i, 2, 1] + a22 = rotations_matrix[i, 2, 2] + + r_rel[i, 0, 0] = a00 * st[0, 0] + a01 * st[1, 0] + a02 * st[2, 0] + r_rel[i, 0, 1] = a00 * st[0, 1] + a01 * st[1, 1] + a02 * st[2, 1] + r_rel[i, 0, 2] = a00 * st[0, 2] + a01 * st[1, 2] + a02 * st[2, 2] + r_rel[i, 1, 0] = a10 * st[0, 0] + a11 * st[1, 0] + a12 * st[2, 0] + r_rel[i, 1, 1] = a10 * st[0, 1] + a11 * st[1, 1] + a12 * st[2, 1] + r_rel[i, 1, 2] = a10 * st[0, 2] + a11 * st[1, 2] + a12 * st[2, 2] + r_rel[i, 2, 0] = a20 * st[0, 0] + a21 * st[1, 0] + a22 * st[2, 0] + r_rel[i, 2, 1] = a20 * st[0, 1] + a21 * st[1, 1] + a22 * st[2, 1] + r_rel[i, 2, 2] = a20 * st[0, 2] + a21 * st[1, 2] + a22 * st[2, 2] + + if out_is_euler: + d_euler = matrix_to_euler_zyx_batch_nb(r_rel) + return canonicalize_euler_zyx_batch_nb(d_euler) + + out6 = np.empty((n, 6), dtype=np.float64) + for i in prange(n): + out6[i, 0] = r_rel[i, 0, 0] + out6[i, 1] = r_rel[i, 0, 1] + out6[i, 2] = r_rel[i, 0, 2] + out6[i, 3] = r_rel[i, 1, 0] + out6[i, 4] = r_rel[i, 1, 1] + out6[i, 5] = r_rel[i, 1, 2] + return out6 + + +def _collect_relative_rotation( + actions: np.ndarray, + states: np.ndarray, + index_range: list[int], + action_chunk: int, +) -> np.ndarray: + start, end = index_range + max_start = max(0, len(actions) - action_chunk) + chunks = [] + for anchor_idx in tqdm( + range(max_start), + desc=" relative rotation anchors", + leave=False, + ): + action_clip = actions[anchor_idx : anchor_idx + action_chunk, start:end] + proprio_clip = states[anchor_idx, start:end] + rel = _compute_delta_from_state_and_abs_rot( + action_clip.astype(np.float64), proprio_clip.astype(np.float64) + ).astype(np.float32) + chunks.append(rel) + return np.concatenate(chunks, axis=0) + + +def collect_dof_vectors_from_arrays( + states: np.ndarray, + actions: np.ndarray, + propri_ranges: dict[str, list[int]], + action_ranges: dict[str, list[int]], + action_chunk: int = 32, +) -> dict[str, np.ndarray]: + vectors: dict[str, np.ndarray] = {} + + absolute_action_keys = { + key: index_range + for key, index_range in action_ranges.items() + if not key.endswith("_relative") + } + relative_action_keys = { + key: index_range + for key, index_range in action_ranges.items() + if key.endswith("_relative") + } - print(f"Reading the following fields:{non_image_columns}") - fast_dataset = dataset.hf_dataset.select_columns(non_image_columns) + for sub_key, index_range in propri_ranges.items(): + start, end = index_range + vectors[sub_key] = states[:, start:end] + + for sub_key, index_range in absolute_action_keys.items(): + start, end = index_range + vectors[sub_key] = actions[:, start:end] + + if relative_action_keys: + logging.info( + "Computing relative action slices (chunk=%d, anchors=%d)...", + action_chunk, + max(0, len(actions) - action_chunk), + ) + for sub_key, index_range in tqdm( + relative_action_keys.items(), desc="Relative action keys" + ): + if "rotation" in sub_key: + vectors[sub_key] = _collect_relative_rotation( + actions, states, index_range, action_chunk + ) + else: + vectors[sub_key] = _collect_relative_cartesian( + actions, states, index_range, action_chunk + ) - for i in tqdm(range(len(fast_dataset))): - sample = fast_dataset[i] - action = sample["action"] # torch.Tensor - propri = sample["observation.state"] + return vectors - for key, action_keys in trajectory_keys.items(): - for action_key, action_range in action_keys.items(): - if key == "action": - frames[repo_id][action_key].append( - action[action_range[0] : action_range[1]].numpy().tolist() - ) - else: - frames[repo_id][action_key].append( - propri[action_range[0] : action_range[1]].numpy().tolist() - ) - return frames +def compute_norm_stats_with_dof_config( + data_root: Path, + output_path: Path, + propri_ranges: dict[str, list[int]], + action_ranges: dict[str, list[int]], + state_key: str = "observation.state", + action_key: str = "action", + action_chunk: int = 32, +) -> dict[str, dict]: + table, state_key, action_key = _load_state_action_table( + data_root, state_key, action_key + ) + states, actions = _table_to_arrays(table, state_key, action_key) + norm_stats = { + state_key: compute_vector_stats(states), + action_key: compute_vector_stats(actions), + } + vectors = collect_dof_vectors_from_arrays( + states=states, + actions=actions, + propri_ranges=propri_ranges, + action_ranges=action_ranges, + action_chunk=action_chunk, + ) -def compute_action_normalizer( - repo_id: str, trajectory_keys: Dict, base_dir: Path, output_dir: Path -) -> None: - """ - Compute action normalizer statistics for all robot_ids. - """ - logging.info("Starting action normalizer computation...") + for sub_key, index_range in propri_ranges.items(): + if sub_key not in vectors: + logging.warning("No samples collected for propri key %s, skipping", sub_key) + continue + slice_stats = compute_vector_stats(vectors[sub_key]) + _apply_slice_stats(norm_stats[state_key], index_range, slice_stats) + logging.info(" %s (agent_pos): dim=%d", sub_key, len(slice_stats["mean"])) + + for sub_key, index_range in action_ranges.items(): + if sub_key not in vectors: + logging.warning("No samples collected for action key %s, skipping", sub_key) + continue + slice_stats = compute_vector_stats(vectors[sub_key]) + _apply_slice_stats(norm_stats[action_key], index_range, slice_stats) + mode = "relative" if sub_key.endswith("_relative") else "absolute" + logging.info(" %s (dof, %s): dim=%d", sub_key, mode, len(slice_stats["mean"])) + + write_norm_stats(output_path, norm_stats) + return norm_stats + + +def load_vectors( + data_root: Path, + state_key: str, + action_key: str, +) -> tuple[np.ndarray, np.ndarray]: + table, state_key, action_key = _load_state_action_table( + data_root, state_key, action_key + ) + return _table_to_arrays(table, state_key, action_key) + + +def compute_norm_stats( + data_root: Path, + output_path: Path, + state_key: str = "observation.state", + action_key: str = "action", + train_ctx: TrainConfigContext | None = None, +) -> dict[str, dict]: + if train_ctx is None: + states, actions = load_vectors(data_root, state_key, action_key) + norm_stats = { + state_key: compute_vector_stats(states), + action_key: compute_vector_stats(actions), + } + write_norm_stats(output_path, norm_stats) + return norm_stats + + return compute_norm_stats_with_dof_config( + data_root=data_root, + output_path=output_path, + propri_ranges=train_ctx.propri_ranges, + action_ranges=train_ctx.action_ranges, + state_key=state_key, + action_key=action_key, + action_chunk=train_ctx.action_chunk, + ) - frames = load_lerobot_dataset(repo_id, trajectory_keys, base_dir) - # Compute statistics - stats = compute_action_statistics(frames) +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Compute norm stats for a local LeRobot v3 dataset.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog="""\ +examples: + %(prog)s --train_config workspace/example/fintune_bus2602/libero.yml + %(prog)s --train_config workspace/example/fintune_bus2602/cvpr_example.yml \\ + --data_root /path/to/repo_id --output_path /path/to/norm_stats_path + %(prog)s --data_root /path/to/lerobot_dataset --output_path /path/to/norm_stats.json +""", + ) + parser.add_argument( + "--train_config", + type=str, + default=None, + help=( + "Training YAML config (e.g. cvpr_example.yml). When set, reads " + "data.lerobot_config.repo_id, data.norm_stats_path, task.dof_config, " + "task.agent_pos_config and task.action_horizon. Action keys ending " + "with '_relative' use the same relative-pose logic as lerobot loader." + ), + ) + parser.add_argument( + "--data_root", + type=str, + default=None, + help="Local LeRobot dataset directory (overrides train config)", + ) + parser.add_argument( + "--output_path", + type=str, + default=None, + help="Output json path (overrides train config)", + ) + parser.add_argument( + "--state_key", + type=str, + default=None, + help="Dataset column for proprioception (overrides train config)", + ) + parser.add_argument( + "--action_key", + type=str, + default=None, + help="Dataset column for action (overrides train config)", + ) + return parser.parse_args() - # Save statistics for each robot_id - output_dir = Path(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - # for robot_id, robot_stats in stats.items(): - # output_file = output_dir / f"{robot_id}_action_stats.json" - # write_json(output_file, robot_stats) - # logging.info(f"Saved action statistics for {robot_id} to {output_file}") +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") + args = parse_args() + + train_ctx: TrainConfigContext | None = None + if args.train_config: + config_path = Path(args.train_config) + if not config_path.exists(): + raise FileNotFoundError(f"train config not found: {config_path}") + train_ctx = parse_train_config(load_train_config(config_path)) + logging.info("train_config: %s", config_path) + logging.info(" dof_config keys: %s", list(train_ctx.action_ranges)) + logging.info(" agent_pos_config keys: %s", list(train_ctx.propri_ranges)) + logging.info(" action_chunk: %d", train_ctx.action_chunk) + + data_root = Path(args.data_root or (train_ctx.data_root if train_ctx else "")) + + output_path = Path(args.output_path or (train_ctx.output_path if train_ctx else "")) + state_key = args.state_key or ( + train_ctx.state_key if train_ctx else "observation.state" + ) + action_key = args.action_key or (train_ctx.action_key if train_ctx else "action") - # Also save a combined file - combined_output = output_dir / "all_robots_action_stats.json" - write_json(combined_output, stats) - logging.info(f"Saved combined action statistics to {combined_output}") + if not data_root.exists(): + raise FileNotFoundError(f"Dataset not found: {data_root}") + logging.info("dataset: %s", data_root) + logging.info("output: %s", output_path) -def main() -> None: + norm_stats = compute_norm_stats( + data_root=data_root, + output_path=output_path, + state_key=state_key, + action_key=action_key, + train_ctx=train_ctx, + ) - repo_id = "xxx" # your dataset name - data_root_path = "/path/to/lerobot/dataset" - output_stats_dir = "/path/to/save/action_stats" - trajectory_keys = { # your dataset keys - "propri": { - "follow_right_ee_cartesian_pos": [0, 3], - "follow_right_ee_rotation": [3, 6], - "follow_right_gripper": [6, 7], - }, - "action": { - "master_right_ee_cartesian_pos": [0, 3], - "master_right_ee_rotation": [3, 6], - "master_right_gripper": [6, 7], - }, - } + for key, stats in norm_stats.items(): + logging.info(" %s: dim=%d", key, len(stats["mean"])) - compute_action_normalizer( - repo_id, trajectory_keys, data_root_path, output_stats_dir - ) - logging.info("Action normalizer computation completed.") + logging.info("Saved norm stats to %s", output_path) if __name__ == "__main__": diff --git a/scripts/draw_openloop_plot.py b/scripts/draw_openloop_plot.py index 07b3d4f7..cfbf2816 100644 --- a/scripts/draw_openloop_plot.py +++ b/scripts/draw_openloop_plot.py @@ -1,125 +1,1052 @@ +#!/usr/bin/env python3 +"""Open-loop evaluation over WebSocket using LeRobot-format datasets. + +Loads episodes from a local LeRobot v3 dataset, sends observations to a running +Wall-X websocket server, collects predicted action chunks, and plots them against +ground-truth trajectories. + +Reference: ``infer_openloop_websocket.py`` (websocket client + open-loop loop). + +For single-arm LIBERO checkpoints, start the server with raw model output:: + + python -m wall_x._vendor.harrix.serving.launch_serving \\ + --env X2ROBOT --port 32194 \\ + --no-serialize-actions \\ + model-config:server-model-config \\ + --model-config.checkpoint-path /path/to/ckpt \\ + --model-config.train-config-path /path/to/libero.yml \\ + --model-config.action-horizon 10 \\ + --model-config.robot-type desktop +""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import json +import logging import os +import sys +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import cv2 +import msgpack +import msgpack_numpy as m +import numpy as np +import websockets import yaml -import torch -import argparse -from tqdm import tqdm -import matplotlib.pyplot as plt -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction -from wall_x.data.load_lerobot_dataset import load_test_dataset, get_data_configs -from wall_x.model.model_utils import register_normalizers -import copy +m.patch() -def load_config(config_path): - """Load configuration from YAML file.""" - with open(config_path, "r") as f: - config = yaml.load(f, Loader=yaml.FullLoader) +logger = logging.getLogger(__name__) - config["data"]["model_type"] = config.get("model_type") +_DEFAULT_CAM_WS_KEYS = { + "observation.images.faceImg": "face_view", + "observation.images.rightImg": "right_wrist_view", + "observation.images.leftImg": "left_wrist_view", + "observation.images.move1Img": "move1_view", +} - return config +@dataclass +class EpisodeArrays: + """Episode state/action plus lazily decoded camera frames.""" -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--pred_horizon", type=int, default=32) - parser.add_argument("--origin_action_dim", type=int, default=14) - args = parser.parse_args() - - origin_action_dim = args.origin_action_dim - pred_horizon = args.pred_horizon - - # get train config - model_path = "/path/to/your/checkpoint" - action_tokenizer_path = "/path/to/Models/fast" - save_dir = "/path/to/save/dir" - path = f"{model_path}/config.yml" - config = load_config(path) - - normalizer_action, normalizer_propri = register_normalizers(config, model_path) - - # load model with customized robot config - model = Qwen2_5_VLMoEForAction.from_pretrained( - model_path, train_config=config, action_tokenizer_path=action_tokenizer_path - ) - - model.set_normalizer( - copy.deepcopy(normalizer_action), copy.deepcopy(normalizer_propri) - ) - model.eval() - model = model.to("cuda") - model.to_bfloat16_for_selected_params() - - # get test dataloader - dataload_config = get_data_configs(config["data"]) - lerobot_config = dataload_config.get("lerobot_config", {}) - dataset = load_test_dataset( - config, lerobot_config, normalizer_action, normalizer_propri, seed=42 - ) - dataloader = dataset.get_dataloader() - # dataloader = dataset.get_train_dataloader() - - total_frames = len(dataloader) - - predict_mode = "fast" if config.get("use_fast_tokenizer", False) else "diffusion" - action_dim = 14 if predict_mode == "diffusion" else origin_action_dim - gt_traj = torch.zeros((total_frames, origin_action_dim)) - pred_traj = torch.zeros((total_frames, origin_action_dim)) - - # use tqdm to show the progress - for idx, batch in tqdm( - enumerate(dataloader), total=total_frames, desc="predicting" - ): - if idx % pred_horizon == 0 and idx + pred_horizon < total_frames: - batch = batch.to("cuda") - with torch.no_grad(): - outputs = model( - **batch, - action_dim=action_dim, - action_horizon=pred_horizon, - mode="predict", - predict_mode=predict_mode, + episode_index: int + instruction: str + states: np.ndarray # (T, D_state) + actions: np.ndarray # (T, D_action) + camera_keys: list[str] + frame_offset: int = 0 + num_steps: int = 0 + _dataset: Any = field(default=None, repr=False) + _image_cache: dict[int, dict[str, np.ndarray]] = field( + default_factory=dict, repr=False + ) + + def get_frame_images(self, frame_idx: int) -> dict[str, np.ndarray]: + """Decode camera frames for one observation index (cached).""" + if frame_idx in self._image_cache: + return self._image_cache[frame_idx] + if self._dataset is None: + raise RuntimeError("Episode image loader is not initialized.") + if not self._image_cache: + logger.info("Decoding camera frames on demand (first obs_idx=%d)", frame_idx) + item = self._dataset[frame_idx] + images = { + cam_key: tensor_to_rgb_uint8(item[cam_key]) + for cam_key in self.camera_keys + if cam_key in item + } + self._image_cache[frame_idx] = images + return images + + +def _repo_root() -> Path: + return Path(__file__).resolve().parents[1] + + +def load_train_config(path: str | Path) -> dict[str, Any]: + with open(path, "r", encoding="utf-8") as f: + cfg = yaml.safe_load(f) or {} + if "data" not in cfg: + cfg["data"] = {} + cfg["data"]["model_type"] = cfg.get("model_type") + return cfg + + +def parse_int_list(value: str | None) -> list[int] | None: + if value is None or not value.strip(): + return None + return [int(x.strip()) for x in value.split(",") if x.strip()] + + +def _task_block(train_config: dict[str, Any]) -> dict[str, Any]: + return train_config.get("task") or {} + + +def _dof_config(train_config: dict[str, Any]) -> dict[str, int]: + return train_config.get("dof_config") or _task_block(train_config).get("dof_config") or {} + + +def _agent_pos_config(train_config: dict[str, Any]) -> dict[str, int]: + return ( + train_config.get("agent_pos_config") + or _task_block(train_config).get("agent_pos_config") + or {} + ) + + +def real_vector_dim(config_block: dict[str, int]) -> int: + """Sum dof/agent dims excluding ``action_padding``.""" + return sum(d for k, d in config_block.items() if k != "action_padding" and d > 0) + + +def slice_by_config(vec: np.ndarray, config_block: dict[str, int]) -> np.ndarray: + """Keep only non-padding dims from a 1-D vector (state / single action row).""" + vec = np.asarray(vec, dtype=np.float32).reshape(-1) + if not config_block: + return vec + parts: list[np.ndarray] = [] + start = 0 + for key, dim in config_block.items(): + if key == "action_padding": + start += dim + continue + end = start + dim + if end <= vec.shape[0]: + parts.append(vec[start:end]) + start = end + if parts: + return np.concatenate(parts, axis=0).astype(np.float32) + return vec + + +def strip_padding_columns(arr: np.ndarray, config_block: dict[str, int]) -> np.ndarray: + """Drop ``action_padding`` columns from a [T, D] action chunk.""" + arr = np.asarray(arr, dtype=np.float32) + if arr.ndim == 1: + return slice_by_config(arr, config_block) + if not config_block: + return arr + parts: list[np.ndarray] = [] + start = 0 + for key, dim in config_block.items(): + if key == "action_padding": + start += dim + continue + parts.append(arr[:, start : start + dim]) + start += dim + if parts: + return np.concatenate(parts, axis=1).astype(np.float32) + return arr + + +def tensor_to_rgb_uint8(img: Any) -> np.ndarray: + arr = np.asarray(img) + if arr.ndim == 3 and arr.shape[0] in (1, 3) and arr.shape[0] != arr.shape[-1]: + arr = np.transpose(arr, (1, 2, 0)) + if np.issubdtype(arr.dtype, np.floating): + arr = np.clip(arr, 0.0, 1.0) * 255.0 + return arr.astype(np.uint8) + + +def encode_image_rgb(image: np.ndarray) -> str: + """JPEG base64; input must be RGB uint8.""" + bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) + ok, buffer = cv2.imencode(".jpg", bgr) + if not ok: + raise RuntimeError("cv2.imencode failed") + return base64.b64encode(buffer).decode("utf-8") + + +def build_camera_ws_mapping(train_config: dict[str, Any]) -> dict[str, str]: + key_mappings = (train_config.get("data") or {}).get("key_mappings") or {} + cam_map = key_mappings.get("camera") or {} + if cam_map: + return dict(cam_map) + return dict(_DEFAULT_CAM_WS_KEYS) + + +def resolve_feature_keys( + train_config: dict[str, Any], + state_key: str | None, + action_key: str | None, +) -> tuple[str, str]: + key_mappings = (train_config.get("data") or {}).get("key_mappings") or {} + sk = state_key or key_mappings.get("state") or "observation.state" + ak = action_key or key_mappings.get("action") or "action" + return sk, ak + + +def _arm_follow_pos_from_agent_cfg( + state_vec: np.ndarray, + agent_cfg: dict[str, int], + arm_prefix: str, +) -> list[float] | None: + """Pack one arm into websocket ``follow{1,2}_pos`` layout: pos3 + rpy3 + grip1.""" + from wall_x._vendor.x2robot_utils import geometry as geom + + pos = rot = grip = None + start = 0 + for key, dim in agent_cfg.items(): + if key == "action_padding": + start += int(dim) + continue + end = start + int(dim) + chunk = ( + state_vec[start:end] + if end <= state_vec.shape[0] + else np.zeros(int(dim), dtype=np.float32) + ) + start = end + if not key.startswith(arm_prefix): + continue + if "cartesian_pos" in key: + pos = chunk.reshape(-1) + elif "rotation_6d" in key.lower(): + rot = geom.so3_to_euler_zyx_batch_nb(chunk.reshape(1, -1)).reshape(-1) + elif "rotation" in key: + rot = chunk.reshape(-1) + elif "gripper" in key: + grip = chunk.reshape(-1) + + if pos is None or rot is None or grip is None: + return None + packed = np.concatenate([pos[:3], rot[:3], grip[:1]], axis=0).astype(np.float32) + if packed.shape[0] != 7: + return None + return packed.tolist() + + +def build_state_payload( + state_vec: np.ndarray, + train_config: dict[str, Any], +) -> dict[str, list[float]]: + """Build websocket ``state`` dict from a flat proprio vector.""" + state_vec = np.asarray(state_vec, dtype=np.float32).reshape(-1) + + # LeRobot dual-arm export in this workflow is already raw follow-pos 14D: + # [L_pos3, L_rot3, L_grip1, R_pos3, R_rot3, R_grip1]. + # Prefer direct passthrough and avoid parsing with agent_pos_config + # (which may be 6D/padded and would mis-slice a 14D state). + if state_vec.shape[0] >= 14: + return { + "follow1_pos": state_vec[:7].tolist(), + "follow2_pos": state_vec[7:14].tolist(), + } + if state_vec.shape[0] >= 7: + return {"follow2_pos": state_vec[:7].tolist()} + + agent_cfg = _agent_pos_config(train_config) + + if agent_cfg: + has_left = any(k.startswith("follow_left_") for k in agent_cfg) + has_right = any(k.startswith("follow_right_") for k in agent_cfg) + payload: dict[str, list[float]] = {} + if has_left: + left = _arm_follow_pos_from_agent_cfg( + state_vec, agent_cfg, "follow_left_" + ) + if left is not None: + payload["follow1_pos"] = left + if has_right: + right = _arm_follow_pos_from_agent_cfg( + state_vec, agent_cfg, "follow_right_" + ) + if right is not None: + payload["follow2_pos"] = right + if payload: + return payload + + # Legacy single-arm LIBERO path: only follow_right_* in agent_pos_config. + real_dim = real_vector_dim(agent_cfg) + if real_dim > 0 and state_vec.shape[0] >= real_dim: + sliced = slice_by_config(state_vec, agent_cfg) + if sliced.shape[0] == 0: + sliced = state_vec[:real_dim] + right = _arm_follow_pos_from_agent_cfg( + state_vec, agent_cfg, "follow_right_" + ) + if right is not None: + return {"follow2_pos": right} + return {"follow2_pos": sliced.tolist()} + raise ValueError(f"Unsupported state dimension: {state_vec.shape[0]}") + + +def decode_predict_action( + predict_action: Any, + train_config: dict[str, Any], +) -> np.ndarray: + """Convert raw ``predict_action`` [H, D_model] to flat action rows per step.""" + if hasattr(predict_action, "detach"): + predict_action = predict_action.detach().cpu().numpy() + pa = np.asarray(predict_action, dtype=np.float32) + if pa.ndim == 3: + pa = pa[0] + + dof_cfg = _dof_config(train_config) + if dof_cfg: + parts: list[np.ndarray] = [] + start = 0 + for key, dim in dof_cfg.items(): + if key == "action_padding": + start += dim + continue + parts.append(pa[:, start : start + dim]) + start += dim + if parts: + return np.concatenate(parts, axis=1).astype(np.float32) + + try: + from wall_x._vendor.harrix.envs.libero_common import decode_chunk + + chunk = decode_chunk(pa, train_config) + if chunk.ndim == 2 and chunk.shape[1] > 0: + return chunk.astype(np.float32) + except Exception: + pass + + return pa.astype(np.float32) + + +def extract_follow_pos_14d_from_response(result: dict[str, Any]) -> np.ndarray: + """Read serialized ``follow{1,2}_pos`` action rows (skip state row 0) as ``[H, 14]``.""" + if "follow2_pos" not in result or "follow1_pos" not in result: + raise KeyError("Response missing follow1_pos/follow2_pos for 14D eval.") + left = np.asarray(result["follow1_pos"], dtype=np.float32) + right = np.asarray(result["follow2_pos"], dtype=np.float32) + if left.ndim == 1: + left = left.reshape(1, -1) + if right.ndim == 1: + right = right.reshape(1, -1) + if left.shape[0] < 2 or right.shape[0] < 2: + raise ValueError( + f"follow_pos response must include state+action rows, got " + f"left={left.shape}, right={right.shape}" + ) + horizon = min(left.shape[0] - 1, right.shape[0] - 1) + return np.concatenate([left[1 : 1 + horizon], right[1 : 1 + horizon]], axis=1) + + +def extract_action_chunk_from_response( + result: dict[str, Any], + train_config: dict[str, Any], +) -> np.ndarray: + if "predict_action" in result: + return decode_predict_action(result["predict_action"], train_config) + + if "follow2_pos" in result and "follow1_pos" in result: + return extract_follow_pos_14d_from_response(result) + + if "follow2_pos" in result: + right = np.asarray(result["follow2_pos"], dtype=np.float32) + if right.ndim == 1: + right = right.reshape(1, -1) + return right[1:] if right.shape[0] > 1 else right + + if "action" in result: + action = np.asarray(result["action"], dtype=np.float32) + if action.ndim == 1: + action = action[np.newaxis, :] + return action + + raise KeyError( + "Response has no action fields. For single-arm LIBERO, restart the server with " + "`--no-serialize-actions` so responses include `predict_action`." + ) + + +def _hf_row_to_numpy(row: Any, key: str) -> np.ndarray: + value = row[key] + if hasattr(value, "numpy"): + value = value.numpy() + return np.asarray(value, dtype=np.float32).reshape(-1) + + +def _episode_instruction(ds: Any, row_index: int = 0) -> str: + row = ds.hf_dataset[row_index] + task_idx = row["task_index"] + if hasattr(task_idx, "item"): + task_idx = task_idx.item() + return str(ds.meta.tasks.iloc[task_idx].name) + + +def resolve_lerobot_dataset_paths(dataset_root: str | Path) -> tuple[str, Path]: + """Return ``(repo_id, root)`` for a local LeRobot dataset directory.""" + root = Path(dataset_root).expanduser().resolve() + if not root.is_dir(): + raise FileNotFoundError(f"LeRobot dataset root not found: {root}") + return root.name, root + + +def open_local_lerobot_dataset( + dataset_root: str | Path, + *, + episodes: list[int] | None = None, + video_backend: str = "pyav", +) -> Any: + """Open a local LeRobot v3 dataset without HuggingFace Hub fallback.""" + import torch + from lerobot.datasets.lerobot_dataset import ( + CODEBASE_VERSION, + LeRobotDataset, + LeRobotDatasetMetadata, + ) + from lerobot.datasets.utils import ( + get_hf_features_from_features, + hf_transform_to_torch, + load_nested_dataset, + ) + from lerobot.datasets.video_utils import get_safe_default_codec + + repo_id, root = resolve_lerobot_dataset_paths(dataset_root) + meta = LeRobotDatasetMetadata(repo_id, root=root) + + features = get_hf_features_from_features(meta.features) + hf_dataset = load_nested_dataset( + root / "data", features=features, episodes=episodes + ) + hf_dataset.set_transform(hf_transform_to_torch) + + if episodes is not None: + available = { + int(ep.item()) if hasattr(ep, "item") else int(ep) + for ep in hf_dataset.unique("episode_index") + } + missing = set(episodes) - available + if missing: + raise ValueError( + f"Episodes {sorted(missing)} not found under {root}. " + f"Available in loaded parquet: {sorted(available)[:10]}" + f"{'...' if len(available) > 10 else ''}" + ) + + if meta.video_keys: + check_eps = ( + episodes if episodes is not None else list(range(meta.total_episodes)) + ) + for ep_idx in check_eps: + for vid_key in meta.video_keys: + video_path = root / meta.get_video_file_path(ep_idx, vid_key) + if not video_path.exists(): + raise FileNotFoundError( + "LeRobot camera videos are missing for open-loop eval.\n" + f" dataset root: {root}\n" + f" first missing file: {video_path}\n" + "Parquet state/action may exist, but this script also needs " + "mp4 files under videos/." + ) + + ds = LeRobotDataset.__new__(LeRobotDataset) + ds.repo_id = repo_id + ds.root = root + ds.image_transforms = None + ds.delta_timestamps = None + ds.episodes = episodes + ds.tolerance_s = 1e-4 + ds.revision = CODEBASE_VERSION + ds.video_backend = video_backend or get_safe_default_codec() + ds.delta_indices = None + ds.meta = meta + ds.hf_dataset = hf_dataset + ds._lazy_loading = False + ds._absolute_to_relative_idx = None + if episodes is not None: + ds._absolute_to_relative_idx = { + abs_idx.item() if isinstance(abs_idx, torch.Tensor) else abs_idx: rel_idx + for rel_idx, abs_idx in enumerate(hf_dataset["index"]) + } + ds.image_writer = None + ds.episode_buffer = None + ds.writer = None + ds.latest_episode = None + ds._current_file_start_frame = None + ds._streaming_encoder = None + ds.batch_encoding_size = 1 + ds.episodes_since_last_encoding = 0 + return ds + + +def load_lerobot_metadata(dataset_root: str | Path) -> Any: + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata + + repo_id, root = resolve_lerobot_dataset_paths(dataset_root) + return LeRobotDatasetMetadata(repo_id, root=root) + + +def plan_eval_frame_range( + num_steps: int, + start_idx: int, + step_stride: int, + action_horizon: int, + max_inferences: int | None, +) -> tuple[int, int, int]: + """Return ``(tabular_start, tabular_end, expected_video_decodes)``.""" + tabular_start = max(0, int(start_idx)) + if max_inferences is None: + return tabular_start, num_steps, -1 + + tabular_end = tabular_start + idx = tabular_start + infer_count = 0 + video_decodes = 0 + while idx <= num_steps - action_horizon - 1: + if infer_count >= max_inferences: + break + record_n = min(step_stride, action_horizon, num_steps - idx) + if record_n <= 0: + break + tabular_end = max(tabular_end, idx + record_n) + video_decodes += 1 + idx += record_n + infer_count += 1 + return tabular_start, min(num_steps, tabular_end), video_decodes + + +def load_episode_arrays( + dataset_root: str | Path, + episode_index: int, + state_key: str, + action_key: str, + camera_keys: list[str], + *, + frame_start: int = 0, + frame_end: int | None = None, + preload_all_images: bool = False, +) -> EpisodeArrays: + ds = open_local_lerobot_dataset( + dataset_root, + episodes=[episode_index], + video_backend="pyav", + ) + num_steps = len(ds.hf_dataset) + if num_steps == 0: + raise ValueError(f"Episode {episode_index} is empty under {dataset_root}") + + start = max(0, int(frame_start)) + end = num_steps if frame_end is None else min(num_steps, int(frame_end)) + states = [ + _hf_row_to_numpy(ds.hf_dataset[i], state_key) for i in range(start, end) + ] + actions = [ + _hf_row_to_numpy(ds.hf_dataset[i], action_key) for i in range(start, end) + ] + instruction = _episode_instruction(ds, start) + + episode = EpisodeArrays( + episode_index=episode_index, + instruction=instruction, + states=np.stack(states, axis=0), + actions=np.stack(actions, axis=0), + camera_keys=list(camera_keys), + frame_offset=start, + num_steps=num_steps, + _dataset=ds, + ) + if preload_all_images: + for frame_idx in range(start, end): + episode.get_frame_images(frame_idx) + return episode + + +def build_obs_payload( + episode: EpisodeArrays, + frame_idx: int, + train_config: dict[str, Any], + cam_ws_mapping: dict[str, str], + extra_view_keys: list[str], +) -> dict[str, Any]: + local_idx = frame_idx - episode.frame_offset + if local_idx < 0 or local_idx >= episode.states.shape[0]: + raise IndexError( + f"Frame {frame_idx} is outside loaded tabular range " + f"[{episode.frame_offset}, {episode.frame_offset + episode.states.shape[0]})." + ) + state = episode.states[local_idx] + frame_images = episode.get_frame_images(frame_idx) + views: dict[str, str] = {} + + ref_shape: tuple[int, int, int] | None = None + for cam_key, ws_key in cam_ws_mapping.items(): + if cam_key not in frame_images: + continue + rgb = frame_images[cam_key] + ref_shape = rgb.shape + views[ws_key] = encode_image_rgb(rgb) + + for ws_key in extra_view_keys: + if ws_key in views: + continue + if ref_shape is None: + ref_shape = (256, 256, 3) + views[ws_key] = encode_image_rgb(np.zeros(ref_shape, dtype=np.uint8)) + + return { + "state": build_state_payload(state, train_config), + "views": views, + "instruction": episode.instruction, + } + + +def _dim_label(i: int, dim: int) -> str: + arm7 = [ + "pos_x", + "pos_y", + "pos_z", + "rot_x", + "rot_y", + "rot_z", + "gripper", + ] + if dim == 14: + prefix = "L_" if i < 7 else "R_" + return prefix + arm7[i % 7] + if dim == 7: + return arm7[i] if i < len(arm7) else f"dim_{i}" + if dim == 12: + prefix = "L_" if i < 6 else "R_" + return f"{prefix}rot6_{i % 6}" + if dim == 10: + names = [ + "pos_x", + "pos_y", + "pos_z", + "rot6_0", + "rot6_1", + "rot6_2", + "rot6_3", + "rot6_4", + "rot6_5", + "gripper", + ] + return names[i] if i < len(names) else f"dim_{i}" + return f"dim_{i}" + + +def follow_pos_14d_to_rotation_6d(rows: list | np.ndarray) -> np.ndarray: + """Convert ``follow1_pos + follow2_pos`` rows from euler rotation to rot6D only.""" + from wall_x._vendor.x2robot_utils import geometry as geom + + arr = np.asarray(rows, dtype=np.float32) + if arr.ndim == 1: + arr = arr[np.newaxis, :] + if arr.shape[1] < 14: + raise ValueError(f"Expected follow_pos_14d rows, got shape {arr.shape}") + + left_rot6 = geom.euler_to_matrix_zyx_6d_nb(arr[:, 3:6]).astype(np.float32) + right_rot6 = geom.euler_to_matrix_zyx_6d_nb(arr[:, 10:13]).astype(np.float32) + return np.concatenate([left_rot6, right_rot6], axis=1).astype(np.float32) + + +def plot_openloop( + action_pred_list: list | np.ndarray, + action_gt_list: list | np.ndarray, + save_path: Path, + *, + title: str = "", + action_l1: float | None = None, +) -> None: + """Plot aligned GT vs predicted action trajectories.""" + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + gt_rows = np.asarray(action_gt_list, dtype=np.float32) + pred_rows = np.asarray(action_pred_list, dtype=np.float32) + assert gt_rows.shape == pred_rows.shape, ( + f"shape mismatch: pred={pred_rows.shape} gt={gt_rows.shape}" + ) + n, dim = gt_rows.shape + + fig, axes = plt.subplots(dim, 1, figsize=(12, 3.5 * dim), sharex=True) + if dim == 1: + axes = [axes] + + for i, ax in enumerate(axes): + ax.plot(gt_rows[:, i], label="Ground Truth", color="blue", linewidth=1.5) + ax.plot(pred_rows[:, i], label="Model Output", color="orange", linewidth=1.5) + ax.set_title(_dim_label(i, dim)) + ax.set_ylabel("Action Value") + ax.grid(True, alpha=0.25) + ax.legend(loc="best") + + axes[-1].set_xlabel("Time Step") + axes[0].set_xticks(np.arange(0, n, step=max(1, min(10, max(n // 15, 1))))) + l1_note = f"Mean L1: {action_l1:.6f}" if action_l1 is not None else "" + if title and l1_note: + fig.suptitle(f"{title}\n{l1_note}") + elif title: + fig.suptitle(title) + elif l1_note: + fig.suptitle(l1_note) + fig.tight_layout(rect=[0, 0, 1, 0.97]) + + save_path.parent.mkdir(parents=True, exist_ok=True) + out = save_path.with_suffix(".jpg") if save_path.suffix != ".jpg" else save_path + fig.savefig(out, dpi=200) + plt.close(fig) + logger.info("Saved plot -> %s", out) + + +def compute_second_diff(arr: np.ndarray) -> np.ndarray: + d1 = arr[1:] - arr[:-1] + return d1[1:] - d1[:-1] + + +async def run_openloop_eval( + uri: str, + dataset_root: str, + train_config_path: str, + save_dir: str, + episode_indices: list[int] | None, + start_ratio: float, + stride: int | None, + max_inferences: int | None, + state_key: str | None, + action_key: str | None, + extra_view_keys: list[str], + preload_all_images: bool = False, + plot_rotation_6d: bool = False, +) -> None: + train_config = load_train_config(train_config_path) + state_key, action_key = resolve_feature_keys(train_config, state_key, action_key) + cam_ws_mapping = build_camera_ws_mapping(train_config) + camera_keys = list(cam_ws_mapping.keys()) + + action_horizon = ( + train_config.get("action_horizon_flow") + or _task_block(train_config).get("action_horizon_flow") + or _task_block(train_config).get("action_horizon") + or train_config.get("action_horizon") + or 10 + ) + step_stride = stride if stride is not None else int(action_horizon) + + meta = load_lerobot_metadata(dataset_root) + if episode_indices is None: + episode_indices = [0] + episode_indices = [i for i in episode_indices if 0 <= i < meta.total_episodes] + if not episode_indices: + raise ValueError("No valid episode indices to evaluate.") + + save_root = Path(save_dir) + save_root.mkdir(parents=True, exist_ok=True) + + async with websockets.connect( + uri, + ping_interval=None, + ping_timeout=None, + max_size=None, + ) as websocket: + metadata = msgpack.unpackb(await websocket.recv()) + logger.info("Connected to %s, server metadata: %s", uri, metadata) + + for ep_idx in episode_indices: + ep_meta = meta.episodes[ep_idx] + num_steps = int( + ep_meta.get("length") + if isinstance(ep_meta, dict) + else getattr(ep_meta, "length", 0) + ) + if num_steps <= 0: + raise ValueError(f"Episode {ep_idx} has invalid length metadata.") + + start_idx = int(start_ratio * num_steps) + tabular_start, tabular_end, expected_decodes = plan_eval_frame_range( + num_steps, + start_idx, + step_stride, + int(action_horizon), + max_inferences, + ) + logger.info( + "Loading episode %d tabular frames [%d, %d) / %d " + "(images on demand%s)...", + ep_idx, + tabular_start, + tabular_end, + num_steps, + f", ~{expected_decodes} obs frames" if expected_decodes >= 0 else "", + ) + episode = load_episode_arrays( + dataset_root, + ep_idx, + state_key=state_key, + action_key=action_key, + camera_keys=camera_keys, + frame_start=tabular_start, + frame_end=tabular_end, + preload_all_images=preload_all_images, + ) + logger.info( + "Episode %d tabular ready: %d frames loaded, video decode per inference", + ep_idx, + episode.states.shape[0], + ) + + gt_full = np.asarray(episode.actions, dtype=np.float32) + if gt_full.ndim == 1: + gt_full = gt_full[np.newaxis, :] + if gt_full.shape[1] != 14: + raise ValueError( + f"Expected LeRobot action dim=14, got {gt_full.shape}. " + "This script now assumes dataset actions and server outputs are both 14D." ) - pred_traj[idx : idx + pred_horizon] = ( - outputs["predict_action"][:, :, :origin_action_dim] - .detach() - .cpu() - .squeeze(0) + pred_full = np.full_like(gt_full, np.nan) + aligned_gt: list[list[float]] = [] + aligned_pred: list[list[float]] = [] + obs_infer_points: list[tuple[int, int]] = [] + + idx = start_idx + infer_count = 0 + while idx <= num_steps - action_horizon - 1: + if max_inferences is not None and infer_count >= max_inferences: + break + + payload = build_obs_payload( + episode, + idx, + train_config, + cam_ws_mapping, + extra_view_keys=extra_view_keys, + ) + await websocket.send(msgpack.packb(payload, use_bin_type=True)) + raw = await websocket.recv() + if isinstance(raw, str): + hint = "" + if "get_serialized_actions" in raw: + hint = ( + "\n\nHint: restart the server with `--no-serialize-actions` " + "so responses include `predict_action`." + ) + raise RuntimeError(f"Server error at frame {idx}:\n{raw}{hint}") + + result = msgpack.unpackb(raw, raw=False) + pred_chunk = extract_action_chunk_from_response(result, train_config) + record_n = min( + step_stride, action_horizon, num_steps - idx, pred_chunk.shape[0] ) + if record_n <= 0: + break - # Denormalize ground truth actions - gt_action_chunk = batch["action_chunk"][:, :, :origin_action_dim] - dof_mask = batch["dof_mask"].to(gt_action_chunk.dtype) - denormalized_gt = ( - model.action_preprocessor.normalizer_action.unnormalize_data( - gt_action_chunk, - [lerobot_config.get("repo_id", "physical-intelligence/libero")], - dof_mask, - ).squeeze(0) + local_idx = idx - episode.frame_offset + gt_chunk = gt_full[local_idx : local_idx + record_n] + pred_chunk = pred_chunk[:record_n] + record_n = min(record_n, pred_chunk.shape[0], gt_chunk.shape[0]) + if record_n <= 0: + break + + row0 = len(aligned_gt) + obs_infer_points.append((row0, idx)) + + pred_full[local_idx : local_idx + record_n] = pred_chunk[:record_n] + aligned_gt.extend(gt_chunk[:record_n].tolist()) + aligned_pred.extend(pred_chunk[:record_n].tolist()) + + logger.info( + "episode=%d obs_idx=%d row0=%d record_n=%d pred_shape=%s gt_dim=%d", + ep_idx, + idx, + row0, + record_n, + pred_chunk.shape, + gt_chunk.shape[1], + ) + + idx += record_n + infer_count += 1 + + if not obs_infer_points: + logger.warning("Episode %d: no inference steps executed.", ep_idx) + continue + + ep_tag = f"ep{ep_idx}" + valid_mask = ~np.isnan(pred_full).any(axis=1) + action_l1 = float(np.mean(np.abs(pred_full[valid_mask] - gt_full[valid_mask]))) + summary = { + "episode_index": ep_idx, + "instruction": episode.instruction, + "num_steps": num_steps, + "obs_infer_points": obs_infer_points, + "action_horizon": action_horizon, + "stride": step_stride, + "mean_l1": action_l1, + "state_key": state_key, + "action_key": action_key, + } + summary_path = save_root / f"{ep_tag}_summary.json" + summary_path.write_text( + json.dumps(summary, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", ) - gt_traj[idx : idx + pred_horizon] = denormalized_gt.detach().cpu() + logger.info("Episode %d mean L1 = %.6f", ep_idx, action_l1) + + ml = min(len(aligned_gt), len(aligned_pred)) + if ml < 3: + logger.warning("Episode %d: too few aligned rows (%d) to plot.", ep_idx, ml) + continue - gt_traj_np = gt_traj.numpy() - pred_traj_np = pred_traj.numpy() + gt_ml = np.asarray(aligned_gt[:ml], dtype=np.float32) + pred_ml = np.asarray(aligned_pred[:ml], dtype=np.float32) + title = f"Episode {ep_idx}: {episode.instruction[:80]}" + plot_openloop( + pred_ml, + gt_ml, + save_root / ep_tag, + title=title, + action_l1=action_l1, + ) + if plot_rotation_6d: + gt_rot6d = follow_pos_14d_to_rotation_6d(gt_ml) + pred_rot6d = follow_pos_14d_to_rotation_6d(pred_ml) + rot6d_l1 = float(np.mean(np.abs(pred_rot6d - gt_rot6d))) + plot_openloop( + pred_rot6d, + gt_rot6d, + save_root / f"{ep_tag}_rot6d", + title=f"{title} (rotation 6D)", + action_l1=rot6d_l1, + ) - timesteps = gt_traj.shape[0] - fig, axs = plt.subplots( - origin_action_dim, 1, figsize=(15, 5 * origin_action_dim), sharex=True +def parse_args() -> argparse.Namespace: + repo = _repo_root() + parser = argparse.ArgumentParser( + description="Open-loop websocket inference on LeRobot datasets.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--uri", default="ws://127.0.0.1:32194") + parser.add_argument( + "--dataset-root", + default="/x2robot_v2/share/yangping/data/lerobot/benchmark/libero_all", + help="Local LeRobot dataset root.", + ) + parser.add_argument( + "--train-config", + default=str(repo / "workspace/example/fintune_bus2602/libero2_6d.yml"), + help="Training YAML for key mappings / dof layout.", + ) + parser.add_argument( + "--save-dir", + default=str(repo / "openloop_lerobot_plots"), + ) + parser.add_argument( + "--episode-indices", + default="0", + help="Comma-separated episode indices, e.g. 0,1,2", + ) + parser.add_argument( + "--start-ratio", + type=float, + default=0.0, + help="Start open-loop from this fraction of the episode length.", + ) + parser.add_argument( + "--stride", + type=int, + default=None, + help="Frames to advance between inferences (default: action_horizon).", + ) + parser.add_argument( + "--max-inferences", + type=int, + default=None, + help=( + "Cap inference requests per episode; also limits tabular/video loading " + "to the evaluated frame range." + ), + ) + parser.add_argument( + "--preload-all-images", + action="store_true", + help="Decode all camera frames up front (slow; old behavior).", + ) + parser.add_argument( + "--plot-rotation-6d", + action="store_true", + help="Also save a rotation-6D-only plot as ep*_rot6d.jpg.", + ) + parser.add_argument( + "--state-key", + default=None, + help="LeRobot state feature key (default: from train config key_mappings).", + ) + parser.add_argument( + "--action-key", + default=None, + help="LeRobot action feature key (default: from train config key_mappings).", + ) + parser.add_argument( + "--extra-view-keys", + default="left_wrist_view", + help="Websocket view keys to fill with black images when missing from data.", + ) + parser.add_argument("--log-level", default="INFO") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + logging.basicConfig( + level=getattr(logging, args.log_level.upper(), logging.INFO), + format="%(asctime)s - %(levelname)s - %(message)s", ) - fig.suptitle("Action Comparison for lerobot", fontsize=16) - for i in range(origin_action_dim): - axs[i].plot(range(timesteps), gt_traj_np[:, i], label="Ground Truth") - axs[i].plot(range(timesteps), pred_traj_np[:, i], label="Prediction") - axs[i].set_ylabel(f"Action Dim {i+1}") - axs[i].legend() - axs[i].grid(True) + repo_root = _repo_root() + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) - axs[-1].set_xlabel("Timestep") - plt.tight_layout(rect=[0, 0.03, 1, 0.95]) - os.makedirs(save_dir, exist_ok=True) - save_path = os.path.join(save_dir, "lerobot_comparison.png") - plt.savefig(save_path) - print(f"Saved plot to {save_path}") - plt.close() + extra_view_keys = [k.strip() for k in args.extra_view_keys.split(",") if k.strip()] + episode_indices = parse_int_list(args.episode_indices) + + try: + asyncio.run( + run_openloop_eval( + uri=args.uri, + dataset_root=args.dataset_root, + train_config_path=args.train_config, + save_dir=args.save_dir, + episode_indices=episode_indices, + start_ratio=args.start_ratio, + stride=args.stride, + max_inferences=args.max_inferences, + state_key=args.state_key, + action_key=args.action_key, + extra_view_keys=extra_view_keys, + preload_all_images=args.preload_all_images, + plot_rotation_6d=args.plot_rotation_6d, + ) + ) + except KeyboardInterrupt: + logger.info("Stopped by user.") + return 130 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/fake_inference.py b/scripts/fake_inference.py index 428aec7a..a8516a3f 100644 --- a/scripts/fake_inference.py +++ b/scripts/fake_inference.py @@ -1,85 +1,103 @@ -import torch -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction - -model_path = "/path/to/model" -model = Qwen2_5_VLMoEForAction.from_pretrained(model_path) -model.eval() - -# Gen Fake data -batch_size = 1 -seq_length = 50 - -torch.manual_seed(0) -fake_input_ids = torch.randint( - 0, len(model.processor.tokenizer), (batch_size, seq_length), dtype=torch.long -) -fake_attention_mask = torch.ones((batch_size, seq_length), dtype=torch.long) -fake_moe_token_types = torch.zeros((batch_size, seq_length), dtype=torch.long) -fake_position_ids = ( - torch.arange(seq_length, dtype=torch.long).unsqueeze(0).expand(batch_size, -1) -) -fake_proprioception = torch.randn((batch_size, 1, 20), dtype=torch.float32) -fake_agent_pos_mask = torch.ones((batch_size, 1, 20), dtype=torch.float32) -fake_dof_mask = torch.ones((batch_size, 32, 20), dtype=torch.float32) -fake_dataset_names = ["x2_normal"] - - -device = "cuda" - -model = model.to(device) -model = model.bfloat16() - -fake_input_ids = fake_input_ids.to(device) -fake_attention_mask = fake_attention_mask.to(device) -fake_moe_token_types = fake_moe_token_types.to(device) -fake_position_ids = fake_position_ids.to(device) -fake_proprioception = fake_proprioception.to(device).bfloat16() -fake_agent_pos_mask = fake_agent_pos_mask.to(device).bfloat16() -fake_dof_mask = fake_dof_mask.to(device).bfloat16() - -try: - with torch.no_grad(): - outputs = model( - input_ids=fake_input_ids, - attention_mask=fake_attention_mask, - moe_token_types=fake_moe_token_types, - position_ids=fake_position_ids, - proprioception=fake_proprioception, - agent_pos_mask=fake_agent_pos_mask, - dof_mask=fake_dof_mask, - dataset_names=fake_dataset_names, - mode="validate", - ) - - print("✅ Fake inference test successful!") - print(f"Output logits shape: {outputs.logits.shape}") - print(f"Output logits dtype: {outputs.logits.dtype}") - print(f"Output logits device: {outputs.logits.device}") - - # Check if output is reasonable - if outputs.logits.shape == (batch_size, seq_length, model.config.vocab_size): - print("✅ Output shape correct") - else: - print("❌ Output shape incorrect") - - if not torch.isnan(outputs.logits).any(): - print("✅ Output contains no NaN values") - else: - print("❌ Output contains NaN values") - - if not torch.isinf(outputs.logits).any(): - print("✅ Output contains no infinity values") - else: - print("❌ Output contains infinity values") - - print("Output logits statistics:") - print(f" Min value: {outputs.logits.min().item():.4f}") - print(f" Max value: {outputs.logits.max().item():.4f}") - print(f" Mean: {outputs.logits.mean().item():.4f}") - print(f" Standard deviation: {outputs.logits.std().item():.4f}") - -except Exception as e: - print(f"❌ Fake inference test failed: {e}") - import traceback - - traceback.print_exc() +#!/usr/bin/env python3 +"""Run one Wall-X VLA inference pass through the harrix adapter. + +This is a lightweight smoke test for the inference path. It builds a synthetic +LIBERO-style observation, loads the checkpoint through harrix, and prints the +predicted action chunk shape. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import numpy as np + + +def _ensure_local_harrix_on_path() -> None: + repo_root = Path(__file__).resolve().parents[1] + harrix_python = repo_root / "third_party" / "harrix" / "python" + if harrix_python.is_dir(): + sys.path.insert(0, str(harrix_python)) + + +def _build_fake_observation(seed: int, image_size: int) -> dict: + rng = np.random.default_rng(seed) + return { + "eef_pos": rng.normal(size=(3,)).astype(np.float32), + "eef_axisangle": rng.normal(size=(3,)).astype(np.float32), + "gripper": rng.normal(size=(1,)).astype(np.float32), + "face_view": rng.integers(0, 256, (image_size, image_size, 3), dtype=np.uint8), + "wrist_view": rng.integers(0, 256, (image_size, image_size, 3), dtype=np.uint8), + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--checkpoint-path", required=True, help="Checkpoint directory." + ) + parser.add_argument( + "--train-config-path", + default=None, + help="Optional training config path. Defaults to config.yml/config.yaml next to the checkpoint.", + ) + parser.add_argument("--norm-key", default="libero_all") + parser.add_argument("--architecture", default="qwen2_5") + parser.add_argument("--action-mode", default="flow") + parser.add_argument( + "--cam-names", + nargs="+", + default=["face_view", "right_wrist_view"], + help="Camera names expected by the checkpoint.", + ) + parser.add_argument("--action-horizon", type=int, default=None) + parser.add_argument("--instruction", default="pick up the object") + parser.add_argument("--image-size", type=int, default=128) + parser.add_argument("--seed", type=int, default=0) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + _ensure_local_harrix_on_path() + + from wall_x._vendor.harrix.adapters.registry import build_adapter + from wall_x._vendor.harrix.eval_config import ( + EvalConfig, + LiberoEnvParams, + autofill_from_checkpoint, + ) + + import wall_x._vendor.harrix.adapters # noqa: F401 register model adapters + + cfg = EvalConfig() + cfg.model.checkpoint_path = args.checkpoint_path + cfg.model.train_config_path = args.train_config_path + cfg.model.norm_key = args.norm_key + cfg.model.cam_names = list(args.cam_names) + cfg.model.action_horizon = args.action_horizon + cfg.model.architecture = args.architecture + cfg.model.action_mode = args.action_mode + cfg.env.libero = LiberoEnvParams(num_trials_per_task=1, task_indices=[0]) + cfg = autofill_from_checkpoint(cfg) + + adapter = build_adapter(cfg) + payload = { + "observation": _build_fake_observation(args.seed, args.image_size), + "instruction": args.instruction, + "noise": None, + } + actions = adapter.predict_batch([payload]) + action = np.asarray(actions[0]) + + print("Fake inference succeeded.") + print(f"action shape: {action.shape}") + print(f"action dtype: {action.dtype}") + print(f"action min/max: {float(action.min()):.6f} / {float(action.max()):.6f}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/infer_libero.py b/scripts/infer_libero.py index 4b107251..d96a3b61 100644 --- a/scripts/infer_libero.py +++ b/scripts/infer_libero.py @@ -1,222 +1,233 @@ -import argparse -import time -import os +#!/usr/bin/env python3 +"""Run LIBERO evaluation through harrix. -# from wall_x.utils.baseline_utils import check_baseline_dump, update_baseline -from wall_x.infer.utils_libero import set_seed_everywhere, TaskSuite, TASK_MAX_STEPS -from wall_x.infer.infer_config import InferConfig -from wall_x.infer.env_libero import LiberoRobotEnv +The script accepts either a full harrix EvalConfig YAML or a checkpoint path +plus common command-line overrides. It intentionally bypasses the legacy +Wall-X inference stack. +""" +from __future__ import annotations -if __name__ == "__main__": - args = argparse.ArgumentParser(description="Wall-X Libero evaluation script") - args.add_argument("--seed", type=int, default=42, help="Random seed") - args.add_argument("--id", type=int, default=None, help="Unique index id") - args.add_argument("--name", type=int, default=None, help="Launch command name") - args.add_argument( - "--baseline_path", type=str, default=None, help="Path to baseline record table" - ) - args.add_argument( - "--update_baseline", - type=bool, - default=False, - help="Whether to update baseline table", - ) - args.add_argument( - "--mode", type=str, default="flow", choices=["flow", "ar"], help="Running mode" - ) - args.add_argument( - "--checkpoint_path", type=str, required=True, help="Model checkpoint path" +import argparse +import sys +import tempfile +from pathlib import Path + +import yaml + + +LIBERO_DEFAULT_MAX_INFER_TIMES = { + "libero_spatial": 22, + "libero_object": 28, + "libero_goal": 30, + "libero_10": 52, + "libero_90": 40, +} + + +def _ensure_local_harrix_on_path() -> None: + repo_root = Path(__file__).resolve().parents[1] + harrix_python = repo_root / "third_party" / "harrix" / "python" + if harrix_python.is_dir(): + sys.path.insert(0, str(harrix_python)) + + +def _parse_task_indices(value: str | None) -> list[int] | None: + if value is None or value.strip() == "": + return None + return [int(x) for x in value.split(",") if x.strip()] + + +def _resolve_max_infer_times( + task_suite_name: str | None, max_infer_times: int | None +) -> int: + if max_infer_times is not None: + return max_infer_times + suite = task_suite_name or "libero_spatial" + return LIBERO_DEFAULT_MAX_INFER_TIMES.get(suite, 22) + + +def _load_or_build_raw_config(args: argparse.Namespace) -> dict: + if args.config is not None: + with open(args.config, "r") as f: + raw = yaml.safe_load(f) or {} + model = raw.setdefault("model", {}) + env = raw.setdefault("env", {}) + libero = env.setdefault("libero", {}) + runtime = raw.setdefault("runtime", {}) + debug = raw.setdefault("debug", {}) + if args.checkpoint_path is not None: + model["checkpoint_path"] = args.checkpoint_path + if args.train_config_path is not None: + model["train_config_path"] = args.train_config_path + task_indices = _parse_task_indices(args.task_indices) + if task_indices is not None: + libero["task_indices"] = task_indices + if args.max_infer_times is not None or libero.get("max_infer_times") is None: + libero["max_infer_times"] = _resolve_max_infer_times( + libero.get("task_suite_name", args.task_suite_name), + args.max_infer_times, + ) + if args.smoke: + libero["task_indices"] = [0] + libero["num_trials_per_task"] = 5 + runtime["num_workers"] = 1 + runtime["max_batch_size"] = 1 + if args.deterministic_model: + debug["deterministic_model"] = True + return raw + else: + if args.checkpoint_path is None: + raise ValueError("--checkpoint-path is required when --config is not set") + max_infer_times = _resolve_max_infer_times( + args.task_suite_name, args.max_infer_times + ) + raw = { + "model": { + "checkpoint_path": args.checkpoint_path, + "norm_key": args.norm_key, + "cam_names": args.cam_names, + "architecture": args.architecture, + "action_mode": args.action_mode, + }, + "env": { + "type": "libero", + "seed": args.seed, + "libero": { + "task_suite_name": args.task_suite_name, + "initial_states_path": args.initial_states_path, + "num_trials_per_task": args.num_trials_per_task, + "max_infer_times": max_infer_times, + "skip_intermediate_render": args.skip_intermediate_render, + }, + }, + "runtime": { + "num_workers": args.num_workers, + "max_batch_size": args.max_batch_size, + "ws_port": args.ws_port, + "log_dir": args.log_dir, + "driver_mode": args.driver_mode, + }, + "debug": {"deterministic_model": args.deterministic_model}, + } + + model = raw.setdefault("model", {}) + env = raw.setdefault("env", {}) + libero = env.setdefault("libero", {}) + runtime = raw.setdefault("runtime", {}) + debug = raw.setdefault("debug", {}) + + if args.checkpoint_path is not None: + model["checkpoint_path"] = args.checkpoint_path + if args.train_config_path is not None: + model["train_config_path"] = args.train_config_path + if args.norm_key is not None: + model["norm_key"] = args.norm_key + if args.cam_names is not None: + model["cam_names"] = args.cam_names + if args.action_horizon is not None: + model["action_horizon"] = args.action_horizon + if args.architecture is not None: + model["architecture"] = args.architecture + if args.action_mode is not None: + model["action_mode"] = args.action_mode + + env["type"] = "libero" + env["seed"] = args.seed + libero["task_suite_name"] = args.task_suite_name + libero["initial_states_path"] = args.initial_states_path + libero["num_trials_per_task"] = args.num_trials_per_task + libero["max_infer_times"] = _resolve_max_infer_times( + args.task_suite_name, args.max_infer_times ) - args.add_argument( - "--train_config_path", - type=str, - required=False, - default=None, - help="Path to training config .yml file", + libero["skip_intermediate_render"] = args.skip_intermediate_render + task_indices = _parse_task_indices(args.task_indices) + if task_indices is not None: + libero["task_indices"] = task_indices + if args.smoke: + libero["task_indices"] = [0] + libero["num_trials_per_task"] = 5 + runtime["num_workers"] = 1 + runtime["max_batch_size"] = 1 + + runtime["num_workers"] = ( + args.num_workers if not args.smoke else runtime["num_workers"] ) - args.add_argument( - "--norm_key", - type=str, - default="physical-intelligence/libero", - help="Key for normalization statistics", + runtime["max_batch_size"] = ( + args.max_batch_size if not args.smoke else runtime["max_batch_size"] ) - args.add_argument( - "--cam_names", - nargs="+", - default=["face_view", "right_wrist_view"], - help="List of camera names (e.g., --cam_names face_view right_wrist_view)", + runtime["ws_port"] = args.ws_port + runtime["log_dir"] = args.log_dir + runtime["driver_mode"] = args.driver_mode + debug["deterministic_model"] = args.deterministic_model + return raw + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--config", default=None, help="Optional harrix EvalConfig YAML." ) - args.add_argument( - "--task_suite_name", - type=str, - default=TaskSuite.LIBERO_SPATIAL, - choices=[e.value for e in TaskSuite], - help="Libero task suite to load", + parser.add_argument("--checkpoint-path", default=None) + parser.add_argument("--train-config-path", default=None) + parser.add_argument("--norm-key", default="libero_all") + parser.add_argument("--architecture", default="qwen2_5") + parser.add_argument("--action-mode", default="flow") + parser.add_argument( + "--cam-names", nargs="+", default=["face_view", "right_wrist_view"] ) - args.add_argument( - "--initial_states_path", - type=str, - default="DEFAULT", - help="Path to initial states .json file, or 'DEFAULT' to use default states.", + parser.add_argument("--action-horizon", type=int, default=None) + parser.add_argument("--task-suite-name", default="libero_spatial") + parser.add_argument("--initial-states-path", default="DEFAULT") + parser.add_argument("--num-trials-per-task", type=int, default=50) + parser.add_argument( + "--task-indices", default=None, help="Comma-separated task ids." ) - args.add_argument( - "--num_trials_per_task", + parser.add_argument( + "--max-infer-times", type=int, - default=50, - help="Number of evaluation episodes to run per task", + default=None, + help=( + "Number of model action chunks per episode. Defaults are suite-specific " + "and match the internal LIBERO evaluator." + ), ) - args.add_argument( - "--rollout_dir", - type=str, - default="./rollouts", - help="Directory to save rollout videos", + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument("--max-batch-size", type=int, default=1) + parser.add_argument("--ws-port", type=int, default=8765) + parser.add_argument("--log-dir", default="/tmp/harrix_libero_eval") + parser.add_argument("--driver-mode", choices=["in_process"], default="in_process") + parser.add_argument("--smoke", action="store_true") + parser.add_argument("--deterministic-model", action="store_true") + parser.add_argument( + "--skip-intermediate-render", + action=argparse.BooleanOptionalAction, + default=True, ) - args = args.parse_args() + return parser.parse_args() - print(f"Using random seed: {args.seed}") - set_seed_everywhere(args.seed) - print("Initializing InferConfig...") - if args.train_config_path is None: - args.train_config_path = os.path.join(args.checkpoint_path, "config.yml") +def main() -> int: + args = parse_args() + _ensure_local_harrix_on_path() - config = InferConfig( - checkpoint_path=args.checkpoint_path, - train_config_path=args.train_config_path, - norm_key=args.norm_key, - cam_names=args.cam_names, - ) - if args.mode == "flow": - config.action_horizon = config.train_config.get("data", {}).get( - "action_horizon_flow", 10 - ) - elif args.mode == "ar": - config.action_horizon = config.train_config.get("data", {}).get( - "action_horizon_ar", 10 - ) - else: - raise ValueError(f"Invalid mode: {args.mode}") - config.model_device = "cuda" + from wall_x._vendor.harrix.eval_config import autofill_from_checkpoint, load_eval_config - print("Initializing LiberoRobotEnv (Evaluator)...") + raw = _load_or_build_raw_config(args) - config.action_dim = 7 - config.pred_horizon = 10 + with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as f: + yaml.safe_dump(raw, f, sort_keys=False) + tmp_config = f.name - evaluator = LiberoRobotEnv( - config=config, - task_suite_name=args.task_suite_name, - initial_states_path=args.initial_states_path, - rollout_dir=args.rollout_dir, - seed=args.seed, - ) + cfg = autofill_from_checkpoint(load_eval_config(tmp_config)) + if cfg.runtime.driver_mode != "in_process": + raise ValueError("Only driver_mode='in_process' is supported") + from wall_x._vendor.harrix.drivers.inproc import run - print(f"\n{'='*20} Starting Evaluation {'='*20}") - print(f"Task suite: {args.task_suite_name}") - print(f"Number of tasks: {evaluator.num_tasks}") - print(f"Trials per task: {args.num_trials_per_task}") - print(f"Initial states: {args.initial_states_path}") - print(f"Videos will be saved to: {evaluator.rollout_dir}") - print(f"{'='*50}\n") - - total_successes = 0 - total_episodes_run = 0 - start_time = time.time() - - for task_id in range(evaluator.num_tasks): - task_successes = 0 - task_episodes_attempted = 0 - - libero_env_instance = None - task_desc = "" - initial_states = None - - max_infer_times = TASK_MAX_STEPS[args.task_suite_name] - print( - f"{args.task_suite_name} TASK_MAX_STEPS: {TASK_MAX_STEPS[args.task_suite_name]}" - ) - for ep_idx in range(args.num_trials_per_task): - print(f" > Running trial {ep_idx + 1} / {args.num_trials_per_task}...") - - try: - print(f"\nCreating environment for Task {task_id}...") - libero_env_instance, task_desc, initial_states = ( - evaluator.create_env_for_task(task_id) - ) - print( - f"--- Starting task {task_id + 1} / {evaluator.num_tasks}: {task_desc} ---" - ) - except Exception as e: - print( - f"\n[CRITICAL ERROR] Failed to create environment for task {task_id}: {e}. Skipping entire task." - ) - continue - - task_episodes_attempted += 1 - total_episodes_run += 1 - - success = False - try: - if args.mode == "flow": - success = evaluator.run_infer_flow_action( - env=libero_env_instance, - task_id=task_id, - task_desc=task_desc, - default_initial_states=initial_states, - episode_idx=ep_idx, - max_infer_times=max_infer_times, - ) - elif args.mode == "ar": - success = evaluator.run_infer_ar_action( - env=libero_env_instance, - task_id=task_id, - task_desc=task_desc, - default_initial_states=initial_states, - episode_idx=ep_idx, - max_infer_times=max_infer_times, - ) - except Exception as e: - print(f" [EXCEPTION] Episode run error: {e}") - - if success: - task_successes += 1 - total_successes += 1 - print(" > Trial result: SUCCESS") - else: - print(" > Trial result: FAILURE") - - if task_episodes_attempted > 0: - print( - f" > Task {task_id} current success rate: {task_successes / task_episodes_attempted * 100:.1f}% ({task_successes}/{task_episodes_attempted})" - ) - if total_episodes_run > 0: - print( - f" > Overall current success rate: {total_successes / total_episodes_run * 100:.1f}% ({total_successes}/{total_episodes_run})" - ) - - task_success_rate = ( - task_successes / task_episodes_attempted - if task_episodes_attempted > 0 - else 0 - ) - print(f"\n--- Task {task_id} ({task_desc}) Summary ---") - print( - f"Success rate: {task_success_rate * 100:.1f}% ({task_successes}/{task_episodes_attempted})" - ) - print(f"{'-'*40}\n") + run(cfg) + return 0 - end_time = time.time() - total_time = end_time - start_time - final_success_rate = ( - total_successes / total_episodes_run if total_episodes_run > 0 else 0 - ) - print(f"\n{'='*20} Final Evaluation Summary {'='*20}") - print(f"Total runtime: {total_time:.2f} seconds ({total_time / 60:.1f} minutes)") - print(f"Total trials run: {total_episodes_run}") - print(f"Total successes: {total_successes}") - print(f"Overall success rate: {final_success_rate * 100:.2f}%") - print(f"{'='*56}") - - print("Evaluation completed.") +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/infer_robochallenge.py b/scripts/infer_robochallenge.py deleted file mode 100644 index f25b3d32..00000000 --- a/scripts/infer_robochallenge.py +++ /dev/null @@ -1,1164 +0,0 @@ -import os -from scipy.fft import dct -from scipy.fft import idct -import yaml -import torch -import numpy as np -import dataclasses -import copy -import json -from PIL import Image -from safetensors.torch import load_file -from qwen_vl_utils.vision_process import smart_resize -from transformers import BatchFeature, AutoProcessor - -from wall_x.model.action_head import Normalizer -from wall_x.utils.constant import action_statistic_dof as default_action_statistic_dof -from numba import jit, prange - -try: - from spatial_tokenizer.spatial_tokenizer_kdisk import SpatialActionTokenizer -except ImportError: - SpatialActionTokenizer = None - -device = "cuda" - -dof_config = { - "follow_left_ee_cartesian_pos": 3, - "follow_left_ee_rotation": 3, - "follow_left_gripper": 1, - "follow_right_ee_cartesian_pos": 3, - "follow_right_ee_rotation": 3, - "follow_right_gripper": 1, - "head_actions": 2, - "height": 1, - "car_pose": 3, - "velocity_decomposed": 3, -} -_CAM_NAME_MAPPING = { - "face_view": "front view", - "left_wrist_view": "left wrist view", - "right_wrist_view": "right wrist view", - "move1_view": "move view", - "move2_view": "move view", - "wall_view": "wall view", - "top_view": "top view", - "side_view": "side view", - "global_view": "global view", -} -camera_to_view_mapping = { - "camera_front": "face_view", - "camera_left": "left_wrist_view", - "camera_right": "right_wrist_view", - "camera_side": "side_view", - "camera_global": "global_view", -} - -action_key_mapping = { - "follow_left_ee_cartesian_pos": "follow1_pos[:3]", - "follow_left_ee_rotation": "follow1_pos[3:6]", - "follow_left_gripper": "follow1_pos[6:7]", - "follow_right_ee_cartesian_pos": "follow2_pos[:3]", - "follow_right_ee_rotation": "follow2_pos[3:6]", - "follow_right_gripper": "follow2_pos[6:7]", - "head_actions": "head_pos", - "height": "lift", - "velocity_decomposed": "velocity_decomposed", - "follow_left_arm_joint_cur": "follow1_joints_cur[-1:]", - "follow_right_arm_joint_cur": "follow2_joints_cur[-1:]", - "follow_left_arm_joint_pos": "follow1_pos", - "follow_right_arm_joint_pos": "follow2_pos", -} - -dim_dof_config = { - "right_xyz": {"rpy": (0, 3), "so3": (0, 3)}, - "right_rot": {"rpy": (3, 6), "so3": (3, 9)}, - "right_gripper": {"rpy": (6, 7), "so3": (9, 10)}, -} -SINGLE_ARM_DIM = 7 - - -@jit(nopython=True, parallel=True) -def euler_to_matrix_zyx_batch_nb(eulers): - N = eulers.shape[0] - R = np.empty((N, 3, 3), dtype=np.float64) - for i in prange(N): - roll = eulers[i, 0] - pitch = eulers[i, 1] - yaw = eulers[i, 2] - - cy, sy = np.cos(yaw), np.sin(yaw) - cp, sp = np.cos(pitch), np.sin(pitch) - cr, sr = np.cos(roll), np.sin(roll) - - R[i, 0, 0] = cy * cp - R[i, 0, 1] = cy * sp * sr - sy * cr - R[i, 0, 2] = cy * sp * cr + sy * sr - - R[i, 1, 0] = sy * cp - R[i, 1, 1] = sy * sp * sr + cy * cr - R[i, 1, 2] = sy * sp * cr - cy * sr - - R[i, 2, 0] = -sp - R[i, 2, 1] = cp * sr - R[i, 2, 2] = cp * cr - return R - - -@jit(nopython=True, parallel=True) -def compose_state_and_delta_to_abs_rpy(delta, state): - """ - Input: - delta: (N,3) -> Δrpy(ZYX) or (N,6) -> Δ6D (first two rows flattened) - state: (3,) -> rpy(ZYX) or (6,) -> 6D (first two rows flattened) - Output: - abs_rpy: (N,3) Absolute pose rpy(ZYX, radians), normalized to (-π, π] - """ - if delta.shape[-1] == 3: - R_delta = euler_to_matrix_zyx_batch_nb(delta) # (N,3,3) - elif delta.shape[-1] == 6: - R_delta = so3_to_matrix_batch_nb(delta) # (N,3,3) - else: - raise ValueError(f"delta last dim must be 3 or 6, got {delta.shape[-1]}") - - if state.shape[-1] == 3: - R_state = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0] # (3,3) - elif state.shape[-1] == 6: - R_state = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] # (3,3) - else: - raise ValueError(f"state last dim must be 3 or 6, got {state.shape[-1]}") - - N = R_delta.shape[0] - R_abs = np.empty((N, 3, 3), dtype=np.float64) - - S00 = R_state[0, 0] - S01 = R_state[0, 1] - S02 = R_state[0, 2] - S10 = R_state[1, 0] - S11 = R_state[1, 1] - S12 = R_state[1, 2] - S20 = R_state[2, 0] - S21 = R_state[2, 1] - S22 = R_state[2, 2] - - for i in prange(N): - A00 = R_delta[i, 0, 0] - A01 = R_delta[i, 0, 1] - A02 = R_delta[i, 0, 2] - A10 = R_delta[i, 1, 0] - A11 = R_delta[i, 1, 1] - A12 = R_delta[i, 1, 2] - A20 = R_delta[i, 2, 0] - A21 = R_delta[i, 2, 1] - A22 = R_delta[i, 2, 2] - - R_abs[i, 0, 0] = A00 * S00 + A01 * S10 + A02 * S20 - R_abs[i, 0, 1] = A00 * S01 + A01 * S11 + A02 * S21 - R_abs[i, 0, 2] = A00 * S02 + A01 * S12 + A02 * S22 - - R_abs[i, 1, 0] = A10 * S00 + A11 * S10 + A12 * S20 - R_abs[i, 1, 1] = A10 * S01 + A11 * S11 + A12 * S21 - R_abs[i, 1, 2] = A10 * S02 + A11 * S12 + A12 * S22 - - R_abs[i, 2, 0] = A20 * S00 + A21 * S10 + A22 * S20 - R_abs[i, 2, 1] = A20 * S01 + A21 * S11 + A22 * S21 - R_abs[i, 2, 2] = A20 * S02 + A21 * S12 + A22 * S22 - - -@jit(nopython=True, parallel=True) -def so3_to_matrix_batch_nb(batch_so3): - N = batch_so3.shape[0] - R_all = np.empty((N, 3, 3), dtype=np.float64) - eps = 1e-12 - for i in prange(N): - r1x, r1y, r1z = batch_so3[i, 0], batch_so3[i, 1], batch_so3[i, 2] - r2x, r2y, r2z = batch_so3[i, 3], batch_so3[i, 4], batch_so3[i, 5] - - # normalize r1 - n1 = np.sqrt(r1x * r1x + r1y * r1y + r1z * r1z) + eps - r1x /= n1 - r1y /= n1 - r1z /= n1 - - # orthogonalize r2 to r1, then normalize - dot12 = r1x * r2x + r1y * r2y + r1z * r2z - r2x -= dot12 * r1x - r2y -= dot12 * r1y - r2z -= dot12 * r1z - n2 = np.sqrt(r2x * r2x + r2y * r2y + r2z * r2z) + eps - r2x /= n2 - r2y /= n2 - r2z /= n2 - - # r3 = r1 x r2 - r3x = r1y * r2z - r1z * r2y - r3y = r1z * r2x - r1x * r2z - r3z = r1x * r2y - r1y * r2x - - R_all[i, 0, 0] = r1x - R_all[i, 0, 1] = r1y - R_all[i, 0, 2] = r1z - R_all[i, 1, 0] = r2x - R_all[i, 1, 1] = r2y - R_all[i, 1, 2] = r2z - R_all[i, 2, 0] = r3x - R_all[i, 2, 1] = r3y - R_all[i, 2, 2] = r3z - return R_all - - -@jit(nopython=True, parallel=True) -def matrix_to_euler_zyx_batch_nb(Rs): - """ - R = Rz(yaw) * Ry(pitch) * Rx(roll) - extract: - pitch = asin(-R[2,0]) - roll = atan2(R[2,1], R[2,2]) - yaw = atan2(R[1,0], R[0,0]) - """ - N = Rs.shape[0] - eulers = np.empty((N, 3), dtype=np.float64) - for i in prange(N): - r00 = Rs[i, 0, 0] - # r01 = Rs[i, 0, 1] - # r02 = Rs[i, 0, 2] - r10 = Rs[i, 1, 0] - # r11 = Rs[i, 1, 1] - # r12 = Rs[i, 1, 2] - r20 = Rs[i, 2, 0] - r21 = Rs[i, 2, 1] - r22 = Rs[i, 2, 2] - - x = -r20 - if x > 1.0: - x = 1.0 - elif x < -1.0: - x = -1.0 - - pitch = np.arcsin(x) - roll = np.arctan2(r21, r22) - yaw = np.arctan2(r10, r00) - - eulers[i, 0] = roll - eulers[i, 1] = pitch - eulers[i, 2] = yaw - return eulers - - -@jit(nopython=True, parallel=True) -def canonicalize_euler_zyx_batch_nb(rpy_batch): - """ - Batch ZYX Euler Angle Normalization (Parallel Version) - Input: rpy_batch (N, 3) [roll, pitch, yaw] (radians) - Output: out (N, 3) Constrained to the same branch with each component in (-π, π] - Rules: - 1) First, wrap each component to (-π, π] - 2) If p > π/2: p = π - p; r += π; y += π - If p <= -π/2: p = -π - p; r += π; y += π - 3) Finally, wrap each component to (-π, π] again. - """ - N = rpy_batch.shape[0] - out = np.empty_like(rpy_batch) - two_pi = 2.0 * np.pi - - for i in prange(N): - r = rpy_batch[i, 0] - p = rpy_batch[i, 1] - y = rpy_batch[i, 2] - - r = (r + np.pi) % two_pi - np.pi - p = (p + np.pi) % two_pi - np.pi - y = (y + np.pi) % two_pi - np.pi - - if p > np.pi / 2.0: - p = np.pi - p - r = r + np.pi - y = y + np.pi - elif p <= -np.pi / 2.0: - p = -np.pi - p - r = r + np.pi - y = y + np.pi - - r = (r + np.pi) % two_pi - np.pi - p = (p + np.pi) % two_pi - np.pi - y = (y + np.pi) % two_pi - np.pi - - out[i, 0] = r - out[i, 1] = p - out[i, 2] = y - - return out - - -def so3_to_euler_zyx_batch_nb(batch_so3): - matrix = so3_to_matrix_batch_nb(batch_so3) - eulers = matrix_to_euler_zyx_batch_nb(matrix) - return canonicalize_euler_zyx_batch_nb(eulers) - - -def update_model_config(train_config, model_config): - model_config.use_state_string_representation = train_config["data"].get( - "use_state_string_representation", False - ) - model_config.flow_loss_weight = train_config.get("flow_loss_weight", 1.0) - - model_config.dof_config = train_config["dof_config"] - model_config.agent_pos_config = train_config["agent_pos_config"] - - model_config.action_horizon_flow = train_config["data"].get( - "action_horizon_flow", 32 - ) - - if train_config.get("_attn_implementation", None) is not None: - model_config._attn_implementation = train_config["_attn_implementation"] - - return model_config - - -def move_to_cuda(obj, device=device): - 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 - - -def extract_components(data, config, is_rpy): - mode = "rpy" if is_rpy else "so3" - result = {} - for key, slice_config in config.items(): - slice_range = slice_config[mode] - result[key] = data[:, slice_range[0] : slice_range[1]] - return result - - -@dataclasses.dataclass -class WallxInferArgs: - config_path: str | None = None - checkpoint_path: str | None = None - action_mode: str = "diffusion" # "ar" or "diffusion" - - max_time_step: int = 1000 - action_start_ratio: float = 0 - action_end_ratio: float = 0.6 - - model_action_dim: int = 14 - action_horizon: int = 32 - - action_dim: int = 14 - - interpolate_action: bool = False - interpolate_multiplier: int = 1 - turtle_as_desktop: bool = False - generate_subtask: bool = False - subtask_interval: int = 0 - with_cur: bool = False - state_str: bool = True ### NOTE - wostate: bool = False - delta_action: bool = False - state_rpy: bool = True - action_rpy: bool = True - - dataset_name: str = "robochallenge_aloha" - use_hard_prompt: bool = True - dct_scale: float = -1 - - -class WallxModelWrapper: - def __init__(self, args: WallxInferArgs): - self.args = args - self.get_model_and_processor() - self.action_predict_mode = "ar" if args.action_mode == "ar" else "diffusion" - print("action_predict_mode", self.action_predict_mode, flush=True) - - def get_model_and_processor(self): - if self.args.config_path is None: - self.args.config_path = os.path.join( - self.args.checkpoint_path, "config.yml" - ) - with open(self.args.config_path, "r") as f: - config = yaml.load(f, Loader=yaml.FullLoader) - self.config = config - self.dof_config = config["dof_config"] - self.agent_pos_config = config["agent_pos_config"] - self.obs_action_keys = [ - "follow_left_ee_cartesian_pos", - "follow_left_ee_rotation", - "follow_left_gripper", - "follow_right_ee_cartesian_pos", - "follow_right_ee_rotation", - "follow_right_gripper", - ] - print("obs_action_keys", self.obs_action_keys, flush=True) - self.action_tokenizer_type = config.get("action_tokenizer_type", None) - config_path = config["qwen_vl_act_config_path"] - self.processor = AutoProcessor.from_pretrained( - config["processor_path"], use_fast=True - ) - self.processor.tokenizer.padding_side = "left" - new_tokens = ["<|propri|>", "<|action|>"] - - # load fast tokenizer - self.action_tokenizer_type = config.get("action_tokenizer_type", None) - print("self.action_tokenizer_type", self.action_tokenizer_type, flush=True) - self.action_tokenizer = None - self.action_mapper = None - if self.action_tokenizer_type: - # fast - if self.action_tokenizer_type == "fast": - print("Using fast tokenizer") - self.action_tokenizer = AutoProcessor.from_pretrained( - config["action_tokenizer_path"], trust_remote_code=True - ) - elif self.action_tokenizer_type == "spatialvla": - print("Using spatialvla tokenizer") - assert ( - SpatialActionTokenizer is not None - ), "SpatialActionTokenizer is not installed" - self.action_tokenizer = SpatialActionTokenizer() - else: - raise ValueError( - f"Unsupported action tokenizer type: {self.action_tokenizer_type}" - ) - new_tokens += [ - f"<|action_token_{i}|>" for i in range(self.action_tokenizer.vocab_size) - ] - - # num_added_tokens = self.processor.tokenizer.add_tokens(new_tokens) - - # define action mapper - if self.action_tokenizer_type: - self.action_mapper = {} - for i in range(self.action_tokenizer.vocab_size): - token = f"<|action_token_{i}|>" - token_id = self.processor.tokenizer.convert_tokens_to_ids(token) - self.action_mapper[token_id] = i - - # action & propri normalizer - self._register_normalizers() - - model_type = config["model_type"] - - if model_type == "qwen2_5": - print("Using qwen2_5 model as base model") - from wall_x.model.qwen2_5_based import ( - Qwen2_5_VLMoEForAction, - Qwen2_5_VLConfig, - ) - - ModelClass = Qwen2_5_VLMoEForAction - ConfigClass = Qwen2_5_VLConfig - - model_config = ConfigClass.from_pretrained(config_path) - model_config = update_model_config(config, model_config) - - print("model_config", model_config, flush=True) - - # if self.args.action_mode == "ar": - # model_config._attn_implementation = "flash_attention_2" - # else - model_config._attn_implementation = "sdpa" - model_config.vision_config._attn_implementation = "flash_attention_2" - - if model_config.model_type == "qwen3_vl": - self.MAX_PIXELS = 16384 * 32 * 32 - self.MIN_PIXELS = 4 * 32 * 32 - self.IMAGE_FACTOR = 32 - elif model_config.model_type == "qwen2_5_vl": - self.MAX_PIXELS = 16384 * 28 * 28 - self.MIN_PIXELS = 4 * 28 * 28 - self.IMAGE_FACTOR = 28 - - model = ModelClass( - model_config, - self.action_tokenizer_type, - self.processor, - self.action_tokenizer, - self.action_mapper, - ) - model.resize_token_embeddings(len(self.processor.tokenizer)) - state_dict = load_file( - self.args.checkpoint_path + "/model.safetensors", device="cpu" - ) - if os.path.exists(os.path.join(self.args.checkpoint_path, "global_step.pth")): - global_step = torch.load( - os.path.join(self.args.checkpoint_path, "global_step.pth") - )["global_step"] - print("Loaded global step:", global_step) - msg = model.load_state_dict(state_dict, strict=False) - print(msg) - model.eval() - model.set_normalizer( - copy.deepcopy(self.normalizer_action), copy.deepcopy(self.normalizer_propri) - ) - model.to(device) - model.to_bfloat16_for_selected_params() - - self.model = model - print("self.args.dataset_name", self.args.dataset_name, flush=True) - print( - "normalizer_action min", - self.normalizer_action.min.__getattr__(self.args.dataset_name), - flush=True, - ) - print( - "normalizer_action delta", - self.normalizer_action.delta.__getattr__(self.args.dataset_name), - flush=True, - ) - - def _register_normalizers(self): - if self.config.get("customized_action_statistic_dof", None): - action_statistic_dof = json.load( - open(self.config["customized_action_statistic_dof"], "r") - ) - else: - action_statistic_dof = default_action_statistic_dof - - if os.path.exists(self.args.checkpoint_path + "/normalizer_action.pth"): - print( - "Loading normalizer_action from checkpoint", - self.args.checkpoint_path + "/normalizer_action.pth", - flush=True, - ) - self.normalizer_action = Normalizer.from_ckpt( - self.args.checkpoint_path + "/normalizer_action.pth" - ) - else: - self.normalizer_action = Normalizer( - action_statistic_dof, - self.config["dof_config"], - min_key=self.config.get("min_key", "min"), - delta_key=self.config.get("delta_key", "delta"), - ) - - # print("action_statistic_dof",action_statistic_dof) - - if os.path.exists(self.args.checkpoint_path + "/normalizer_propri.pth"): - print( - "Loading normalizer_propri from checkpoint", - self.args.checkpoint_path + "/normalizer_propri.pth", - flush=True, - ) - self.normalizer_propri = Normalizer.from_ckpt( - self.args.checkpoint_path + "/normalizer_propri.pth" - ) - else: - self.normalizer_propri = Normalizer( - action_statistic_dof, - self.config["agent_pos_config"], - min_key=self.config.get("min_key", "min"), - delta_key=self.config.get("delta_key", "delta"), - ) - - print("self.args.dataset_name", self.args.dataset_name, flush=True) - print( - "normalizer_propri min", - self.normalizer_propri.min.__getattr__(self.args.dataset_name), - flush=True, - ) - print( - "normalizer_propri delta", - self.normalizer_propri.delta.__getattr__(self.args.dataset_name), - flush=True, - ) - print( - "normalizer_action min", - self.normalizer_action.min.__getattr__(self.args.dataset_name), - flush=True, - ) - print( - "normalizer_action delta", - self.normalizer_action.delta.__getattr__(self.args.dataset_name), - flush=True, - ) - - def get_text_ar(self, instruction, camera_names, norm_state=None, state_mask=None): - role_start_symbol = "<|im_start|>" - role_end_symbol = "<|im_end|>" - vision_start_symbol = "<|vision_start|>" - vision_end_symbol = "<|vision_end|>" - image_pad_symbol = "<|image_pad|>" - propri_symbol = "<|propri|>" - - prologue = f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n" - user_request = f"{role_start_symbol}user\nObservation:" - print("camera_names", camera_names, flush=True) - for cam_name in camera_names: - user_request += f" {_CAM_NAME_MAPPING[cam_name]}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}" - user_request += "\nInstruction:" - if self.args.state_str: - assert norm_state is not None - if isinstance(norm_state, torch.Tensor): - if state_mask is not None: - if isinstance(state_mask, torch.Tensor): - mask_1d = state_mask[0, 0].to( - dtype=torch.bool, device=norm_state.device - ) - else: - mask_1d = torch.as_tensor(state_mask, device=norm_state.device)[ - 0, 0 - ].to(dtype=torch.bool) - - norm_state = norm_state[..., mask_1d] - - norm_state = norm_state.detach().cpu().numpy() - print("norm_state", norm_state, flush=True) - discretized_state = ( - np.digitize(norm_state, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1 - ) - propri = " ".join(map(str, discretized_state[0, 0])) - elif self.args.wostate: - propri = "" - else: - propri = propri_symbol - text_prompt = ( - f"\nPredict the next action in robot action.\nProprioception: {propri}\n" - ) - user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" - assistant_message = f"{role_start_symbol}assistant\n" - text = prologue + user_message + assistant_message - - return text - - def get_text_flow( - self, - instruction, - camera_names, - action_chunk_size, - norm_state=None, - state_mask=None, - ): - role_start_symbol = "<|im_start|>" - role_end_symbol = "<|im_end|>" - vision_start_symbol = "<|vision_start|>" - vision_end_symbol = "<|vision_end|>" - image_pad_symbol = "<|image_pad|>" - propri_symbol = "<|propri|>" - action_symbol = "<|action|>" - action_space = "Rel EEF" if self.args.delta_action else "Abs EEF" - _camera = ", ".join([_CAM_NAME_MAPPING[cam_name] for cam_name in camera_names]) - prologue = f"<|im_start|>system\nYou are an embodied vision-language-action (VLA) model controlling the robot with language instructions.\n Embodiment: {self.args.dataset_name.split('_')[-1]}\n Camera Setup: {_camera},\n Frequency: 32HZ\n Action Space: {action_space}\n<|im_end|>\n" - - user_request = f"{role_start_symbol}user\nObservation:" - print("camera_names", camera_names, flush=True) - for cam_name in camera_names: - user_request += f" {_CAM_NAME_MAPPING[cam_name]}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}" - user_request += "\nInstruction:" - if self.args.state_str: - assert norm_state is not None - if isinstance(norm_state, torch.Tensor): - if state_mask is not None: - if isinstance(state_mask, torch.Tensor): - mask_1d = state_mask[0, 0].to( - dtype=torch.bool, device=norm_state.device - ) - else: - mask_1d = torch.as_tensor(state_mask, device=norm_state.device)[ - 0, 0 - ].to(dtype=torch.bool) - norm_state = norm_state[..., mask_1d] - norm_state = norm_state.detach().cpu().numpy() - discretized_state = ( - np.digitize(norm_state, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1 - ) - propri = " ".join(map(str, discretized_state[0, 0])) - elif self.args.wostate: - propri = "" - else: - propri = propri_symbol - text_prompt = ( - f"\nPredict the next action in robot action.\nProprioception: {propri}\n" - ) - user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" - assistant_message = f"{role_start_symbol}assistant\n" - action = f"{action_symbol * action_chunk_size}" - text = prologue + user_message + assistant_message + action - - return text - - def resize_images(self, observation): - image_inputs = [] - view_candidates = [ - "face_view", - "left_wrist_view", - "right_wrist_view", - "side_view", - "global_view", - ] - for key in observation.keys(): - if key not in view_candidates: - print("!!! key not in view_candidates", key, flush=True) - continue - # 1. Get the original image - current_obs = observation[key] - img_pil = Image.fromarray(current_obs) - orig_width, orig_height = img_pil.size - - # 2. Apply resolution limits (if the configuration is not -1) - target_size = 256 - if target_size != -1: - # Logic for maintaining aspect ratio constraints - if orig_width > orig_height: - new_width = target_size - new_height = int(target_size * orig_height / orig_width) - else: - new_height = target_size - new_width = int(target_size * orig_width / orig_height) - img_pil = img_pil.resize((new_width, new_height)) - - # 3. Apply intelligent scaling - current_width, current_height = img_pil.size - resized_height, resized_width = smart_resize( - current_height, - current_width, - factor=self.IMAGE_FACTOR, - min_pixels=self.MIN_PIXELS, - max_pixels=self.MAX_PIXELS, - ) - resized_img = img_pil.resize((resized_width, resized_height)) - print("resized_img", resized_img.size, flush=True) - - image_inputs.append(resized_img) - - return image_inputs - - def _construct_input( - self, - observation, - instruction, - camera_names, - valid_action_dim=7, - mode="ar", - single_image=False, - ): - additional_inputs = {} - - agent_pos = torch.from_numpy(observation["agent_pos"]) - agent_pos_mask = torch.from_numpy(observation["agent_pos_mask"]) - dof_mask = torch.from_numpy(observation["dof_mask"]) - additional_inputs["dof_mask"] = dof_mask - print("before normalizing agent_pos", agent_pos, flush=True) - - if self.normalizer_propri is not None: - - agent_pos = self.normalizer_propri.normalize_data( - agent_pos, [self.args.dataset_name] - ) - additional_inputs["proprioception"] = agent_pos - additional_inputs["agent_pos_mask"] = agent_pos_mask - - print( - f"normalizing agent_pos: {agent_pos}, {self.args.dataset_name}", - flush=True, - ) - print("agent_pos_mask", agent_pos_mask, flush=True) - print("dof_mask", dof_mask[0, 0], flush=True) - if mode == "ar": - text = self.get_text_ar( - instruction, camera_names, agent_pos, agent_pos_mask - ) - elif mode == "diffusion": - text = self.get_text_flow( - instruction, - camera_names, - self.args.action_horizon, - agent_pos, - agent_pos_mask, - ) - elif mode == "subtask": - text = self.get_text_subtask(instruction, single_view=single_image) - else: - raise ValueError(f"Invalid mode: {mode}") - text = [text] - - image_inputs = self.resize_images(observation) - if single_image: - image_inputs = [ - image_inputs[0] - ] # single view subtask/vqa use head view only - image_inputs = self.processor.image_processor( - images=image_inputs, videos=None, return_tensors="pt" - ) - image_grid_thw = image_inputs["image_grid_thw"] - # Processing image placeholder tokens in the text - if image_grid_thw is not None: - merge_length = self.processor.image_processor.merge_size**2 - index = 0 - for i in range(len(text)): - while "<|image_pad|>" in text[i]: - # Replace image placeholders with actual quantities. - text[i] = text[i].replace( - "<|image_pad|>", - "<|placeholder|>" - * (image_grid_thw[index].prod() // merge_length), - 1, - ) - index += 1 - text[i] = text[i].replace("<|placeholder|>", "<|image_pad|>") - - text_inputs = self.processor.tokenizer( - text, return_tensors="pt", padding=True, truncation=True, max_length=1024 - ) - inputs = BatchFeature(data={**text_inputs, **image_inputs}) - - action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") - additional_inputs["moe_token_types"] = inputs.input_ids == action_token_id - additional_inputs["dataset_names"] = [self.args.dataset_name] - - inputs.update(additional_inputs) - inputs = move_to_cuda(inputs, device) - print("inputs", inputs.keys(), flush=True) - for k in inputs.keys(): - if isinstance(inputs[k], torch.Tensor): - print(k, inputs[k].shape, flush=True) - return inputs - - def preprocess(self, state, views, valid_action_dim): - # print("state",state, flush=True) - # state: dict, keys: follow1_pos, follow2_pos - model_action_dim = sum(self.dof_config.values()) - - if valid_action_dim not in (SINGLE_ARM_DIM, 2 * SINGLE_ARM_DIM): - raise ValueError( - f"Invalid valid_action_dim: {valid_action_dim}, expect 7 or 14" - ) - - # 1) First, prepare a container of size (1, 1, D), where D = model_action_dim. - agent_data = np.zeros((1, 1, model_action_dim), dtype=np.float32) - agent_pos_mask = np.zeros((1, 1, model_action_dim), dtype=np.float32) - dof_mask = np.zeros( - (1, self.args.action_horizon, model_action_dim), dtype=np.float32 - ) - - # 2) Determine the interval to be filled [start:end) - if valid_action_dim == SINGLE_ARM_DIM: - start = 0 if model_action_dim == SINGLE_ARM_DIM else SINGLE_ARM_DIM - end = start + SINGLE_ARM_DIM - - if end > model_action_dim: - raise ValueError( - f"model_action_dim={model_action_dim} too small for valid_action_dim=7 " - f"(need end={end})" - ) - - follow2 = np.asarray(state["follow2_pos"], dtype=np.float32).reshape( - 1, 1, SINGLE_ARM_DIM - ) - agent_data[:, :, start:end] = follow2 - agent_pos_mask[:, :, start:end] = 1 - dof_mask[:, :, start:end] = 1 - - else: # valid_action_dim == 14 - end = 2 * SINGLE_ARM_DIM - if end > model_action_dim: - raise ValueError( - f"model_action_dim={model_action_dim} too small for valid_action_dim=14" - ) - - follow1 = np.asarray(state["follow1_pos"], dtype=np.float32).reshape( - 1, 1, SINGLE_ARM_DIM - ) - follow2 = np.asarray(state["follow2_pos"], dtype=np.float32).reshape( - 1, 1, SINGLE_ARM_DIM - ) - agent_data[:, :, :end] = np.concatenate([follow1, follow2], axis=-1) - agent_pos_mask[:, :, :end] = 1 - dof_mask[:, :, :end] = 1 - - observation = { - camera_to_view_mapping[key]: views[key][0] for key in views.keys() - } - observation["agent_pos"] = agent_data - observation["agent_pos_mask"] = agent_pos_mask - observation["dof_mask"] = dof_mask - return observation - - def model_output_process(self, action_pred, state): - - if not self.args.delta_action and self.args.action_rpy: - return action_pred - - pred_components = extract_components( - action_pred, dim_dof_config, self.args.action_rpy - ) - - pred_right_xyz = pred_components["right_xyz"] - pred_right_rot = pred_components["right_rot"] - pred_right_gripper = pred_components["right_gripper"] - if "left_xyz" in pred_components: - pred_left_xyz = pred_components["left_xyz"] - pred_left_rot = pred_components["left_rot"] - pred_left_gripper = pred_components["left_gripper"] - else: - pred_left_xyz = np.zeros((self.args.action_horizon, 3)) - pred_left_rot = np.zeros((self.args.action_horizon, 3)) - pred_left_gripper = np.zeros((self.args.action_horizon, 1)) - - post_action_pred = np.zeros((self.args.action_horizon, self.args.action_dim)) - if self.args.delta_action: - assert ( - self.args.action_dim == 14 - ), "Delta robot support and testing are not yet available." - - state_components = extract_components( - state, dim_dof_config, self.args.state_rpy - ) - left_xyz = state_components["left_xyz"] - left_rot = state_components["left_rot"] - right_xyz = state_components["right_xyz"] - right_rot = state_components["right_rot"] - - post_action_pred[:, :3] = pred_left_xyz + left_xyz - post_action_pred[:, 3:6] = compose_state_and_delta_to_abs_rpy( - pred_left_rot, left_rot[0] - ) - post_action_pred[:, 6:7] = pred_left_gripper - post_action_pred[:, 7:10] = pred_right_xyz + right_xyz - post_action_pred[:, 10:13] = compose_state_and_delta_to_abs_rpy( - pred_right_rot, right_rot[0] - ) - post_action_pred[:, 13:14] = pred_right_gripper - - elif not self.args.action_rpy: - post_action_pred[:, :3] = pred_left_xyz - post_action_pred[:, 3:6] = so3_to_euler_zyx_batch_nb(pred_left_rot) - post_action_pred[:, 6:7] = pred_left_gripper - post_action_pred[:, 7:10] = pred_right_xyz - post_action_pred[:, 10:13] = so3_to_euler_zyx_batch_nb(pred_right_rot) - post_action_pred[:, 13:14] = pred_right_gripper - else: - post_action_pred = action_pred - return post_action_pred - - def postprocess(self, action_pred, interpolate_multiplier=None): - - if interpolate_multiplier is None: - interpolate_multiplier = self.args.interpolate_multiplier - - if isinstance(action_pred, torch.Tensor): - action_pred = action_pred.to(torch.float32).cpu().squeeze(0).numpy() - left_action_pred = action_pred[:, :7] # (32, 7) - right_action_pred = action_pred[:, 7:14] # (32, 7) - - start_frame = int(self.args.action_start_ratio * len(left_action_pred)) - end_frame = int(self.args.action_end_ratio * len(left_action_pred)) - left_action_pred = left_action_pred[start_frame:end_frame] - right_action_pred = right_action_pred[start_frame:end_frame] - - print("left_action_pred", left_action_pred[-1], flush=True) - print("right_action_pred", right_action_pred[-1], flush=True) - - left_action_pred = left_action_pred.tolist() - right_action_pred = right_action_pred.tolist() - - serialized_actions = { - "follow1_pos": left_action_pred, - "follow2_pos": right_action_pred, - ## for joint-control - # "follow1_joints":left_action_pred, - # "follow2_joints":right_action_pred, - } - - return serialized_actions - - def predict_action_rtc( - self, - state, - views, - instruction=None, - valid_action_dim=7, - update_subtask=False, - action_predict_mode=None, - ): - if action_predict_mode is not None: - self.action_predict_mode = action_predict_mode - - observation = self.preprocess(state, views, valid_action_dim) - print("use instruction", instruction, flush=True) - camera_names = [camera_to_view_mapping[key] for key in views.keys()] - # camera_names = ["right_wrist_view", "global_view", "side_view"] - print("mode:", self.action_predict_mode, flush=True) - inputs = self._construct_input( - observation, - instruction, - camera_names=camera_names, - valid_action_dim=valid_action_dim, - mode=self.action_predict_mode, - ) - model_action_dim = sum(self.dof_config.values()) - padding = torch.zeros((1, model_action_dim)) - norm_padding = self.normalizer_action.normalize_data( - padding, [self.args.dataset_name] - ) - inputs["padding_action"] = norm_padding - inputs = move_to_cuda(inputs, device="cuda:0") - agent_data = observation["agent_pos"][..., : self.args.model_action_dim][0] - print("before generate_flow_action", flush=True) - print(inputs.keys(), flush=True) - print(inputs["dataset_names"], flush=True) - print(self.processor.tokenizer.decode(inputs["input_ids"][0]), flush=True) - print(inputs["agent_pos_mask"][0], flush=True) - print(inputs["dof_mask"][0, 0], flush=True) - if self.action_predict_mode == "ar": - action_pred = self.generate_ar_action(inputs) - else: - action_pred = self.generate_flow_action(inputs) - print("after generate_flow_action", flush=True) - if isinstance(action_pred, torch.Tensor): - action_pred = action_pred.float().cpu().squeeze(0).numpy() - - if action_pred is None: - return None - - if self.args.dct_scale > 0: - scale = self.args.dct_scale - dct_coeff = dct(action_pred, axis=0, norm="ortho") - dct_coeff = np.around(dct_coeff * scale) - action_pred = idct(dct_coeff / scale, axis=0, norm="ortho") - - if action_pred.shape[-1] == 7: - print("Before concat action_pred", action_pred.shape, flush=True) - right_action_pred = action_pred - left_action_pred = np.zeros_like(right_action_pred) - action_pred = np.concatenate([left_action_pred, right_action_pred], axis=1) - print("After concat action_pred", action_pred.shape, flush=True) - # unnorm action_pred - # print("action_pred", action_pred[:, 3], flush=True) - action_pred = ( - self.normalizer_action.unnormalize_data( - torch.tensor(action_pred).unsqueeze(0), [self.args.dataset_name] - ) - .squeeze(0) - .numpy() - ) - print("After unnormalize_data action_pred", action_pred.shape, flush=True) - - action_pred = self.model_output_process(action_pred, agent_data) - action_pred = self.postprocess(action_pred) - - return action_pred - - def generate_flow_action( - self, - inputs, - last_action_chunk=None, - max_guidance_weight=20.0, - num_inference_timesteps=10, - sigma_action=0.2, - ): - model_action_dim = sum(self.dof_config.values()) - if last_action_chunk is None: - output = self.model.generate_flow_action( - action_horizon=self.args.action_horizon, - action_dim=model_action_dim, - num_inference_timesteps=num_inference_timesteps, - unnorm=False, - **inputs, - ) - else: - output = self.model.generate_flow_action_rtc( - action_horizon=self.args.action_horizon, - action_dim=model_action_dim, - num_inference_timesteps=num_inference_timesteps, - inference_delay=self.args.rtc_inference_delay, - execution_horizon=self.args.rtc_execution_horizon - 1, - max_guidance_weight=max_guidance_weight, - last_action_chunk=last_action_chunk, - sigma_action=sigma_action, - unnorm=False, - **inputs, - ) - action_pred = output["predict_action"] # (b, action_horizon, action_dim) - return action_pred - - def generate_text(self, inputs): - return self.model.generate_text(**inputs) - - def _preprocess_ar_batch(self, batch): - input_ids = batch["input_ids"] - attention_mask = batch["attention_mask"] - moe_token_types = batch["moe_token_types"] - labels = batch.get("labels", None) - prefix_length = batch.get("prefix_length", None) - - generation_prompt_ids = torch.tensor( - [151644, 77091], device=input_ids.device, dtype=input_ids.dtype - ) # <|im_start|>assistant - matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & ( - input_ids[0, 1:] == generation_prompt_ids[1] - ) - if matches.any(): - split_pos = torch.nonzero(matches, as_tuple=True)[0][0].item() - # construct output ids - gt_output_ids = input_ids[:, split_pos + 3 : prefix_length] - # remove output part from input - input_ids = input_ids[:, : split_pos + 3] - moe_token_types = moe_token_types[:, : split_pos + 3] - if attention_mask is not None: - attention_mask = attention_mask[:, : split_pos + 3] - if labels is not None: - labels = labels[:, split_pos + 3 : prefix_length] - - batch.update( - { - "input_ids": input_ids, - "attention_mask": attention_mask, - "moe_token_types": moe_token_types, - "labels": labels, - "gt_output_ids": gt_output_ids, - "prefix_length": split_pos + 3, - } - ) - - return batch - - def generate_ar_action(self, inputs): - # batch = self._preprocess_ar_batch(batch=inputs) - action_pred = None - count = 0 - while action_pred is None: - if count > 5: - # raise ValueError("re-generate ar action failed") - return None - count += 1 - output = self.model.generate_ar_action( - # action_dim=args.action_dim, - action_dim=14, - action_horizon=self.args.action_horizon, - unnorm=False, - **inputs, - ) - action_pred = output["predict_action"] - - action_pred = action_pred[0] - return action_pred - - -class WallxInfer: - def __init__(self, args: WallxInferArgs): - self.args = args - self.model_wrapper = WallxModelWrapper(args) - - def run_infer_robochallenge( - self, state, views, instruction, valid_action_dim=7, action_predict_mode=None - ): - action_pred = self.model_wrapper.predict_action_rtc( - state=state, - views=views, - instruction=instruction, - valid_action_dim=valid_action_dim, - action_predict_mode=action_predict_mode, - ) - return action_pred - - -if __name__ == "__main__": - args = WallxInferArgs() - Infer = WallxInfer(args) - Infer.run_infer() diff --git a/scripts/merge_sharded_weights.py b/scripts/merge_sharded_weights.py index 314e228b..d58b9add 100644 --- a/scripts/merge_sharded_weights.py +++ b/scripts/merge_sharded_weights.py @@ -70,65 +70,16 @@ def load_sharded_checkpoint(checkpoint_dir: str) -> Dict[str, torch.Tensor]: except AttributeError as e: if "StorageMeta" in str(e): - print(f"[ERROR] StorageMeta compatibility issue: {e}") - print("[INFO] Attempting alternative loading method...") - return load_checkpoint_alternative(checkpoint_dir) - else: - raise - - -def load_checkpoint_alternative(checkpoint_dir: str) -> Dict[str, torch.Tensor]: - """ - Alternative method to load checkpoint by directly reading shard files. - - Args: - checkpoint_dir: Path to directory containing .distcp files - - Returns: - Dictionary of merged model state - """ - checkpoint_path = Path(checkpoint_dir) - - # Find all shard files - shard_files = sorted(checkpoint_path.glob("*.distcp")) - - if not shard_files: - raise FileNotFoundError(f"No .distcp files found in {checkpoint_dir}") - - print(f"[INFO] Found {len(shard_files)} shard files") - - # Load all shards - merged_state = {} - - for shard_file in shard_files: - print(f"[INFO] Loading shard: {shard_file.name}") - try: - shard_data = torch.load(shard_file, map_location="cpu") - - # Merge the shard into the state dict - if isinstance(shard_data, dict): - for key, value in shard_data.items(): - if isinstance(value, torch.Tensor): - if key in merged_state: - # Handle duplicates - concatenate or overwrite based on shape - print(f"[WARNING] Duplicate key found: {key}") - merged_state[key] = value - elif isinstance(value, dict): - # Nested dict structure - for subkey, subvalue in value.items(): - full_key = f"{key}.{subkey}" if key else subkey - if isinstance(subvalue, torch.Tensor): - merged_state[full_key] = subvalue - - except Exception as e: - print(f"[WARNING] Failed to load shard {shard_file.name}: {e}") - continue - - if not merged_state: - raise RuntimeError("Failed to load any checkpoint data from shards") - - print(f"[INFO] Loaded {len(merged_state)} tensors from shards") - return merged_state + raise RuntimeError( + "Unable to load this DCP checkpoint because its metadata uses " + "StorageMeta from a different PyTorch version. The previous " + "manual .distcp fallback was removed because FSDP shards cannot " + "be reconstructed by directly loading shard files and overwriting " + "duplicate keys. Please run this script with a PyTorch version " + "compatible with the checkpoint writer, or re-save the checkpoint " + "with the current PyTorch version." + ) from e + raise def save_merged_checkpoint( diff --git a/scripts/normalize.py b/scripts/normalize.py index 75a3f98e..9b1ef008 100644 --- a/scripts/normalize.py +++ b/scripts/normalize.py @@ -1,4 +1,6 @@ -# This file is copied from openpi +# Adapted from Physical-Intelligence/openpi: +# https://github.com/Physical-Intelligence/openpi/blob/650c5b0283a49c42784fb5055a0507da2c6d347d/src/openpi/shared/normalize.py +# Upstream license: Apache-2.0. import json import pathlib diff --git a/scripts/run_libero.sh b/scripts/run_libero.sh new file mode 100755 index 00000000..4c5b8f25 --- /dev/null +++ b/scripts/run_libero.sh @@ -0,0 +1,204 @@ +#!/usr/bin/env bash +set -euo pipefail + +# LIBERO evaluation launcher. Configure with environment variables: +# +# CHECKPOINT_PATH=/path/to/checkpoint bash scripts/run_libero.sh +# bash scripts/run_libero.sh /path/to/checkpoint +# CHECKPOINT_PATH=/path/to/checkpoint SMOKE=1 bash scripts/run_libero.sh +# CONFIG=/path/to/eval.yaml bash scripts/run_libero.sh +# +# Optional knobs: +# CUDA_ID=0 +# DRIVER_MODE=in_process +# NUM_WORKERS=1 +# MAX_BATCH_SIZE=1 +# ALL_SUITES=1 # run all standard LIBERO suites (40 tasks) +# TASK_SUITES="libero_spatial ..." # custom suite list (space- or comma-separated) +# TASK_SUITE_NAME=libero_spatial # single suite when ALL_SUITES=0 and TASK_SUITES unset +# NUM_TRIALS_PER_TASK=50 +# TASK_INDICES=0,1,2 # omit to run every task in the suite +# MAX_INFER_TIMES=52 +# NORM_KEY=libero_all +# ROLLOUT_BASE=/path/to/rollout # per-suite logs under ${ROLLOUT_BASE}/${suite}/ +# LOG_DIR=/tmp/harrix_libero_eval # overrides ROLLOUT_BASE when set (single suite) + +export PYTHONPATH="/path/to/LIBERO:${PYTHONPATH:-}" + +export ALL_SUITES=1 +# export TASK_SUITE_NAME=libero_10 +# export CUDA_ID=0 +export NUM_WORKERS=10 +export NUM_TRIALS_PER_TASK=20 +export CHECKPOINT_PATH=/path/to/checkpoint +export ROLLOUT_BASE=/path/to/rollout + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" + +DEFAULT_ALL_SUITES=( + libero_spatial + libero_object + libero_goal + libero_10 +) + +if [[ $# -gt 1 ]]; then + echo "Usage: bash scripts/run_libero.sh [CHECKPOINT_PATH]" >&2 + exit 2 +fi +if [[ $# -eq 1 ]]; then + CHECKPOINT_PATH="$1" +fi + +export CUDA_VISIBLE_DEVICES="${CUDA_ID:-0}" + +# MuJoCo offscreen rendering via NVIDIA EGL (required on headless GPU nodes). +export MUJOCO_GL=egl +export PYOPENGL_PLATFORM=egl +EGL_VENDOR_DIR="${HOME}/.config/glvnd/egl_vendor.d" +EGL_VENDOR_FILE="${EGL_VENDOR_DIR}/10_nvidia.json" +if [[ ! -f "${EGL_VENDOR_FILE}" ]]; then + mkdir -p "${EGL_VENDOR_DIR}" + cat > "${EGL_VENDOR_FILE}" <<'EOF' +{ + "file_format_version" : "1.0.0", + "ICD" : { + "library_path" : "libEGL_nvidia.so.0" + } +} +EOF +fi +export __EGL_VENDOR_LIBRARY_FILENAMES="${EGL_VENDOR_FILE}" +export LD_LIBRARY_PATH="/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib:/usr/local/nvidia/lib64:${LD_LIBRARY_PATH:-}" + +# After CUDA_VISIBLE_DEVICES remapping, MuJoCo only sees devices from index 0. +export MUJOCO_EGL_DEVICE_ID="${MUJOCO_EGL_DEVICE_ID:-0}" +if [[ -d "${REPO_ROOT}/third_party/harrix/python" ]]; then + export PYTHONPATH="${REPO_ROOT}/third_party/harrix/python:${PYTHONPATH:-}" +fi +export PYTHONPATH="${REPO_ROOT}:${PYTHONPATH:-}" + +resolve_task_suites() { + local suites=() + if [[ "${ALL_SUITES:-0}" == "1" ]]; then + if [[ -n "${TASK_SUITES:-}" ]]; then + TASK_SUITES="${TASK_SUITES//,/ }" + read -r -a suites <<< "${TASK_SUITES}" + else + suites=("${DEFAULT_ALL_SUITES[@]}") + fi + elif [[ -n "${TASK_SUITES:-}" ]]; then + TASK_SUITES="${TASK_SUITES//,/ }" + read -r -a suites <<< "${TASK_SUITES}" + else + suites=("${TASK_SUITE_NAME:-libero_spatial}") + fi + echo "${suites[@]}" +} + +resolve_log_dirs() { + local suite="$1" + local multi_suite="$2" + local log_dir rollout_dir + + if [[ -n "${LOG_DIR:-}" ]]; then + if [[ "${multi_suite}" == "1" ]]; then + log_dir="${LOG_DIR}/${suite}" + else + log_dir="${LOG_DIR}" + fi + elif [[ -n "${ROLLOUT_BASE:-}" ]]; then + if [[ "${multi_suite}" == "1" ]]; then + log_dir="${ROLLOUT_BASE}/${suite}/json" + else + log_dir="${ROLLOUT_BASE}/json" + fi + else + log_dir="/tmp/harrix_libero_eval" + fi + + if [[ -n "${WALLX_ROLLOUT_DIR:-}" ]]; then + if [[ "${multi_suite}" == "1" ]]; then + rollout_dir="${WALLX_ROLLOUT_DIR}/${suite}" + else + rollout_dir="${WALLX_ROLLOUT_DIR}" + fi + elif [[ -n "${ROLLOUT_BASE:-}" ]]; then + if [[ "${multi_suite}" == "1" ]]; then + rollout_dir="${ROLLOUT_BASE}/${suite}/videos" + else + rollout_dir="${ROLLOUT_BASE}/videos" + fi + else + rollout_dir="" + fi + + mkdir -p "${log_dir}" + if [[ -n "${rollout_dir}" ]]; then + mkdir -p "${rollout_dir}" + fi + LOG_DIR_RESOLVED="${log_dir}" + WALLX_ROLLOUT_DIR_RESOLVED="${rollout_dir}" +} + +run_suite() { + local suite="$1" + local multi_suite="$2" + + resolve_log_dirs "${suite}" "${multi_suite}" + if [[ -n "${WALLX_ROLLOUT_DIR_RESOLVED}" ]]; then + export WALLX_ROLLOUT_DIR="${WALLX_ROLLOUT_DIR_RESOLVED}" + else + unset WALLX_ROLLOUT_DIR + fi + + local args=() + if [[ -n "${CONFIG:-}" ]]; then + args+=(--config "${CONFIG}") + else + if [[ -z "${CHECKPOINT_PATH:-}" ]]; then + echo "CHECKPOINT_PATH is required when CONFIG is not set." >&2 + exit 2 + fi + args+=(--checkpoint-path "${CHECKPOINT_PATH}") + fi + + if [[ -n "${TRAIN_CONFIG_PATH:-}" ]]; then + args+=(--train-config-path "${TRAIN_CONFIG_PATH}") + fi + if [[ "${SMOKE:-0}" == "1" ]]; then + args+=(--smoke) + fi + if [[ "${DET:-0}" == "1" ]]; then + args+=(--deterministic-model) + fi + if [[ -n "${TASK_INDICES:-}" ]]; then + args+=(--task-indices "${TASK_INDICES}") + fi + if [[ -n "${MAX_INFER_TIMES:-}" ]]; then + args+=(--max-infer-times "${MAX_INFER_TIMES}") + fi + + args+=(--driver-mode "${DRIVER_MODE:-in_process}") + args+=(--num-workers "${NUM_WORKERS:-1}") + args+=(--max-batch-size "${MAX_BATCH_SIZE:-1}") + args+=(--task-suite-name "${suite}") + args+=(--num-trials-per-task "${NUM_TRIALS_PER_TASK:-50}") + args+=(--norm-key "${NORM_KEY:-libero_all}") + args+=(--log-dir "${LOG_DIR_RESOLVED}") + + echo "=== LIBERO eval: suite=${suite} log_dir=${LOG_DIR_RESOLVED} ===" + python scripts/infer_libero.py "${args[@]}" +} + +suites=($(resolve_task_suites)) +multi_suite=0 +if [[ "${#suites[@]}" -gt 1 ]]; then + multi_suite=1 +fi + +cd "${REPO_ROOT}" +for suite in "${suites[@]}"; do + run_suite "${suite}" "${multi_suite}" +done diff --git a/scripts/run_serving.sh b/scripts/run_serving.sh new file mode 100755 index 00000000..d05f2695 --- /dev/null +++ b/scripts/run_serving.sh @@ -0,0 +1,16 @@ +#!/bin/bash + +export ENABLE_CUDA_GRAPH=True +export ENABLE_EXPERIMENTAL_INFERENCE_ENGINE=True + +CKPT_PATH=/path/to/checkpoint + +python -m wall_x._vendor.harrix.serving.launch_serving \ + --env X2ROBOT \ + --port 32195 \ + model-config:server-model-config \ + --model-config.checkpoint-path "$CKPT_PATH" \ + --model-config.action-horizon 32 \ + --model-config.robot-action-interpolate-multiplier 1 \ + --model-config.robot-action-end-ratio 1.0 \ + --model-config.robot-type desktop \ No newline at end of file diff --git a/scripts/vqa_inference.py b/scripts/vqa_inference.py index adbef89a..b8e10a2e 100644 --- a/scripts/vqa_inference.py +++ b/scripts/vqa_inference.py @@ -90,7 +90,7 @@ def generate(self, image: Image.Image, text: str, **kwargs) -> str: # Local Image img = Image.open( - "/x2robot_v2/yangping/github/wall-x/assets/cot_example_frame.png" + "/path/to/cot_example_frame.png" ).convert("RGB") # Internet Image # import requests diff --git a/wall_x/_vendor/__init__.py b/wall_x/_vendor/__init__.py new file mode 100644 index 00000000..e2e697ef --- /dev/null +++ b/wall_x/_vendor/__init__.py @@ -0,0 +1,4 @@ +"""Vendored sister-package sources, bundled at export time so the exported +package is self-contained. Internal development uses ``pip install -e`` +against the live repos instead. +""" diff --git a/wall_x/_vendor/harrix/__init__.py b/wall_x/_vendor/harrix/__init__.py new file mode 100644 index 00000000..1db9b1b8 --- /dev/null +++ b/wall_x/_vendor/harrix/__init__.py @@ -0,0 +1,5 @@ +"""Public harrix inference and evaluation runtime.""" + +__version__ = "0.1.0" + +__all__ = ["__version__"] diff --git a/wall_x/_vendor/harrix/adapters/__init__.py b/wall_x/_vendor/harrix/adapters/__init__.py new file mode 100644 index 00000000..5791ac80 --- /dev/null +++ b/wall_x/_vendor/harrix/adapters/__init__.py @@ -0,0 +1,6 @@ +"""Adapter package entry point. + +Importing ``variants`` triggers variant registration side effects. +""" + +from wall_x._vendor.harrix.adapters import variants # noqa: F401 diff --git a/wall_x/_vendor/harrix/adapters/base.py b/wall_x/_vendor/harrix/adapters/base.py new file mode 100644 index 00000000..68b7f91b --- /dev/null +++ b/wall_x/_vendor/harrix/adapters/base.py @@ -0,0 +1,51 @@ +"""Common inference adapter abstraction. + +The adapter owns all architecture-specific model setup and exposes a single +``predict_batch`` entry point. Environment drivers pass payloads through without +interpreting the env-adapter schema. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +import numpy as np + +from wall_x._vendor.harrix.eval_config import EvalConfig + + +class BaseInferAdapter(ABC): + @abstractmethod + def __init__(self, cfg: EvalConfig) -> None: + """Load processors, model config, checkpoints, and normalizers. + + Subclasses should reject unsupported ``cfg.model.action_mode`` values + during construction. + """ + + @property + @abstractmethod + def chunk_horizon(self) -> int: + """Number of action steps returned by each ``predict_batch`` call.""" + + @property + @abstractmethod + def action_mode(self) -> str: + """Configured inference algorithm, fixed at construction time.""" + + @abstractmethod + def predict_batch(self, payloads: list[dict]) -> list[np.ndarray]: + """Run one batched inference call and return one chunk per payload. + + ``payloads[i]`` schema is defined by the env-adapter pair: + { + "observation": dict[str, np.ndarray], + # Env-defined ndarray bundle, for example LIBERO: + # {"eef_pos":(3,), "eef_axisangle":(3,), "gripper":(1,), + # "face_view":(H,W,3), "wrist_view":(H,W,3)} + "instruction": str, + "noise": np.ndarray | None, + # Flow may pass (chunk_horizon, action_dim); other modes + # may leave this as None. + } + """ diff --git a/wall_x/_vendor/harrix/adapters/qwen_vlact.py b/wall_x/_vendor/harrix/adapters/qwen_vlact.py new file mode 100644 index 00000000..400b71c1 --- /dev/null +++ b/wall_x/_vendor/harrix/adapters/qwen_vlact.py @@ -0,0 +1,465 @@ +"""Shared adapter base for Qwen-VL action models. + +Variant-specific subclasses live in ``harrix.adapters.variants`` and provide +the Wall-X training adapter class via ``_training_adapter``. The shared base +handles checkpoint loading, normalizers, LIBERO observation encoding, prompt +construction, and flow-action inference. +""" + +from __future__ import annotations + +import copy +import logging + +import numpy as np +import torch +from PIL import Image + +from qwen_vl_utils.vision_process import smart_resize + +from wall_x._vendor.x2robot_utils.text_templates import ( + get_prologue_with_embodied_information, + preprocesser_call, +) + +from wall_x._vendor.harrix.adapters.base import BaseInferAdapter +from wall_x._vendor.harrix.envs.libero_common import decode_chunk, encode_proprio +from wall_x._vendor.harrix.eval_config import EvalConfig +from wall_x._vendor.harrix.utils.ckpt_load import ( + load_state_dict, + resolve_checkpoint_dir, +) +from wall_x._vendor.harrix.utils.normalizer import build_normalizers +from wall_x._vendor.harrix.utils.train_config import ( + build_data_config, + build_model_config, + load_train_config_with_ckpt_overlay, + normalize_train_config_for_inference, + register_data_backend, + resolve_state_bins, + resolve_use_state_string_representation, +) +from wall_x.trainer.trainer_utils import load_wallx_processors + + +_SUPPORTED_MODES = {"flow", "ar", "dllm", "vqa", "subtask"} +_IMPLEMENTED_MODES = {"flow"} + +# Special tokens used by the Qwen-VL action model. +_ROLE_START = "<|im_start|>" +_ROLE_END = "<|im_end|>" +_VISION_START = "<|vision_start|>" +_VISION_END = "<|vision_end|>" +_IMAGE_PAD = "<|image_pad|>" +_PROPRI = "<|propri|>" +_ACTION = "<|action|>" +logger = logging.getLogger(__name__) + + +_PUBLIC_CAMERA_LABELS = { + "face_view": "front view", + "right_wrist_view": "right wrist view", + "left_wrist_view": "left wrist view", +} + + +def _camera_label(cam_name: str) -> str: + return _PUBLIC_CAMERA_LABELS.get(cam_name, cam_name.replace("_", " ")) + + +def _normalizer_width(normalizer, norm_key: str) -> int | None: + if normalizer is None or norm_key not in normalizer.delta: + return None + return int(normalizer.delta[norm_key].shape[0]) + + +def _normalize_real_prefix(normalizer, tensor, dataset_names): + if normalizer is None: + return tensor + widths = [_normalizer_width(normalizer, name) for name in dataset_names] + if any(width is None for width in widths) or len(set(widths)) != 1: + return normalizer.normalize_data(tensor, dataset_names) + width = widths[0] + if tensor.shape[-1] == width: + return normalizer.normalize_data(tensor, dataset_names) + if tensor.shape[-1] < width: + raise ValueError( + f"normalizer width {width} exceeds tensor dim {tensor.shape[-1]}" + ) + out = tensor.clone() + out[..., :width] = normalizer.normalize_data( + tensor[..., :width], dataset_names + ) + out[..., width:] = 0 + return out + +class QwenVLActInferAdapter(BaseInferAdapter): + """Shared constructor and batched flow inference implementation.""" + + # ---- subclass hook ---- + + @classmethod + def _training_adapter(cls): + """Return the training-side ModelAdapter subclass.""" + raise NotImplementedError(f"{cls.__name__} must override _training_adapter()") + + # ---- ctor ---- + + def __init__(self, cfg: EvalConfig) -> None: + mode = cfg.model.action_mode + if mode not in _SUPPORTED_MODES: + raise ValueError( + f"{type(self).__name__}: unknown action_mode={mode!r}, " + f"supported={sorted(_SUPPORTED_MODES)}" + ) + if mode not in _IMPLEMENTED_MODES: + raise NotImplementedError( + f"{type(self).__name__}: action_mode={mode!r} is not implemented; " + "currently supported=" + f"{sorted(_IMPLEMENTED_MODES)}" + ) + self._action_mode = mode + + ta = self._training_adapter() + device = "cuda" + self._device = device + self._checkpoint_path = resolve_checkpoint_dir(cfg.model.checkpoint_path) + + # 1) train_config with checkpoint overlays, then data backend registration. + train_config = load_train_config_with_ckpt_overlay( + cfg.model.train_config_path, self._checkpoint_path + ) + train_config = normalize_train_config_for_inference( + train_config, cfg.model.train_config_path + ) + register_data_backend(train_config) + self._train_config = train_config + + # 2) normalizers first — action_tokenizer setup needs them during processor load. + normalizer_action, normalizer_propri, resolved_norm_key = build_normalizers( + self._checkpoint_path, train_config, cfg.model.norm_key + ) + self._normalizer_action = normalizer_action + self._normalizer_propri = normalizer_propri + self._norm_key = resolved_norm_key + if normalizer_propri is not None and resolved_norm_key in normalizer_propri.delta: + train_config["_libero_proprio_norm_dim"] = int( + normalizer_propri.delta[resolved_norm_key].shape[0] + ) + + # 3) data_config. Image resizing needs resolution/image_factor/min/max. + self._data_config = build_data_config(cfg.model.train_config_path, train_config) + + # 4) HF model config + ConfigClass = ta.config_class() + self._model_config = build_model_config( + ConfigClass, + self._checkpoint_path, + train_config, + cfg.model.train_config_path, + ) + + # 5) processor + tokenizer_mixin (may extend vocab via action_tokenizer) + procs = load_wallx_processors( + train_config, normalizer=normalizer_action, device=device + ) + self._processor = procs["processor"] + self._tokenizer_mixin = procs.get("tokenizer_mixin") + logger.info( + "processor vocab size: %s (action_tokenizer_type=%r)", + len(self._processor.tokenizer), + train_config.get("action_tokenizer_type"), + ) + + # 6) model + ModelClass = ta.inference_model_class() + self._model_class = ModelClass + model = ModelClass(self._model_config, self._processor, self._tokenizer_mixin) + model.resize_token_embeddings(len(self._processor.tokenizer)) + model.to_bfloat16_for_selected_params() + + # 7) checkpoint weights and model finalization. + state_dict = load_state_dict(self._checkpoint_path, ModelClass) + embed_key = "model.embed_tokens.weight" + if embed_key in state_dict: + ckpt_vocab = state_dict[embed_key].shape[0] + cur_vocab = model.model.embed_tokens.weight.shape[0] + if cur_vocab != ckpt_vocab: + logger.info( + "resize_token_embeddings from %d to %d to match checkpoint", + cur_vocab, + ckpt_vocab, + ) + model.resize_token_embeddings(ckpt_vocab) + msg = model.load_state_dict(state_dict, strict=False) + logger.info( + "%s load_state_dict: missing=%s unexpected=%s", + type(self).__name__, + len(msg.missing_keys), + len(msg.unexpected_keys), + ) + model.set_normalizer( + copy.deepcopy(normalizer_action), + copy.deepcopy(normalizer_propri), + ) + model.eval() + model.to(device) + model.to_bfloat16_for_selected_params() + self._model = model + + # 8) cached runtime fields + self._cam_names = list(cfg.model.cam_names) + self._action_horizon = int(cfg.model.action_horizon) + self._action_dim = sum(train_config["dof_config"].values()) + self._num_inference_timesteps = 10 + self._robot_id = "10000" + + # ---- BaseInferAdapter API ---- + + @property + def chunk_horizon(self) -> int: + return self._action_horizon + + @property + def action_mode(self) -> str: + return self._action_mode + + def predict_batch(self, payloads: list[dict]) -> list[np.ndarray]: + if self._action_mode == "flow": + return self._flow_batch(payloads) + raise NotImplementedError( + f"{type(self).__name__}.predict_batch: action_mode={self._action_mode!r} " + "is not dispatched" + ) + + # ---- batched flow inference ---- + + def _flow_batch(self, payloads: list[dict]) -> list[np.ndarray]: + # 1) Encode each observation into proprioception, masks, and views. + observations: list[dict] = [] + instructions: list[str] = [] + noises: list[np.ndarray | None] = [] + any_noise = False + for p in payloads: + observations.append( + encode_proprio( + p["observation"], self._train_config, self._action_horizon + ) + ) + instructions.append(p["instruction"]) + n = p.get("noise") + if n is not None: + any_noise = True + noises.append(n) + + # 2) Stack noise when provided; all-None delegates sampling to the model. + if any_noise: + if any(n is None for n in noises): + raise ValueError( + "payload noise must be provided for every payload or for none" + ) + batch_noise = torch.stack( + [torch.from_numpy(n).to(dtype=torch.float32) for n in noises], dim=0 + ) + else: + batch_noise = None + + # 3) Prompts and model inputs. + prefix_list, postfix_list = [], [] + for ins in instructions: + prefix, postfix = self._get_flow_prompt(ins) + prefix_list.append(prefix) + postfix_list.append(postfix) + batch_inputs = self._construct_model_input( + observations, prefix_list, postfix_list + ) + + # 4) Normalized zero action as the flow starting point. + padding = torch.zeros( + ( + len(batch_inputs["dataset_names"]), + 1, + self._action_dim, + ), + dtype=torch.float32, + ) + padding_action = _normalize_real_prefix( + self._normalizer_action, padding, batch_inputs["dataset_names"] + ).to(batch_inputs["input_ids"].device) + + # 5) Flow forward. + model_output = self._model.generate_flow_action( + action_horizon=self._action_horizon, + action_dim=self._action_dim, + num_inference_timesteps=self._num_inference_timesteps, + padding_action=padding_action, + noise=batch_noise, + **batch_inputs, + ) + + # 6) Decode one action chunk per payload. + predict_action = model_output["predict_action"] # (B, H, D_action) + if isinstance(predict_action, torch.Tensor): + predict_action = predict_action.detach().cpu().numpy() + return [ + decode_chunk(predict_action[i : i + 1], self._train_config) + for i in range(len(payloads)) + ] + + # ---- prompt template ---- + + def _get_flow_prompt(self, instruction: str) -> tuple[str, str]: + """Build the flow-action prompt.""" + if self._train_config["data"].get("use_embodied_system_prompt_ratio", 0) > 0: + robot_id = ( + self._robot_id if self._norm_key in ("x2_normal", "ex_normal") else 0 + ) + cam_name_mapping = {cn: cn for cn in self._cam_names} + prologue = get_prologue_with_embodied_information( + dataset_name=self._norm_key, + cam_mapping=cam_name_mapping, + robot_id=robot_id, + uid="", + config=self._data_config, + ) + else: + prologue = f"{_ROLE_START}system\nYou are a helpful assistant.{_ROLE_END}\n" + user_request = f"{_ROLE_START}user\nObservation:" + for cn in self._cam_names: + user_request += ( + f" {_camera_label(cn)}: " f"{_VISION_START}{_IMAGE_PAD}{_VISION_END}" + ) + user_request += "\nInstruction:" + text_prompt = ( + f"\nPredict the next action in robot action.\nProprioception: {_PROPRI}\n" + ) + user_message = f"{user_request} {instruction}{text_prompt}{_ROLE_END}\n" + assistant_message = f"{_ROLE_START}assistant\n" + flow_action = _ACTION * self._action_horizon + + prefix_text = prologue + user_message + assistant_message + postfix_text = flow_action + return prefix_text, postfix_text + + # ---- batched model input construction ---- + + def _construct_model_input( + self, + observations: list[dict], + prefix_list: list[str], + postfix_list: list[str], + ) -> dict: + """Build model inputs for the flow inference path.""" + batch_size = len(observations) + dataset_names = [self._norm_key] * batch_size + + # Proprioception and masks are prepared as ndarrays in encode_proprio. + agent_pos = torch.cat( + [torch.from_numpy(o["proprioception"]) for o in observations], dim=0 + ) + agent_pos_mask = torch.cat( + [torch.from_numpy(o["agent_pos_mask"]) for o in observations], dim=0 + ) + dof_mask = torch.cat( + [torch.from_numpy(o["dof_mask"]) for o in observations], dim=0 + ) + agent_pos = _normalize_real_prefix( + self._normalizer_propri, agent_pos, dataset_names + ) + + # Resize images per sample. + image_inputs: list[torch.Tensor] = [] + all_image_sizes: list[tuple[int, int]] = [] + for o in observations: + for cn in self._cam_names: + if cn not in o: + continue + tensor = self._resize_image(o[cn], cn) + image_inputs.append(tensor) + # Tensor (H, W, C) to PIL-compatible (W, H). + all_image_sizes.append((tensor.shape[1], tensor.shape[0])) + + inputs = preprocesser_call( + processor=self._processor, + prefix_text=prefix_list, + postfix_text=postfix_list, + images=image_inputs, + videos=None, + padding=True, + truncation=True, + return_tensors="pt", + max_length=1000, + pad_to_128_multiple=False, + pad_prefix_to_same_length=False, + norm_state=( + agent_pos + if resolve_use_state_string_representation(self._train_config) + else None + ), + agent_pos_mask=agent_pos_mask, + state_augmentation_prob=0.0, + state_drop_prob=0.0, + state_augmentation_ratio=0.0, + state_bins=resolve_state_bins(self._train_config), + inference_mode=True, + ) + + action_token_id = self._processor.tokenizer.convert_tokens_to_ids(_ACTION) + moe_token_types = inputs["input_ids"] == action_token_id + + extra = { + "proprioception": agent_pos.detach(), + "agent_pos_mask": agent_pos_mask, + "dof_mask": dof_mask, + "image_size": all_image_sizes, + "moe_token_types": moe_token_types, + "dataset_names": dataset_names, + } + inputs.update(extra) + return _move_to_device(inputs, self._device) + + # ---- image resizing ---- + + def _resize_image(self, img: np.ndarray, cam_name: str) -> torch.Tensor: + """Resize one image with the train-time image config.""" + if isinstance(img, np.ndarray): + pil = Image.fromarray(img) + elif isinstance(img, Image.Image): + pil = img + else: + raise ValueError(f"unsupported image type: {type(img)}") + orig_w, orig_h = pil.size + + target = self._data_config.resolution.get(cam_name, -1) + if target != -1: + if orig_w > orig_h: + new_w, new_h = target, int(target * orig_h / orig_w) + else: + new_h, new_w = target, int(target * orig_w / orig_h) + pil = pil.resize((new_w, new_h)) + + cur_w, cur_h = pil.size + resized_h, resized_w = smart_resize( + cur_h, + cur_w, + factor=self._data_config.image_factor, + min_pixels=self._data_config.min_pixels, + max_pixels=self._data_config.max_pixels, + ) + resized = pil.resize((resized_w, resized_h)) + return torch.from_numpy(np.array(resized)).to(self._device) + + +def _move_to_device(obj, device): + """Recursively move tensors inside common containers to ``device``.""" + from transformers import BatchFeature + + if isinstance(obj, torch.Tensor): + return obj.to(device) + if isinstance(obj, (dict, BatchFeature)): + return {k: _move_to_device(v, device) for k, v in obj.items()} + if isinstance(obj, list): + return [_move_to_device(v, device) for v in obj] + if isinstance(obj, tuple): + return tuple(_move_to_device(v, device) for v in obj) + return obj diff --git a/wall_x/_vendor/harrix/adapters/registry.py b/wall_x/_vendor/harrix/adapters/registry.py new file mode 100644 index 00000000..ddc0ec4d --- /dev/null +++ b/wall_x/_vendor/harrix/adapters/registry.py @@ -0,0 +1,39 @@ +"""Inference adapter registry and factory.""" + +from __future__ import annotations + +from typing import Type + +from wall_x._vendor.harrix.adapters.base import BaseInferAdapter +from wall_x._vendor.harrix.eval_config import EvalConfig + + +ADAPTER_REGISTRY: dict[str, Type[BaseInferAdapter]] = {} + + +def register_adapter(name: str): + def deco(cls: Type[BaseInferAdapter]): + if name in ADAPTER_REGISTRY: + raise ValueError( + f"adapter {name!r} is already registered " + f"(cls={ADAPTER_REGISTRY[name].__name__})" + ) + ADAPTER_REGISTRY[name] = cls + return cls + + return deco + + +def build_adapter(cfg: EvalConfig) -> BaseInferAdapter: + arch = cfg.model.architecture + cls = ADAPTER_REGISTRY.get(arch) + if cls is None: + raise ValueError( + f"unknown architecture={arch!r}, registered={sorted(ADAPTER_REGISTRY)}" + ) + return cls(cfg) + + +def registered_architectures() -> list[str]: + """Return registered architecture names for diagnostics.""" + return sorted(ADAPTER_REGISTRY) diff --git a/wall_x/_vendor/harrix/adapters/variants/__init__.py b/wall_x/_vendor/harrix/adapters/variants/__init__.py new file mode 100644 index 00000000..7219446b --- /dev/null +++ b/wall_x/_vendor/harrix/adapters/variants/__init__.py @@ -0,0 +1,15 @@ +"""Register inference adapter variants through import side effects.""" + +from __future__ import annotations + +import importlib +import pkgutil + + +for module_info in pkgutil.iter_modules(__path__): + if module_info.name.startswith("_"): + continue + try: + importlib.import_module(f"{__name__}.{module_info.name}") + except ImportError: + continue diff --git a/wall_x/_vendor/harrix/adapters/variants/qwen2_5.py b/wall_x/_vendor/harrix/adapters/variants/qwen2_5.py new file mode 100644 index 00000000..036409b5 --- /dev/null +++ b/wall_x/_vendor/harrix/adapters/variants/qwen2_5.py @@ -0,0 +1,17 @@ +"""Qwen2.5 inference adapter registration. + +The shared inference logic lives in ``harrix.adapters.qwen_vlact``. This module +binds the Qwen2.5 variant key to Wall-X's training-side adapter class. +""" + +from wall_x._vendor.harrix.adapters.qwen_vlact import QwenVLActInferAdapter +from wall_x._vendor.harrix.adapters.registry import register_adapter + + +@register_adapter("qwen2_5") +class Qwen2_5InferAdapter(QwenVLActInferAdapter): + @classmethod + def _training_adapter(cls): + from wall_x.model.qact.qwen2_5.adapter import Qwen2_5Adapter + + return Qwen2_5Adapter diff --git a/wall_x/fusions/__init__.py b/wall_x/_vendor/harrix/drivers/__init__.py similarity index 100% rename from wall_x/fusions/__init__.py rename to wall_x/_vendor/harrix/drivers/__init__.py diff --git a/wall_x/_vendor/harrix/drivers/inproc/__init__.py b/wall_x/_vendor/harrix/drivers/inproc/__init__.py new file mode 100644 index 00000000..5813d0cb --- /dev/null +++ b/wall_x/_vendor/harrix/drivers/inproc/__init__.py @@ -0,0 +1,197 @@ +"""In-process driver: DirectModelHandle + N x SubprocEnvHandle. + +The model lives in the driver process while simulator envs live in subprocesses. +Evaluation proceeds in chunk-level lockstep: + 1. Seed the driver process and construct the model handle. + 2. Start one subprocess env handle per worker. + 3. Claim a task-local frame from JobState. + 4. Reset envs, batch active observations, run model.predict_batch, then + fan out action chunks to env subprocesses. + 5. Complete all episodes in the frame and move to the next frame. +""" + +from __future__ import annotations + +import json +import logging +import os +import time + +import numpy as np + +from wall_x._vendor.harrix.eval_config import EvalConfig +from wall_x._vendor.harrix.drivers.inproc.env_handle import SubprocEnvHandle +from wall_x._vendor.harrix.drivers.inproc.model_handle import DirectModelHandle +from wall_x._vendor.harrix.drivers.job_state import JobState + +logger = logging.getLogger(__name__) + + +def run(cfg: EvalConfig) -> None: + import wall_x._vendor.harrix.envs # noqa: F401 trigger env register + import wall_x._vendor.harrix.adapters # noqa: F401 trigger adapter register + from wall_x._vendor.harrix.envs.registry import enumerate_episodes_for + from wall_x._vendor.harrix.utils.seed import set_seed_everywhere + + # 1) Seed the driver process. + set_seed_everywhere(cfg.env.seed) + + # 2) Build DirectModelHandle; the model is loaded in the driver process. + logger.info("Constructing DirectModelHandle in the driver process") + t_model = time.time() + model_handle = DirectModelHandle(cfg) + logger.info("Model loaded in %.1fs", time.time() - t_model) + + # 3) Start env subprocesses. + logger.info("Starting %s SubprocEnvHandle(s)", cfg.runtime.num_workers) + env_handles = [ + SubprocEnvHandle(cfg, worker_id=i) for i in range(cfg.runtime.num_workers) + ] + + # 4) JobState runs in frame-sync mode for lockstep evaluation. + os.makedirs(cfg.runtime.log_dir, exist_ok=True) + log_path = os.path.join(cfg.runtime.log_dir, "state.jsonl") + report_path = os.path.join(cfg.runtime.log_dir, "report.json") + episodes = enumerate_episodes_for(cfg) + logger.info("env.type=%r; scheduled %s episodes", cfg.env.type, len(episodes)) + state = JobState( + episodes, + log_path, + batch_sync_mode=True, + batch_size=cfg.runtime.num_workers, + ) + + # 5) Main loop, one frame at a time. + t0 = time.time() + frame_idx = 0 + while not state.is_drained(): + frame_eps = state.claim_frame() + if not frame_eps: + # Previous frame still has in-flight episodes. + time.sleep(0.05) + continue + + results = _run_frame(frame_eps, model_handle, env_handles, cfg, frame_idx) + for ep, res in zip(frame_eps, results): + if "_error" in res: + state.fail(ep, str(res["_error"])) + else: + state.complete(ep, res) + frame_idx += 1 + + elapsed = time.time() - t0 + logger.info("All episodes finished in %.1fs (%.1f min)", elapsed, elapsed / 60) + + state.dump_final(report_path) + with open(report_path) as f: + report = json.load(f) + overall = report["overall"] + logger.info( + "attempted=%s, successes=%s, success_rate=%.2f%%, failed=%s", + overall["attempted"], + overall["successes"], + overall["success_rate"] * 100, + overall["failed"], + ) + logger.info("Report: %s", report_path) + logger.info("State log: %s", log_path) + + for h in env_handles: + h.shutdown() + model_handle.shutdown() + + +def _run_frame( + frame_eps: list, + model_handle: DirectModelHandle, + env_handles: list, + cfg: EvalConfig, + frame_idx: int, +) -> list[dict]: + """Run one task-local frame with chunk-level lockstep. + + Active workers are batched together at each chunk boundary. Workers that + already finished no longer participate in later forwards. + """ + from wall_x._vendor.harrix.envs.libero_common import encode_raw_obs + + n = len(frame_eps) + t_frame_start = time.time() + + # ---- a) reset: fan out, then gather ---- + for i in range(n): + env_handles[i].submit_reset(tuple(frame_eps[i])) + initials = [env_handles[i].wait_reset() for i in range(n)] + + obs_list = [r["obs"] for r in initials] + instr_list = [r["instruction"] for r in initials] + status = [ + { + "done": False, + "success": False, + "steps": 0, + "task_desc": initials[i].get("task_desc", ""), + } + for i in range(n) + ] + + max_rounds = cfg.env.libero.max_infer_times + + # ---- b) chunk lockstep ---- + for round_idx in range(max_rounds): + active = [i for i in range(n) if not status[i]["done"]] + if not active: + break + + payloads = [ + { + "observation": encode_raw_obs(obs_list[i]), + "instruction": instr_list[i], + "noise": None, + } + for i in active + ] + chunks = model_handle.predict_batch(payloads) + + # fan-out submit + for k, i in enumerate(active): + env_handles[i].submit_execute_chunk(chunks[k]) + # gather + for k, i in enumerate(active): + try: + r = env_handles[i].wait_execute_chunk() + except Exception as e: + status[i]["_error"] = str(e) + status[i]["done"] = True + continue + obs_list[i] = r["obs"] + status[i]["steps"] += r["steps"] + if r["done"]: + status[i]["done"] = True + status[i]["success"] = True + + for i in range(n): + env_handles[i].submit_finalize_episode(status[i]["success"]) + for i in range(n): + env_handles[i].wait_finalize_episode() + + elapsed_frame = time.time() - t_frame_start + logger.info( + "frame=%s n=%s succ=%s/%s elapsed=%.1fs", + frame_idx, + n, + sum(s["success"] for s in status), + n, + elapsed_frame, + ) + + return [ + { + "success": s["success"], + "steps": s["steps"], + "elapsed_sec": round(elapsed_frame / max(1, n), 3), + "task_desc": s["task_desc"], + **({"_error": s["_error"]} if "_error" in s else {}), + } + for s in status + ] diff --git a/wall_x/_vendor/harrix/drivers/inproc/env_handle.py b/wall_x/_vendor/harrix/drivers/inproc/env_handle.py new file mode 100644 index 00000000..64bbec55 --- /dev/null +++ b/wall_x/_vendor/harrix/drivers/inproc/env_handle.py @@ -0,0 +1,152 @@ +"""Subprocess env handle used by the in-process driver. + +The model remains in the driver process. Each env subprocess receives reset and +execute-chunk commands through a pipe. +""" + +from __future__ import annotations + +import multiprocessing as mp +import os +import random + +import numpy as np + + +def _subproc_main(cfg, worker_id, child_conn): + """Env subprocess entry point.""" + # Spawned children inherit env vars, but set these explicitly for launchers + # that did not configure them. Robosuite validates the EGL id against the + # CUDA_VISIBLE_DEVICES environment string, so keep the same visible id here. + cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") or "0" + os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices + # EGL device index is always 0 after CUDA_VISIBLE_DEVICES remapping. + os.environ["MUJOCO_EGL_DEVICE_ID"] = os.environ.get("MUJOCO_EGL_DEVICE_ID") or "0" + + # Env workers run robosuite only. Do not import torch here; otherwise many + # subprocesses may initialize CUDA contexts and compete with the driver model. + seed = cfg.env.seed + worker_id + np.random.seed(seed) + random.seed(seed) + os.environ["PYTHONHASHSEED"] = str(seed) + + import wall_x._vendor.harrix.envs # noqa: F401 trigger register + from wall_x._vendor.harrix.envs.registry import build_env + + env = build_env(cfg, worker_id) + try: + while True: + try: + cmd, args = child_conn.recv() + except EOFError: + break + try: + if cmd == "reset_episode": + result = env.reset_episode(tuple(args)) + elif cmd == "execute_chunk": + result = env.execute_chunk(args) + elif cmd == "finalize_episode": + env.finalize_episode(bool(args)) + result = None + elif cmd == "shutdown": + env.shutdown() + child_conn.send(("ok", None)) + break + else: + raise ValueError(f"unknown cmd {cmd!r}") + child_conn.send(("ok", result)) + except Exception as e: + import traceback + + child_conn.send(("err", f"{e}\n{traceback.format_exc()}")) + finally: + try: + child_conn.close() + except Exception: + pass + + +class SubprocEnvHandle: + """Synchronous pipe wrapper around one env subprocess. + + Usage: + h.submit_reset(ep_id) + ... + h.wait_reset() # -> {"obs", "instruction", "task_desc"} + h.submit_execute_chunk(actions) + ... + h.wait_execute_chunk() # -> {"obs", "done", "steps"} + """ + + def __init__(self, cfg, worker_id: int): + # Use spawn instead of fork because the driver may already hold a CUDA + # context for the model. + ctx = mp.get_context("spawn") + self._parent_conn, child_conn = ctx.Pipe() + self._proc = ctx.Process( + target=_subproc_main, + args=(cfg, worker_id, child_conn), + daemon=False, + name=f"infer-subproc-w{worker_id}", + ) + self._proc.start() + child_conn.close() + self._has_pending = False + + def _send(self, cmd: str, args): + self._parent_conn.send((cmd, args)) + self._has_pending = True + + def _recv(self): + if not self._has_pending: + raise RuntimeError("no pending request to wait for") + status, payload = self._parent_conn.recv() + self._has_pending = False + if status == "err": + raise RuntimeError(f"subproc env error (w={self._proc.name}): {payload}") + return payload + + def submit_reset(self, ep_id): + self._send("reset_episode", ep_id) + + def wait_reset(self): + return self._recv() + + def submit_execute_chunk(self, actions): + self._send("execute_chunk", np.asarray(actions, dtype=np.float32)) + + def wait_execute_chunk(self): + return self._recv() + + def submit_finalize_episode(self, success: bool): + self._send("finalize_episode", success) + + def wait_finalize_episode(self): + return self._recv() + + def finalize_episode(self, success: bool): + """Blocking convenience wrapper.""" + self.submit_finalize_episode(success) + return self.wait_finalize_episode() + + def reset_episode(self, ep_id): + """Blocking convenience wrapper.""" + self.submit_reset(ep_id) + return self.wait_reset() + + def execute_chunk(self, actions): + """Blocking convenience wrapper.""" + self.submit_execute_chunk(actions) + return self.wait_execute_chunk() + + def shutdown(self) -> None: + try: + self._send("shutdown", None) + self._recv() + except Exception: + pass + if self._proc.is_alive(): + self._proc.join(timeout=5) + if self._proc.is_alive(): + self._proc.terminate() + self._proc.join(timeout=2) diff --git a/wall_x/_vendor/harrix/drivers/inproc/model_handle.py b/wall_x/_vendor/harrix/drivers/inproc/model_handle.py new file mode 100644 index 00000000..ada020b6 --- /dev/null +++ b/wall_x/_vendor/harrix/drivers/inproc/model_handle.py @@ -0,0 +1,25 @@ +"""In-process model handle. + +The model is constructed in the driver process and calls the adapter directly. +""" + +from __future__ import annotations + + +class DirectModelHandle: + def __init__(self, cfg): + # Trigger adapter registration. + import wall_x._vendor.harrix.adapters # noqa: F401 + from wall_x._vendor.harrix.adapters.registry import build_adapter + + self._adapter = build_adapter(cfg) + + @property + def chunk_horizon(self) -> int: + return self._adapter.chunk_horizon + + def predict_batch(self, payloads): + return self._adapter.predict_batch(payloads) + + def shutdown(self) -> None: + pass diff --git a/wall_x/_vendor/harrix/drivers/job_state.py b/wall_x/_vendor/harrix/drivers/job_state.py new file mode 100644 index 00000000..0a9b11cf --- /dev/null +++ b/wall_x/_vendor/harrix/drivers/job_state.py @@ -0,0 +1,343 @@ +"""Episode work queue with JSONL persistence. + +Drivers enumerate episode ids into ``pending``; env handles claim work and +return results through ``complete`` or ``fail``. Every state transition is +appended to ``state.jsonl``, so completed episodes can be skipped on restart. + +Two scheduling modes are supported: +- FIFO: workers claim the next pending episode. +- batch_sync: episodes are grouped into task-local frames. A new frame is not + released until the previous frame is complete, which gives deterministic + lockstep batches at the cost of possible idle workers. +""" + +from __future__ import annotations + +import json +import os +import threading +import time +from typing import Optional + + +class JobState: + def __init__( + self, + all_episodes: list[tuple], + log_path: str, + batch_sync_mode: bool = False, + batch_size: int = 1, + ): + """ + all_episodes: list of (suite_name, task_idx, ep_idx) tuples + log_path: state.jsonl path; existing completed episodes are skipped + batch_sync_mode: enable task-local frame barriers + batch_size: number of episodes per frame + """ + self._lock = threading.Lock() + self._log_path = log_path + self._log_fh = None + self._batch_sync_mode = bool(batch_sync_mode) + self._batch_size = int(batch_size) + + completed_set = ( + self._load_completed(log_path) if os.path.exists(log_path) else set() + ) + remaining: list[tuple] = [ + tuple(ep) for ep in all_episodes if tuple(ep) not in completed_set + ] + + self._in_progress: dict[tuple, int] = {} + self._completed: dict[tuple, dict] = {} + self._failed: dict[tuple, str] = {} + + if self._batch_sync_mode: + self._frames: list[dict] = self._build_frames(remaining, self._batch_size) + self._cur_frame_idx: int = 0 + self._cur_frame_inflight: int = 0 + self._pending = None + else: + self._frames = [] + self._cur_frame_idx = 0 + self._cur_frame_inflight = 0 + self._pending: list[tuple] = remaining + + os.makedirs(os.path.dirname(log_path) or ".", exist_ok=True) + self._log_fh = open(log_path, "a") + + log_entry = { + "event": "session_start", + "pending": ( + len(self._pending) + if not self._batch_sync_mode + else sum(len(f["eps"]) for f in self._frames) + ), + "skipped_completed": len(completed_set), + "batch_sync_mode": self._batch_sync_mode, + } + if self._batch_sync_mode: + log_entry["num_frames"] = len(self._frames) + log_entry["batch_size"] = self._batch_size + self._append_log(log_entry) + + @staticmethod + def _build_frames(remaining: list[tuple], batch_size: int) -> list[dict]: + """Group episodes by ``(suite, task_idx)`` into fixed-size frames. + + Partial tail frames are kept. + """ + frames: list[dict] = [] + cur_key: Optional[tuple] = None + cur_bucket: list[tuple] = [] + for ep in remaining: + key = (ep[0], ep[1]) + if key != cur_key and cur_bucket: + for i in range(0, len(cur_bucket), batch_size): + frames.append( + {"eps": cur_bucket[i : i + batch_size], "next_slot": 0} + ) + cur_bucket = [] + cur_key = key + cur_bucket.append(ep) + if cur_bucket: + for i in range(0, len(cur_bucket), batch_size): + frames.append({"eps": cur_bucket[i : i + batch_size], "next_slot": 0}) + return frames + + @staticmethod + def _load_completed(log_path: str) -> set[tuple]: + completed = set() + with open(log_path, "r") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + rec = json.loads(line) + except json.JSONDecodeError: + continue + if rec.get("event") == "completed" and "ep_id" in rec: + completed.add(tuple(rec["ep_id"])) + return completed + + def _append_log(self, extra: dict): + rec = {"ts": time.strftime("%Y-%m-%dT%H:%M:%S"), **extra} + if "ep_id" in rec and isinstance(rec["ep_id"], tuple): + rec["ep_id"] = list(rec["ep_id"]) + self._log_fh.write(json.dumps(rec, ensure_ascii=False) + "\n") + self._log_fh.flush() + + # ---- claim API ---- + + def claim(self, worker_id: int) -> Optional[list]: + """Claim one episode id. + + FIFO returns ``None`` when the pending queue is empty. In batch-sync + mode, ``None`` can also mean the current frame has been fully issued but + still has in-flight episodes. + """ + with self._lock: + if self._batch_sync_mode: + return self._claim_sync(worker_id) + return self._claim_fifo(worker_id) + + def claim_frame(self) -> list[list]: + """Claim one complete frame for lockstep in-process evaluation. + + In FIFO mode this returns up to ``batch_size`` pending episodes. + """ + with self._lock: + if self._batch_sync_mode: + return self._claim_frame_sync() + return self._claim_frame_fifo() + + def _claim_fifo(self, worker_id: int) -> Optional[list]: + if not self._pending: + return None + ep = self._pending.pop(0) + self._in_progress[ep] = worker_id + self._append_log({"event": "claimed", "ep_id": ep, "worker": worker_id}) + return list(ep) + + def _claim_sync(self, worker_id: int) -> Optional[list]: + while self._cur_frame_idx < len(self._frames): + frame = self._frames[self._cur_frame_idx] + if frame["next_slot"] < len(frame["eps"]): + ep = frame["eps"][frame["next_slot"]] + frame["next_slot"] += 1 + self._cur_frame_inflight += 1 + self._in_progress[ep] = worker_id + self._append_log( + { + "event": "claimed", + "ep_id": ep, + "worker": worker_id, + "frame": self._cur_frame_idx, + } + ) + return list(ep) + if self._cur_frame_inflight > 0: + return None + self._append_log( + { + "event": "frame_done", + "frame": self._cur_frame_idx, + "size": len(frame["eps"]), + } + ) + self._cur_frame_idx += 1 + return None + + def _claim_frame_sync(self) -> list[list]: + # Advance completed frames. + while self._cur_frame_idx < len(self._frames): + frame = self._frames[self._cur_frame_idx] + if frame["next_slot"] < len(frame["eps"]): + break + if self._cur_frame_inflight > 0: + # The previous frame still has in-flight episodes. + return [] + self._append_log( + { + "event": "frame_done", + "frame": self._cur_frame_idx, + "size": len(frame["eps"]), + } + ) + self._cur_frame_idx += 1 + if self._cur_frame_idx >= len(self._frames): + return [] + frame = self._frames[self._cur_frame_idx] + out: list[list] = [] + while frame["next_slot"] < len(frame["eps"]): + ep = frame["eps"][frame["next_slot"]] + frame["next_slot"] += 1 + self._cur_frame_inflight += 1 + self._in_progress[ep] = -1 + self._append_log( + {"event": "claimed", "ep_id": ep, "frame": self._cur_frame_idx} + ) + out.append(list(ep)) + return out + + def _claim_frame_fifo(self) -> list[list]: + if not self._pending: + return [] + n = min(self._batch_size, len(self._pending)) + out: list[list] = [] + for _ in range(n): + ep = self._pending.pop(0) + self._in_progress[ep] = -1 + self._append_log({"event": "claimed", "ep_id": ep}) + out.append(list(ep)) + return out + + # ---- complete / fail ---- + + def complete(self, ep_id: list, result: dict) -> None: + ep = tuple(ep_id) + with self._lock: + if ep in self._in_progress: + self._in_progress.pop(ep, None) + if self._batch_sync_mode: + self._cur_frame_inflight = max(0, self._cur_frame_inflight - 1) + self._completed[ep] = result + self._append_log({"event": "completed", "ep_id": ep, **result}) + + def fail(self, ep_id: list, error: str) -> None: + ep = tuple(ep_id) + with self._lock: + if ep in self._in_progress: + self._in_progress.pop(ep, None) + if self._batch_sync_mode: + self._cur_frame_inflight = max(0, self._cur_frame_inflight - 1) + self._failed[ep] = error + self._append_log({"event": "failed", "ep_id": ep, "error": error}) + + # ---- status queries ---- + + def get_frame_inflight(self) -> int: + with self._lock: + return self._cur_frame_inflight if self._batch_sync_mode else 0 + + def is_drained(self) -> bool: + with self._lock: + if self._batch_sync_mode: + return ( + self._cur_frame_idx >= len(self._frames) + and self._cur_frame_inflight == 0 + ) + return len(self._pending) == 0 and len(self._in_progress) == 0 + + def progress(self) -> dict: + with self._lock: + base = { + "in_progress": len(self._in_progress), + "completed": len(self._completed), + "failed": len(self._failed), + } + if self._batch_sync_mode: + pending = sum( + len(f["eps"]) - f["next_slot"] + for f in self._frames[self._cur_frame_idx :] + ) + base["pending"] = pending + base["frame"] = f"{self._cur_frame_idx}/{len(self._frames)}" + base["frame_inflight"] = self._cur_frame_inflight + else: + base["pending"] = len(self._pending) + return base + + def dump_final(self, report_path: str) -> None: + with self._lock: + per_task: dict[tuple, dict] = {} + for ep, result in self._completed.items(): + key = (ep[0], ep[1]) + d = per_task.setdefault( + key, {"attempted": 0, "successes": 0, "steps": []} + ) + d["attempted"] += 1 + if result.get("success"): + d["successes"] += 1 + if "steps" in result: + d["steps"].append(result["steps"]) + for ep in self._failed: + key = (ep[0], ep[1]) + d = per_task.setdefault( + key, {"attempted": 0, "successes": 0, "steps": []} + ) + d["attempted"] += 1 + + total_attempted = sum(d["attempted"] for d in per_task.values()) + total_successes = sum(d["successes"] for d in per_task.values()) + overall_rate = total_successes / max(1, total_attempted) + + report = { + "overall": { + "attempted": total_attempted, + "successes": total_successes, + "success_rate": overall_rate, + "failed": len(self._failed), + }, + "per_task": { + f"{suite}_t{task_idx}": { + **d, + "success_rate": d["successes"] / max(1, d["attempted"]), + "avg_steps": ( + (sum(d["steps"]) / max(1, len(d["steps"]))) + if d["steps"] + else None + ), + } + for (suite, task_idx), d in sorted(per_task.items()) + }, + } + with open(report_path, "w") as f: + json.dump(report, f, indent=2, ensure_ascii=False) + self._append_log( + { + "event": "session_end", + "report_path": report_path, + "overall_success_rate": overall_rate, + } + ) diff --git a/wall_x/_vendor/harrix/envs/__init__.py b/wall_x/_vendor/harrix/envs/__init__.py new file mode 100644 index 00000000..605263f5 --- /dev/null +++ b/wall_x/_vendor/harrix/envs/__init__.py @@ -0,0 +1,8 @@ +"""Environment registrations for harrix.""" + +_LIBERO_IMPORT_ERROR = None + +try: + from wall_x._vendor.harrix.envs import libero # noqa: F401 +except ModuleNotFoundError as exc: + _LIBERO_IMPORT_ERROR = exc diff --git a/wall_x/_vendor/harrix/envs/base.py b/wall_x/_vendor/harrix/envs/base.py new file mode 100644 index 00000000..5c3aef33 --- /dev/null +++ b/wall_x/_vendor/harrix/envs/base.py @@ -0,0 +1,119 @@ +"""Base environment abstraction. + +Two execution granularities are supported: + +- Episode-level ``run_episode``: caller supplies a predict callback and the env + owns the full episode loop. +- Chunk-level ``reset_episode`` + ``execute_chunk``: caller runs the model + between chunks and feeds action chunks back to the env. + +Subclasses must implement the chunk-level primitives. The default episode loop +is built on top of those primitives. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Callable + +import numpy as np + +from wall_x._vendor.harrix.eval_config import EvalConfig + + +class BaseEnv(ABC): + + @abstractmethod + def __init__(self, cfg: EvalConfig, worker_id: int) -> None: + """Perform env-specific setup.""" + + @classmethod + @abstractmethod + def enumerate_episodes(cls, cfg: EvalConfig) -> list[tuple]: + """Return episode ids to seed JobState before env instances are built. + + JobState treats the returned tuples as opaque ids. + """ + + @property + @abstractmethod + def robot_spec(self) -> dict: + """Return robot metadata for driver/checkpoint validation. + + Fields: + - dof_layout: dict[str, int] + - cam_names: list[str] + - norm_key: str + """ + + # ---- chunk-level primitives ---- + + @abstractmethod + def reset_episode(self, ep_id: tuple) -> dict: + """Start a new episode and return the first fresh observation. + + Returns {"obs": dict, "instruction": str, "task_desc": str (optional)}. + """ + + @abstractmethod + def execute_chunk(self, actions: np.ndarray) -> dict: + """Execute one action chunk with shape ``(H, action_dim)``. + + Returns {"obs": dict, "done": bool, "steps": int}. + """ + + def shutdown(self) -> None: + """Release resources. Subclasses may override.""" + + # ---- episode-level default implementation ---- + + def run_episode( + self, + ep_id: tuple, + predict: Callable[[dict, str, int], np.ndarray], + ) -> dict: + """Episode loop built from reset, predict, and execute_chunk. + + ``predict(observation, instruction, step)`` returns one action chunk. + The env remains unaware of the model transport. + """ + import time + + from wall_x._vendor.harrix.envs.libero_common import encode_raw_obs + + t_ep_start = time.time() + initial = self.reset_episode(ep_id) + obs = initial["obs"] + instruction = initial["instruction"] + task_desc = initial.get("task_desc", "") + + max_rounds = self._max_infer_rounds() + success = False + steps_total = 0 + for round_idx in range(max_rounds): + encoded = encode_raw_obs(obs) + chunk = predict(encoded, instruction, round_idx) + result = self.execute_chunk(chunk) + obs = result["obs"] + steps_total += result["steps"] + if result["done"]: + success = True + break + + self.finalize_episode(success) + return { + "success": bool(success), + "steps": steps_total, + "elapsed_sec": round(time.time() - t_ep_start, 3), + "task_desc": task_desc, + } + + def finalize_episode(self, success: bool) -> None: + """Hook for env-specific cleanup after an episode (e.g. save rollouts).""" + + def _max_infer_rounds(self) -> int: + """Return the maximum number of model chunks for one episode.""" + raise NotImplementedError( + f"{type(self).__name__} must override _max_infer_rounds when using " + "the default run_episode implementation" + ) diff --git a/wall_x/_vendor/harrix/envs/libero.py b/wall_x/_vendor/harrix/envs/libero.py new file mode 100644 index 00000000..885d51ee --- /dev/null +++ b/wall_x/_vendor/harrix/envs/libero.py @@ -0,0 +1,342 @@ +"""LIBERO environment implementation. + +The chunk-level API mirrors standard LIBERO rollout semantics: +- reset_episode enables rendering, sets initial state, and performs warmup steps. +- execute_chunk may skip intermediate image rendering but re-enables rendering + before returning an observation. +""" + +from __future__ import annotations + +import gc +import logging +import os +from typing import Optional + +import numpy as np + +from wall_x._vendor.harrix.envs.base import BaseEnv +from wall_x._vendor.harrix.envs.libero_common import ( + get_rollout_frame, + model_action_to_libero_env, + save_rollout_video, +) +from wall_x._vendor.harrix.envs.libero_sim import ( + create_libero_engine, + find_image_observables, + get_instruction, + get_libero_dummy_action, + get_task_suite, + load_initial_states, + pick_initial_state, + resolve_task_info, + set_render_enabled, +) +from wall_x._vendor.harrix.envs.registry import register_env +from wall_x._vendor.harrix.eval_config import EvalConfig + +logger = logging.getLogger(__name__) + + +@register_env("libero") +class LiberoEnv(BaseEnv): + + def __init__(self, cfg: EvalConfig, worker_id: int) -> None: + self._cfg = cfg + self._libero_cfg = cfg.env.libero + self._worker_id = worker_id + self._seed = cfg.env.seed + self._task_suite_name = self._libero_cfg.task_suite_name + + # Task-suite metadata. + self._task_suite = get_task_suite(self._task_suite_name) + self._num_tasks = self._task_suite.n_tasks + self._custom_initial_states = load_initial_states( + self._libero_cfg.initial_states_path + ) + + # Robosuite engine, lazily rebuilt on task changes. + self._libero_env = None + self._current_task_id: Optional[int] = None + self._rebuild_env_per_episode = self._libero_cfg.rebuild_env_per_episode + + # Render-skip state. + self._skip_intermediate_render = self._libero_cfg.skip_intermediate_render + self._force_render_task_ids = set( + self._libero_cfg.force_render_task_indices or [] + ) + self._effective_skip_render = self._skip_intermediate_render + self._image_obs: list = [] + self._chunk_granular_render_toggle = ( + self._libero_cfg.chunk_granular_render_toggle + ) + self._render_enabled_state = False + + # Optional bit-alignment dump for debugging. + self._bit_dump_dir = os.environ.get("WALLX_BIT_DUMP_DIR", "").strip() or None + if self._bit_dump_dir: + os.makedirs(self._bit_dump_dir, exist_ok=True) + + rollout_dir = (self._libero_cfg.rollout_dir or "").strip() + if not rollout_dir: + rollout_dir = os.environ.get("WALLX_ROLLOUT_DIR", "").strip() + if rollout_dir and os.environ.get("WALLX_DISABLE_ROLLOUT", "0") == "1": + rollout_dir = "" + self._rollout_dir = rollout_dir or None + self._rollout_fps = int(self._libero_cfg.rollout_fps) + if self._rollout_dir: + worker_subdir = f"worker{worker_id}" if cfg.runtime.num_workers > 1 else "" + self._rollout_dir = os.path.join( + self._rollout_dir, + self._task_suite_name, + worker_subdir, + ) + os.makedirs(self._rollout_dir, exist_ok=True) + logger.info("Rollout MP4 saving enabled: %s", self._rollout_dir) + + # Per-episode state consumed by execute_chunk. + self._current_ep: Optional[tuple] = None + self._current_task_desc: str = "" + self._current_instruction: str = "" + self._chunk_counter_in_ep: int = 0 + self._last_obs_for_dump: Optional[dict] = None + self._replay_images: list[np.ndarray] = [] + self._rollout_saved = False + + # ---- BaseEnv API ---- + + @classmethod + def enumerate_episodes(cls, cfg: EvalConfig) -> list[tuple]: + libero_cfg = cfg.env.libero + suite = libero_cfg.task_suite_name + + if libero_cfg.task_indices is not None: + task_indices = [int(x) for x in libero_cfg.task_indices] + else: + ts = get_task_suite(suite) + task_indices = list(range(ts.n_tasks)) + + eps = [] + for tid in task_indices: + for epi in range(libero_cfg.num_trials_per_task): + eps.append((suite, tid, epi)) + return eps + + @property + def robot_spec(self) -> dict: + """Return robot spec for driver/adapter validation.""" + from wall_x._vendor.harrix.utils.train_config import ( + load_train_config_with_ckpt_overlay, + ) + + train_cfg = load_train_config_with_ckpt_overlay( + self._cfg.model.train_config_path, + self._cfg.model.checkpoint_path, + ) + return { + "dof_layout": train_cfg.get("dof_config", {}), + "cam_names": list(self._cfg.model.cam_names), + "norm_key": self._cfg.model.norm_key, + } + + def _max_infer_rounds(self) -> int: + return self._libero_cfg.max_infer_times + + def reset_episode(self, ep_id: tuple) -> dict: + suite, task_id, ep_idx = ep_id + if suite != self._task_suite_name: + raise ValueError( + f"env bound to suite={self._task_suite_name}, got ep with suite={suite}" + ) + + need_rebuild = self._rebuild_env_per_episode or self._current_task_id != task_id + if need_rebuild: + self._rebuild_env(task_id) + + task_desc, default_states = resolve_task_info(self._task_suite, task_id) + init_state = pick_initial_state( + self._libero_cfg.initial_states_path, + self._custom_initial_states, + task_desc, + default_states, + ep_idx, + ) + + if not need_rebuild: + self._libero_env.reset() + obs = self._libero_env.set_init_state(init_state) + if obs is None: + raise RuntimeError("set_init_state returned None") + + set_render_enabled(self._image_obs, True) + self._render_enabled_state = True + + dummy_action = get_libero_dummy_action() + for _ in range(10): + obs, _, _, _ = self._libero_env.step(dummy_action) + + self._current_ep = (task_id, ep_idx) + self._current_task_desc = task_desc + self._current_instruction = get_instruction(task_desc) + self._chunk_counter_in_ep = 0 + self._last_obs_for_dump = obs + self._begin_rollout_capture(obs) + + return { + "obs": obs, + "instruction": self._current_instruction, + "task_desc": task_desc, + } + + def execute_chunk(self, actions: np.ndarray) -> dict: + actions = np.asarray(actions, dtype=np.float32) + H = actions.shape[0] + + if self._bit_dump_dir and self._current_ep is not None: + self._dump_chunk_npz( + (self._task_suite_name, *self._current_ep), + self._chunk_counter_in_ep, + self._last_obs_for_dump or {}, + actions, + ) + self._chunk_counter_in_ep += 1 + + skip_render = self._chunk_skip_render() + # Disable rendering at the chunk start when render-skip is enabled. + if skip_render and self._image_obs: + if self._render_enabled_state: + set_render_enabled(self._image_obs, False) + self._render_enabled_state = False + + last_obs = None + done = False + steps = 0 + for step_idx in range(H): + # Re-enable rendering before the final step to return a fresh image. + if ( + skip_render + and self._image_obs + and step_idx == H - 1 + and not self._render_enabled_state + ): + set_render_enabled(self._image_obs, True) + self._render_enabled_state = True + + action = model_action_to_libero_env(actions[step_idx].reshape(-1)) + obs, _, done_flag, _ = self._libero_env.step(action) + last_obs = obs + steps += 1 + self._append_rollout_frame(obs) + if bool(done_flag): + done = True + # Do not add an extra simulator step on early success; just make + # sure future rendering is enabled. + if skip_render and self._image_obs and not self._render_enabled_state: + set_render_enabled(self._image_obs, True) + self._render_enabled_state = True + break + + self._last_obs_for_dump = last_obs + return {"obs": last_obs, "done": done, "steps": steps} + + def finalize_episode(self, success: bool) -> None: + self._save_episode_rollout(success) + + def shutdown(self) -> None: + if self._libero_env is not None: + try: + self._libero_env.close() + except Exception: + pass + self._libero_env = None + + # ---- internals ---- + + def _chunk_skip_render(self) -> bool: + """Skip intermediate renders unless rollout MP4 saving needs every frame.""" + return self._effective_skip_render and self._rollout_dir is None + + def _begin_rollout_capture(self, obs: dict | None) -> None: + self._replay_images = [] + self._rollout_saved = False + if self._rollout_dir and obs is not None: + if self._image_obs and not self._render_enabled_state: + set_render_enabled(self._image_obs, True) + self._render_enabled_state = True + self._replay_images.append(get_rollout_frame(obs)) + + def _append_rollout_frame(self, obs: dict | None) -> None: + if self._rollout_dir and obs is not None: + self._replay_images.append(get_rollout_frame(obs)) + + def _save_episode_rollout(self, success: bool) -> None: + if ( + not self._rollout_dir + or self._rollout_saved + or not self._replay_images + or self._current_ep is None + ): + return + task_id, ep_idx = self._current_ep + try: + mp4_path = save_rollout_video( + self._rollout_dir, + self._replay_images, + task_id=task_id, + episode_idx=ep_idx, + success=success, + task_description=self._current_task_desc, + fps=self._rollout_fps, + ) + self._rollout_saved = True + if mp4_path: + logger.info("Saved rollout MP4: %s", mp4_path) + except Exception as exc: + logger.warning( + "Failed to save rollout MP4 for task%d ep%d: %s", + task_id, + ep_idx, + exc, + ) + + def _rebuild_env(self, task_id: int) -> None: + if self._libero_env is not None: + try: + self._libero_env.close() + except Exception: + pass + self._libero_env = None + gc.collect() + + self._libero_env = create_libero_engine( + task_id=task_id, + task_suite_name=self._task_suite_name, + resolution=256, + seed=self._seed, + ) + self._current_task_id = task_id + self._image_obs = find_image_observables(self._libero_env) + self._effective_skip_render = ( + self._skip_intermediate_render + and task_id not in self._force_render_task_ids + ) + + def _dump_chunk_npz( + self, ep_id, chunk_idx: int, raw_obs: dict, chunk_actions: np.ndarray + ) -> None: + _, task_id, ep_idx = ep_id + path = os.path.join( + self._bit_dump_dir, f"t{task_id}_ep{ep_idx}_c{chunk_idx}.npz" + ) + fields = {"action_chunk": np.asarray(chunk_actions, dtype=np.float32)} + for k in ( + "robot0_eef_pos", + "robot0_eef_quat", + "robot0_gripper_qpos", + "agentview_image", + "robot0_eye_in_hand_image", + ): + v = raw_obs.get(k) + if v is not None: + fields[k] = np.asarray(v) + np.savez_compressed(path, **fields) diff --git a/wall_x/_vendor/harrix/envs/libero_common.py b/wall_x/_vendor/harrix/envs/libero_common.py new file mode 100644 index 00000000..2bfa2f35 --- /dev/null +++ b/wall_x/_vendor/harrix/envs/libero_common.py @@ -0,0 +1,359 @@ +"""Env-adapter IO helpers for single-arm LIBERO tasks. + +Env code extracts a compact ndarray payload from raw LIBERO observations. +Adapters then build proprioception, masks, and 7-dof right-arm action chunks +from that payload. This module intentionally depends only on NumPy so adapter +processes can import it without importing robosuite or LIBERO. +""" + +from __future__ import annotations + +import math +import os + +import numpy as np + + +# Fallback single-arm LIBERO dof_config when train_config does not provide one. +_LIBERO_FALLBACK_DOF_CONFIG = { + "follow_right_ee_cartesian_pos": 3, + "follow_right_ee_rotation": 3, + "follow_right_gripper": 1, +} + +_LIBERO_FALLBACK_AGENT_POS_CONFIG = dict(_LIBERO_FALLBACK_DOF_CONFIG) +_VIRTUAL_TAIL_KEYS = frozenset(("action_padding",)) + + +def _resolve_dof_config(train_config: dict) -> dict: + return ( + train_config.get("dof_config") + or train_config.get("task", {}).get("dof_config") + or _LIBERO_FALLBACK_DOF_CONFIG + ) + + +def _resolve_agent_pos_config(train_config: dict) -> dict: + return ( + train_config.get("agent_pos_config") + or train_config.get("task", {}).get("agent_pos_config") + or _LIBERO_FALLBACK_AGENT_POS_CONFIG + ) + + +def _move_virtual_keys_to_tail(layout: dict) -> dict: + head = {k: v for k, v in layout.items() if k not in _VIRTUAL_TAIL_KEYS} + tail = {k: v for k, v in layout.items() if k in _VIRTUAL_TAIL_KEYS} + return {**head, **tail} + + +def _effective_agent_pos_config(train_config: dict, state_values: dict) -> dict: + config = _move_virtual_keys_to_tail(dict(_resolve_agent_pos_config(train_config))) + gripper_key = next( + ( + key + for key in config + if key.replace("follow_", "").replace("master_", "") == "right_gripper" + ), + None, + ) + if gripper_key is None: + return config + + old_dim = int(config[gripper_key]) + target_dim = None + override = os.environ.get("WALLX_LIBERO_STATE_GRIPPER_DIM") + if override: + target_dim = int(override) + elif os.environ.get("WALLX_LIBERO_AUTO_STATE_GRIPPER_DIM", "1") != "0": + norm_dim = int(train_config.get("_libero_proprio_norm_dim") or 0) + real_dim = sum(v for k, v in config.items() if k not in _VIRTUAL_TAIL_KEYS) + if norm_dim == real_dim + 1: + target_dim = old_dim + 1 + + available_dim = state_values["right_gripper"].shape[1] + if target_dim is None or target_dim == old_dim or target_dim > available_dim: + return config + + config[gripper_key] = target_dim + delta = target_dim - old_dim + if "action_padding" in config: + config["action_padding"] = max(0, int(config["action_padding"]) - delta) + return config + + +def _build_right_arm_state_values(obs_ndarrays: dict) -> dict[str, np.ndarray]: + """Bare-key state tensors for single-arm LIBERO proprio construction.""" + rot3 = np.asarray(obs_ndarrays["eef_axisangle"], dtype=np.float32).reshape(1, 3) + values: dict[str, np.ndarray] = { + "right_ee_cartesian_pos": np.asarray( + obs_ndarrays["eef_pos"], dtype=np.float32 + ).reshape(1, 3), + "right_ee_rotation": rot3, + "right_gripper": np.asarray(obs_ndarrays["gripper"], dtype=np.float32).reshape( + 1, -1 + ), + } + from wall_x._vendor.x2robot_utils.geometry import euler_to_matrix_zyx_6d_nb + + rot6d = euler_to_matrix_zyx_6d_nb(rot3.astype(np.float64)).reshape(1, 6) + values["right_ee_rotation_6D"] = rot6d.astype(np.float32) + return values + + +# Auxiliary action keys that should be masked out for single-arm LIBERO. +_DOF_MASK_ZERO_KEYS = frozenset( + ( + "follow_left_ee_cartesian_pos", + "follow_left_ee_rotation", + "follow_left_ee_rotation_6D", + "follow_left_gripper", + "head_actions", + "height", + "velocity_decomposed", + "action_padding", + ) +) + + +# ============================================================ +# LIBERO raw observation decoding helpers. +# ============================================================ + + +def _get_libero_image(obs: dict) -> np.ndarray: + """Return the third-person camera image, rotated to match preprocessing.""" + return obs["agentview_image"][::-1, ::-1] + + +def get_rollout_frame(obs: dict) -> np.ndarray: + """Return one RGB frame for rollout MP4 saving.""" + return np.asarray(_get_libero_image(obs), dtype=np.uint8) + + +def _get_libero_wrist_image(obs: dict) -> np.ndarray: + """Return the wrist camera image, rotated to match preprocessing.""" + return obs["robot0_eye_in_hand_image"][::-1, ::-1] + + +def _quat2axisangle(quat) -> np.ndarray: + """Convert an xyzw quaternion to a 3D axis-angle vector.""" + if quat[3] > 1.0: + quat[3] = 1.0 + elif quat[3] < -1.0: + quat[3] = -1.0 + + den = np.sqrt(1.0 - quat[3] * quat[3]) + if math.isclose(den, 0.0): + return np.zeros(3) + + return (quat[:3] * 2.0 * math.acos(quat[3])) / den + + +# ============================================================ +# Env-side to adapter-side observation encoding. +# ============================================================ + + +def encode_raw_obs(raw_obs: dict) -> dict: + """Extract the minimal ndarray payload from a raw LIBERO observation. + + The payload contains three 1-D state arrays and two rotated image arrays. + """ + if "agentview_image" not in raw_obs: + raise KeyError( + "agentview_image missing in raw_obs; render-skip may have returned " + "a stale observation" + ) + return { + "eef_pos": np.asarray(raw_obs["robot0_eef_pos"], dtype=np.float32), + "eef_axisangle": np.asarray( + _quat2axisangle(raw_obs["robot0_eef_quat"]), dtype=np.float32 + ), + "gripper": np.asarray(raw_obs["robot0_gripper_qpos"], dtype=np.float32), + "face_view": _get_libero_image(raw_obs), + "wrist_view": _get_libero_wrist_image(raw_obs), + } + + +# ============================================================ +# Adapter-side proprioception and mask construction. +# ============================================================ + + +def encode_proprio( + obs_ndarrays: dict, + train_config: dict, + action_horizon: int, +) -> dict: + """Convert a single-arm LIBERO ndarray payload into model input fields. + + Returned fields include proprioception, agent_pos_mask, dof_mask, + face_view, and right_wrist_view. + """ + state_values = _build_right_arm_state_values(obs_ndarrays) + agent_pos_config = _effective_agent_pos_config(train_config, state_values) + dof_config = _move_virtual_keys_to_tail(dict(_resolve_dof_config(train_config))) + + propri_parts: list[np.ndarray] = [] + mask_parts: list[np.ndarray] = [] + for key, dim in agent_pos_config.items(): + bare = key.replace("follow_", "").replace("master_", "") + if bare in state_values: + v = state_values[bare] + if bare == "right_gripper" and v.shape[1] > dim: + v = v[:, :dim] + if v.shape[1] != dim: + raise ValueError( + f"agent_pos_config[{key!r}]={dim} does not match " + f"observation dimension {v.shape[1]}" + ) + propri_parts.append(v) + mask_parts.append(np.ones((1, dim), dtype=np.float32)) + else: + propri_parts.append(np.zeros((1, dim), dtype=np.float32)) + mask_parts.append(np.zeros((1, dim), dtype=np.float32)) + + # (1, 1, D) + proprioception = np.concatenate(propri_parts, axis=1)[None] + agent_pos_mask = np.concatenate(mask_parts, axis=1)[None] + + # dof_mask: (1, T, D_action) + total_dof = sum(dof_config.values()) + dof_mask = np.ones((1, action_horizon, total_dof)) + start = 0 + for key, dim in dof_config.items(): + if key in _DOF_MASK_ZERO_KEYS: + dof_mask[:, :, start : start + dim] = 0 + start += dim + + return { + "proprioception": proprioception.astype(np.float32), + "agent_pos_mask": agent_pos_mask.astype(np.float32), + "dof_mask": dof_mask, + "face_view": obs_ndarrays["face_view"], + "right_wrist_view": obs_ndarrays["wrist_view"], + } + + +# ============================================================ +# Adapter-side action decoding. +# ============================================================ + + +def decode_chunk(predict_action: np.ndarray, train_config: dict) -> np.ndarray: + """Extract a 7-dof right-arm chunk from model action output.""" + if predict_action.ndim == 3: + predict_action = predict_action[0] + + dof_config = _resolve_dof_config(train_config) + slices: dict[str, slice] = {} + start = 0 + for key, dim in dof_config.items(): + bare = key.replace("follow_", "").replace("master_", "") + slices[bare] = slice(start, start + dim) + start += dim + + pos = predict_action[:, slices["right_ee_cartesian_pos"]] + grip = predict_action[:, slices["right_gripper"]] + + if "right_ee_rotation_6D" in slices: + from wall_x._vendor.x2robot_utils.geometry import so3_to_euler_zyx_batch_nb + + rot6d = np.asarray( + predict_action[:, slices["right_ee_rotation_6D"]], dtype=np.float64 + ) + rot = so3_to_euler_zyx_batch_nb(rot6d).astype(np.float32) + elif "right_ee_rotation" in slices: + rot = predict_action[:, slices["right_ee_rotation"]] + else: + raise KeyError( + "dof_config has no right-arm rotation slice " + f"(keys={list(slices.keys())})" + ) + + return np.concatenate([pos, rot, grip], axis=1) + + +def gripper_model_to_libero_osc(grip_2d: np.ndarray) -> float: + """Map model gripper output to robosuite OSC_POSE gripper command in [-1, 1].""" + g = np.asarray(grip_2d, dtype=np.float64).reshape(-1) + if g.size == 0: + raise ValueError("empty gripper action") + cmd = float(g[0]) + if os.environ.get("WALLX_LIBERO_GRIPPER_BINARIZE", "0") == "1": + if abs(cmd) < 1e-6: + cmd = -1.0 + else: + cmd = float(np.sign(cmd)) + if os.environ.get("WALLX_LIBERO_INVERT_GRIPPER", "0") == "1": + cmd *= -1.0 + return cmd + + +def _sanitize_task_description(task_description: str, max_len: int = 50) -> str: + return ( + task_description.lower() + .replace(" ", "_") + .replace("\n", "_") + .replace(".", "_")[:max_len] + ) + + +def save_rollout_video( + rollout_dir: str, + rollout_images: list[np.ndarray], + *, + task_id: int, + episode_idx: int, + success: bool, + task_description: str, + fps: int = 30, +) -> str | None: + """Save an MP4 replay of one LIBERO episode.""" + if not rollout_images: + return None + + import imageio + + os.makedirs(rollout_dir, exist_ok=True) + task_slug = _sanitize_task_description(task_description) + mp4_path = os.path.join( + rollout_dir, + f"task{task_id}_ep{episode_idx}--success={int(success)}--{task_slug}.mp4", + ) + writer = imageio.get_writer(mp4_path, fps=fps, macro_block_size=1) + try: + for img in rollout_images: + writer.append_data(np.asarray(img, dtype=np.uint8)) + finally: + writer.close() + return mp4_path + + +def model_action_to_libero_env(action: np.ndarray) -> np.ndarray: + """Convert model output to the robosuite OSC_POSE 7D action. + + Internal LIBERO evaluation passes the model's 7D chunk directly to + ``env.step``. Keep that as the public default; conversion modes remain + available only for ablations through environment variables. + """ + from scipy.spatial.transform import Rotation as R + + a = np.asarray(action, dtype=np.float64).reshape(-1) + if a.size not in (7, 8): + raise ValueError(f"expected 7D or 8D model action, got shape {a.shape}") + + pos_delta = a[:3] + rot_mode = os.environ.get("WALLX_LIBERO_ROT_MODE", "direct").strip() + if rot_mode == "euler_zyx_to_rotvec": + rot_aa = R.from_euler("zyx", a[3:6]).as_rotvec() + elif rot_mode == "direct": + rot_aa = a[3:6] + else: + raise ValueError( + "WALLX_LIBERO_ROT_MODE must be 'euler_zyx_to_rotvec' or 'direct', " + f"got {rot_mode!r}" + ) + grip = gripper_model_to_libero_osc(a[6:8] if a.size == 8 else a[6:7]) + return np.concatenate([pos_delta, rot_aa, np.array([grip], dtype=np.float64)]) diff --git a/wall_x/_vendor/harrix/envs/libero_sim.py b/wall_x/_vendor/harrix/envs/libero_sim.py new file mode 100644 index 00000000..ae510465 --- /dev/null +++ b/wall_x/_vendor/harrix/envs/libero_sim.py @@ -0,0 +1,181 @@ +"""LIBERO benchmark and robosuite engine helpers. + +This module may import robosuite and LIBERO. Adapter-side code should use +``libero_common.py`` instead, which only depends on NumPy. +""" + +from __future__ import annotations + +import json +import logging +import os +from typing import Any, Optional + +import numpy as np +from robosuite.wrappers import VisualizationWrapper + +logger = logging.getLogger(__name__) + + +# ============================================================ +# One-time side effect: auto-create ~/.libero/config.yaml so importing LIBERO +# does not trigger an interactive prompt. +# ============================================================ + + +def _ensure_libero_config() -> None: + import yaml as _yaml + + libero_config_path = os.environ.get( + "LIBERO_CONFIG_PATH", os.path.expanduser("~/.libero") + ) + config_file = os.path.join(libero_config_path, "config.yaml") + + if not os.path.exists(config_file): + os.makedirs(libero_config_path, exist_ok=True) + import libero.libero as _libero_pkg + + benchmark_root = os.path.dirname(os.path.abspath(_libero_pkg.__file__)) + default_paths = { + "benchmark_root": benchmark_root, + "bddl_files": os.path.join(benchmark_root, "./bddl_files"), + "init_states": os.path.join(benchmark_root, "./init_files"), + "datasets": os.path.join(benchmark_root, "../datasets"), + "assets": os.path.join(benchmark_root, "./assets"), + } + with open(config_file, "w") as f: + _yaml.dump(default_paths, f) + logger.info("Auto-created LIBERO config: %s", config_file) + + +_ensure_libero_config() + + +# ============================================================ +# task-suite entry point +# ============================================================ + + +def get_task_suite(task_suite_name: str): + """Load a LIBERO task suite.""" + from libero.libero import benchmark + + return benchmark.get_benchmark_dict()[task_suite_name]() + + +# ============================================================ +# actions +# ============================================================ + + +def get_libero_dummy_action() -> list[float]: + """Return the 7-dof dummy action used for episode warmup.""" + return [0, 0, 0, 0, 0, 0, -1] + + +# ============================================================ +# robosuite engine factory +# ============================================================ + + +def create_libero_engine( + task_id: int, + task_suite_name: str, + resolution: int = 256, + seed: int = 7, +) -> Any: + """Construct one LIBERO robosuite engine.""" + from libero.libero import get_libero_path + from libero.libero.envs import OffScreenRenderEnv + + task_suite = get_task_suite(task_suite_name) + task = task_suite.get_task(task_id) + task_bddl_file = os.path.join( + get_libero_path("bddl_files"), task.problem_folder, task.bddl_file + ) + env = OffScreenRenderEnv( + bddl_file_name=task_bddl_file, + camera_heights=resolution, + camera_widths=resolution, + ) + # The seed still affects object poses even when an initial state is fixed. + env.seed(seed) + env.env = VisualizationWrapper(env.env) + env.env.set_visualization_setting(setting="grippers", visible=False) + return env + + +# ============================================================ +# task metadata / initial states +# ============================================================ + + +def load_initial_states(initial_states_path: str) -> Optional[dict]: + """Load custom initial states, or return None for suite defaults.""" + if initial_states_path == "DEFAULT": + return None + with open(initial_states_path, "r") as f: + return json.load(f) + + +def resolve_task_info(task_suite, task_id: int) -> tuple[str, Any]: + """Return ``(task_desc, default_initial_states)`` for one task id.""" + num_tasks = task_suite.n_tasks + if task_id < 0 or task_id >= num_tasks: + raise ValueError(f"invalid task_id={task_id}, num_tasks={num_tasks}") + task = task_suite.get_task(task_id) + return task.language, task_suite.get_task_init_states(task_id) + + +def pick_initial_state( + initial_states_path: str, + custom_initial_states: Optional[dict], + task_desc: str, + default_states: Any, + episode_idx: int, +) -> np.ndarray: + """Pick one initial state from suite defaults or a custom states file.""" + if initial_states_path == "DEFAULT": + if default_states is None: + raise ValueError("default states missing for DEFAULT mode") + return default_states[episode_idx] + + if custom_initial_states is None: + raise ValueError(f"custom initial states not loaded for {initial_states_path}") + key = task_desc.replace(" ", "_") + ep_key = f"demo_{episode_idx}" + record = custom_initial_states[key][ep_key] + if not record["success"]: + raise ValueError(f"expert demo failed for {ep_key}") + return np.array(record["initial_state"]) + + +def get_instruction(task_desc: str) -> str: + """Return the instruction text for a LIBERO task description.""" + return task_desc + + +# ============================================================ +# render-skip: directly assign obs._enabled to avoid set_enabled() side effects. +# ============================================================ + + +def find_image_observables(env) -> list: + """Find image observables along the env.env wrapper chain.""" + cur = env + seen = set() + while cur is not None and id(cur) not in seen: + seen.add(id(cur)) + if hasattr(cur, "_observables") and isinstance(cur._observables, dict): + return [ + obs + for obs in cur._observables.values() + if getattr(obs, "modality", None) == "image" + ] + cur = getattr(cur, "env", None) + return [] + + +def set_render_enabled(image_obs_list, enabled: bool) -> None: + for obs in image_obs_list: + obs._enabled = enabled diff --git a/wall_x/_vendor/harrix/envs/registry.py b/wall_x/_vendor/harrix/envs/registry.py new file mode 100644 index 00000000..78d455c3 --- /dev/null +++ b/wall_x/_vendor/harrix/envs/registry.py @@ -0,0 +1,53 @@ +"""Environment registry and factory.""" + +from __future__ import annotations + +from typing import Type + +from wall_x._vendor.harrix.eval_config import EvalConfig +from wall_x._vendor.harrix.envs.base import BaseEnv + + +_REGISTRY: dict[str, Type[BaseEnv]] = {} + + +def register_env(name: str): + def deco(cls: Type[BaseEnv]): + if name in _REGISTRY: + raise ValueError( + f"env {name!r} already registered (cls={_REGISTRY[name].__name__})" + ) + _REGISTRY[name] = cls + return cls + + return deco + + +def _get_class(cfg: EvalConfig) -> Type[BaseEnv]: + t = cfg.env.type + cls = _REGISTRY.get(t) + if cls is None: + if t == "libero": + import wall_x._vendor.harrix.envs as _envs + + exc = getattr(_envs, "_LIBERO_IMPORT_ERROR", None) + if exc is not None: + raise RuntimeError( + "LIBERO evaluation dependencies are not installed. " + "Install LIBERO/robosuite and their simulator dependencies " + "before using env.type='libero'." + ) from exc + raise ValueError(f"unknown env type={t!r}, registered: {sorted(_REGISTRY)}") + return cls + + +def build_env(cfg: EvalConfig, worker_id: int) -> BaseEnv: + return _get_class(cfg)(cfg, worker_id) + + +def enumerate_episodes_for(cfg: EvalConfig) -> list[tuple]: + return _get_class(cfg).enumerate_episodes(cfg) + + +def registered_envs() -> list[str]: + return sorted(_REGISTRY) diff --git a/wall_x/_vendor/harrix/eval_config.py b/wall_x/_vendor/harrix/eval_config.py new file mode 100644 index 00000000..1ab59d61 --- /dev/null +++ b/wall_x/_vendor/harrix/eval_config.py @@ -0,0 +1,217 @@ +"""Typed YAML-driven configuration for inference and evaluation. + +The driver, model handle, and environment workers share one ``EvalConfig`` so +configuration is parsed once and then passed through explicitly. + +YAML schema: + model: + checkpoint_path: # checkpoint directory + train_config_path: null # null = /config.yml + norm_key: libero_all + cam_names: [face_view, right_wrist_view] + action_horizon: null # null = read data.action_horizon_flow + architecture: qwen2_5 # adapter registry key + action_mode: flow # flow / ar / dllm / vqa / subtask + env: + type: libero # env registry key + seed: 42 + libero: + task_suite_name: libero_spatial + initial_states_path: DEFAULT + num_trials_per_task: 50 + task_indices: null # null = all tasks + max_infer_times: 22 + skip_intermediate_render: true + force_render_task_indices: null # e.g. [5] always render that task + chunk_granular_render_toggle: false + rebuild_env_per_episode: false + rollout_dir: null # save third-person MP4 replays per episode + rollout_fps: 30 + runtime: + num_workers: 1 + max_batch_size: 1 + ws_port: 8765 + log_dir: /path/to/wallx_log + batch_sync_mode: false + debug: + deterministic_model: false + +Driver flow: + cfg = load_eval_config(yaml_path) + cfg = autofill_from_checkpoint(cfg) + +Unknown YAML fields are rejected to avoid silent typos. +""" + +import dataclasses +import os +from dataclasses import dataclass, field +from typing import Optional + +import yaml + + +@dataclass +class ModelSection: + checkpoint_path: str = "" + train_config_path: Optional[str] = None + norm_key: str = "libero_all" + cam_names: list = field(default_factory=lambda: ["face_view", "right_wrist_view"]) + action_horizon: Optional[int] = None + # Adapter implementation key in harrix.adapters.registry.ADAPTER_REGISTRY. + architecture: str = "qwen2_5" + # Inference algorithm. Each adapter validates its supported subset. + action_mode: str = "flow" + + +@dataclass +class LiberoEnvParams: + """LIBERO-specific env settings used when ``env.type == "libero"``.""" + + task_suite_name: str = "libero_spatial" + initial_states_path: str = "DEFAULT" + num_trials_per_task: int = 50 + task_indices: Optional[list] = None + max_infer_times: int = 22 + # Render-skip is enabled by default. Listed task ids always render every + # simulator step for contact-sensitive tasks. + skip_intermediate_render: bool = True + force_render_task_indices: Optional[list] = None + # If enabled, render observables are toggled only when the chunk boundary + # actually changes the desired state. + chunk_granular_render_toggle: bool = False + # Rebuild the simulator for every episode instead of only on task changes. + # This is mainly a debugging option because it changes the simulator RNG path. + rebuild_env_per_episode: bool = False + # When set, save a third-person MP4 replay for each episode under this directory. + # Falls back to the ``WALLX_ROLLOUT_DIR`` environment variable when null. + rollout_dir: Optional[str] = None + rollout_fps: int = 30 + + +@dataclass +class EnvSection: + # Env implementation key in harrix.envs.registry. + type: str = "libero" + seed: int = 42 + libero: LiberoEnvParams = field(default_factory=LiberoEnvParams) + + +@dataclass +class RuntimeSection: + num_workers: int = 1 + max_batch_size: int = 1 + ws_port: int = 8765 + log_dir: str = "/path/to/wallx_log" + # Run fixed task-local frames instead of dynamic work stealing. This improves + # reproducibility at the cost of possible idle workers inside a frame. + batch_sync_mode: bool = False + # Public Wall-X evaluation uses the in-process driver. Other drivers may be + # enabled by downstream/internal integrations. + driver_mode: str = "in_process" + + +@dataclass +class DebugSection: + # Enable deterministic torch backend options for variance debugging. + deterministic_model: bool = False + + +@dataclass +class EvalConfig: + model: ModelSection = field(default_factory=ModelSection) + env: EnvSection = field(default_factory=EnvSection) + runtime: RuntimeSection = field(default_factory=RuntimeSection) + debug: DebugSection = field(default_factory=DebugSection) + + +def _build_dataclass(cls, raw): + """Build a dataclass from a dict and reject unknown fields.""" + if raw is None: + return cls() + if not isinstance(raw, dict): + raise ValueError(f"Expected dict for {cls.__name__}, got {type(raw).__name__}") + field_names = {f.name for f in dataclasses.fields(cls)} + field_types = {f.name: f.type for f in dataclasses.fields(cls)} + unknown = set(raw.keys()) - field_names + if unknown: + raise ValueError( + f"Unknown field(s) {sorted(unknown)} in {cls.__name__}; " + f"expected one of {sorted(field_names)}" + ) + built = {} + for k, v in raw.items(): + ft = field_types[k] + if ( + isinstance(ft, type) + and dataclasses.is_dataclass(ft) + and isinstance(v, dict) + ): + built[k] = _build_dataclass(ft, v) + else: + built[k] = v + return cls(**built) + + +def load_eval_config(yaml_path: str) -> EvalConfig: + """Load ``EvalConfig`` from YAML and validate checkpoint_path.""" + with open(yaml_path, "r") as f: + raw = yaml.safe_load(f) + if raw is None: + raise ValueError(f"Empty YAML: {yaml_path}") + + cfg = _build_dataclass(EvalConfig, raw) + + if not cfg.model.checkpoint_path: + raise ValueError(f"model.checkpoint_path is required in {yaml_path}") + if not os.path.isdir(cfg.model.checkpoint_path): + raise FileNotFoundError( + f"model.checkpoint_path does not exist: {cfg.model.checkpoint_path}" + ) + + return cfg + + +def _load_train_config_yaml(path: str) -> dict: + """Load a checkpoint-side train YAML. + + Training checkpoints may contain PyYAML-specific tags such as + ``!!python/tuple``; ``safe_load`` cannot parse those. + """ + with open(path, "r") as f: + return yaml.load(f, Loader=yaml.FullLoader) or {} + + +def _read_action_horizon_flow(train_yml: dict) -> int: + task = train_yml.get("task") or {} + data = train_yml.get("data") or {} + return int( + train_yml.get("action_horizon_flow") + or task.get("action_horizon_flow") + or data.get("action_horizon_flow") + or 32 + ) + + +def autofill_from_checkpoint(cfg: EvalConfig) -> EvalConfig: + """Fill omitted model fields from the checkpoint-side train config. + + This resolves ``model.train_config_path`` and ``model.action_horizon``. + """ + if cfg.model.train_config_path is None: + for fname in ("config.yml", "config.yaml"): + cand = os.path.join(cfg.model.checkpoint_path, fname) + if os.path.exists(cand): + cfg.model.train_config_path = cand + break + else: + raise FileNotFoundError( + f"No config.yml/config.yaml in {cfg.model.checkpoint_path}; " + "set model.train_config_path explicitly in YAML" + ) + + if cfg.model.action_horizon is None: + train_yml = _load_train_config_yaml(cfg.model.train_config_path) + cfg.model.action_horizon = _read_action_horizon_flow(train_yml) + + return cfg diff --git a/wall_x/serving/__init__.py b/wall_x/_vendor/harrix/serving-old/__init__.py similarity index 100% rename from wall_x/serving/__init__.py rename to wall_x/_vendor/harrix/serving-old/__init__.py diff --git a/wall_x/_vendor/harrix/serving-old/_wallx_infer/__init__.py b/wall_x/_vendor/harrix/serving-old/_wallx_infer/__init__.py new file mode 100644 index 00000000..cb5d7d60 --- /dev/null +++ b/wall_x/_vendor/harrix/serving-old/_wallx_infer/__init__.py @@ -0,0 +1 @@ +"""Wall-X inference helpers vendored for harrix serving.""" diff --git a/wall_x/infer/base_dataclass.py b/wall_x/_vendor/harrix/serving-old/_wallx_infer/base_dataclass.py similarity index 63% rename from wall_x/infer/base_dataclass.py rename to wall_x/_vendor/harrix/serving-old/_wallx_infer/base_dataclass.py index b550c3b5..0834268e 100644 --- a/wall_x/infer/base_dataclass.py +++ b/wall_x/_vendor/harrix/serving-old/_wallx_infer/base_dataclass.py @@ -1,10 +1,10 @@ -from wall_x.infer.infer_config import InferConfig +from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig from typing import Optional, List from dataclasses import dataclass, field import numpy as np import torch -import wall_x.infer.data_utils as data_utils -from wall_x.infer.logger import InferLogger +import wall_x._vendor.x2robot_utils.geometry as data_utils +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger dof_dims = { "left_ee_cartesian_pos": 3, @@ -13,29 +13,57 @@ "left_ee_rotation_relative": 3, "left_ee_rotation_6D": 6, "left_ee_rotation_6D_relative": 6, - "left_arm_joint_pos": 7, + "left_arm_joint_pos": 7, # 6 joint + 1 gripper "left_gripper": 1, "left_gripper_cur": 1, - "left_arm_joint_cur": 1, + "left_arm_joint_cur": 7, # 1 -> 7 "right_ee_cartesian_pos": 3, "right_ee_cartesian_pos_relative": 3, "right_ee_rotation": 3, "right_ee_rotation_relative": 3, "right_ee_rotation_6D": 6, "right_ee_rotation_6D_relative": 6, - "right_arm_joint_pos": 7, + "right_arm_joint_pos": 7, # 6 joint + 1 gripper "right_gripper": 1, "right_gripper_cur": 1, - "right_arm_joint_cur": 1, + "right_arm_joint_cur": 7, # 1 -> 7 "head_actions": 2, "height": 1, "car_pose": 3, "velocity_decomposed": 3, + "velocity_decomposed_odom": 3, + "head_rotation": 2, # match ex001 + "left_joint": 6, + "left_joint_gripper": 1, + "right_joint": 6, + "right_joint_gripper": 1, + "left_rotation_quat": 4, + "right_rotation_quat": 4, + "left_quaternion": 4, + "right_quaternion": 4, + "left_wrench_ext_local_force": 3, + "left_wrench_ext_local_torque": 3, + "right_wrench_ext_local_force": 3, + "right_wrench_ext_local_torque": 3, + "left_wrench_ext_local_force_from_joint": 3, + "left_wrench_ext_local_torque_from_joint": 3, + "right_wrench_ext_local_force_from_joint": 3, + "right_wrench_ext_local_torque_from_joint": 3, + "left_wrench_ext_world_force": 3, + "left_wrench_ext_world_torque": 3, + "left_wrench_ext_world_force_from_joint": 3, + "left_wrench_ext_world_torque_from_joint": 3, + "right_wrench_ext_world_force": 3, + "right_wrench_ext_world_torque": 3, + "right_wrench_ext_world_force_from_joint": 3, + "right_wrench_ext_world_torque_from_joint": 3, + "left_arm_joint_dev": 7, + "right_arm_joint_dev": 7, } class ComputedDict(dict): - """Smart dictionary that supports registering computation rules and auto-computes None values on get""" + """Dict that registers compute rules and auto-fills None values on get.""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -43,45 +71,45 @@ def __init__(self, *args, **kwargs): def register_compute_rule(self, key, compute_func): """ - Register a computation rule. + Register a compute rule Args: - key: The key that needs computation - compute_func: Computation function that takes self as argument and returns the computed result + key: Key to compute + compute_func: Callable taking self, returns computed value """ self._compute_rules[key] = compute_func def get(self, key, default=None): - """Override get method to support auto-computation""" + """Override get to support lazy computation.""" value = super().get(key, default) - # If value is None and there's a compute rule, try to compute + # If value is None and a rule exists, try to compute if value is None and key in self._compute_rules: try: computed_value = self._compute_rules[key](self) if computed_value is not None: - # Cache the computed result + # Cache computed value self[key] = computed_value return computed_value except Exception: - pass # If computation fails, return None or default + pass # On failure return None or default return value if value is not None else default def __getitem__(self, key): - """Override [] operator to support auto-computation""" + """Override [] to support lazy computation.""" value = super().get(key, None) - # If value is None and there's a compute rule, try to compute + # If value is None and a rule exists, try to compute if value is None and key in self._compute_rules: try: computed_value = self._compute_rules[key](self) if computed_value is not None: - # Cache the computed result + # Cache computed value self[key] = computed_value return computed_value except Exception: - pass # If computation fails, raise original KeyError or return None + pass # On failure raise KeyError or return None if key in self: return super().__getitem__(key) @@ -94,7 +122,7 @@ class RobotStateActionData: data: ComputedDict = field( default_factory=lambda: ComputedDict( { - # State (formerly pose) - using state_ prefix + # State (formerly pose) - state_ prefix "state_left_ee_cartesian_pos": None, # (1, 3) "state_left_ee_rotation": None, # (1, 3) "state_left_ee_rotation_6D": None, @@ -110,10 +138,42 @@ class RobotStateActionData: "state_right_gripper_cur": None, "state_right_arm_joint_cur": None, # (1, 1) "state_head_actions": None, + "state_head_rotation": None, # match ex001 "state_height": None, "state_car_pose": None, "state_velocity_decomposed": None, - # Action - using action_ prefix + "state_velocity_decomposed_odom": None, + # support joint control + "state_left_joint": None, + "state_left_joint_gripper": None, + "state_right_joint": None, + "state_right_joint_gripper": None, + # support quaternion control + "state_left_quaternion": None, + "state_right_quaternion": None, + "state_left_rotation_quat": None, + "state_right_rotation_quat": None, + # support wrench/force-torque observation + "state_left_wrench_ext_local_force": None, + "state_left_wrench_ext_local_torque": None, + "state_right_wrench_ext_local_force": None, + "state_right_wrench_ext_local_torque": None, + "state_left_wrench_ext_local_force_from_joint": None, + "state_left_wrench_ext_local_torque_from_joint": None, + "state_right_wrench_ext_local_force_from_joint": None, + "state_right_wrench_ext_local_torque_from_joint": None, + "state_left_wrench_ext_world_force": None, + "state_left_wrench_ext_world_torque": None, + "state_left_wrench_ext_world_force_from_joint": None, + "state_left_wrench_ext_world_torque_from_joint": None, + "state_right_wrench_ext_world_force": None, + "state_right_wrench_ext_world_torque": None, + "state_right_wrench_ext_world_force_from_joint": None, + "state_right_wrench_ext_world_torque_from_joint": None, + # support joint deviation + "state_left_arm_joint_dev": None, + "state_right_arm_joint_dev": None, + # Action - action_ prefix "action_left_ee_cartesian_pos": None, "action_left_ee_cartesian_pos_relative": None, "action_left_ee_rotation": None, @@ -131,9 +191,21 @@ class RobotStateActionData: "action_right_gripper": None, "action_right_arm_joint_pos": None, "action_head_actions": None, + "action_head_rotation": None, # match ex001 "action_height": None, "action_car_pose": None, "action_velocity_decomposed": None, + "action_velocity_decomposed_odom": None, + # support joint control + "action_left_joint": None, + "action_left_joint_gripper": None, + "action_right_joint": None, + "action_right_joint_gripper": None, + # support quaternion control + "action_left_quaternion": None, + "action_right_quaternion": None, + "action_left_rotation_quat": None, + "action_right_rotation_quat": None, } ) ) @@ -141,8 +213,8 @@ class RobotStateActionData: logger = InferLogger.get_robot_logger("RobotStateActionData") def __post_init__(self): - """Register computation rules""" - # State computation rules - euler angles -> 6D rotation + """Register a compute rule""" + # State rules: euler angles -> 6D rotation self.data.register_compute_rule( "state_left_ee_rotation_6D", lambda d: ( @@ -160,7 +232,7 @@ def __post_init__(self): ), ) - # Action computation rules - absolute position computed from relative + state + # Action rules: absolute position from relative + state self.data.register_compute_rule( "action_left_ee_cartesian_pos", lambda d: ( @@ -182,7 +254,7 @@ def __post_init__(self): ), ) - # Action computation rules - get absolute rpy + # Action rules: derive absolute RPY self.data.register_compute_rule( # delta rpy -> abs rpy "action_left_ee_rotation", lambda d: ( @@ -203,10 +275,10 @@ def __post_init__(self): else None ), ) - self.data.register_compute_rule( # delta 6D -> abs 6D -> abs rpy + self.data.register_compute_rule( # delta 6D -> abs 6D "action_left_ee_rotation_6D", lambda d: ( - data_utils.compose_state_and_delta_to_abs_rpy( + data_utils.compose_state_and_delta_to_abs_6d( d["action_left_ee_rotation_6D_relative"], d["state_left_ee_rotation_6D"][0], ) @@ -236,10 +308,10 @@ def __post_init__(self): else None ), ) - self.data.register_compute_rule( # delta 6D -> abs 6D -> abs rpy + self.data.register_compute_rule( # delta 6D -> abs 6D "action_right_ee_rotation_6D", lambda d: ( - data_utils.compose_state_and_delta_to_abs_rpy( + data_utils.compose_state_and_delta_to_abs_6d( d["action_right_ee_rotation_6D_relative"], d["state_right_ee_rotation_6D"][0], ) @@ -251,11 +323,17 @@ def __post_init__(self): def get_agent_pos(self, obs_action_keys=None): if obs_action_keys is None: - obs_action_keys = self.config.train_config["data"]["obs_action_keys"] + obs_action_keys = self.config.train_config["agent_pos_config"].keys() agent_pose_data = [] for key in obs_action_keys: - # Remove follow_ or master_ prefix + # action_padding is a virtual key used to pad agent_pos width. + if key == "action_padding": + dim = self.config.train_config["agent_pos_config"][key] + agent_pose_data.append(np.zeros((1, dim))) + continue + + # Strip follow_ or master_ prefix if key.startswith("follow_"): key = key.replace("follow_", "") elif key.startswith("master_"): @@ -265,10 +343,10 @@ def get_agent_pos(self, obs_action_keys=None): state_key = f"state_{key}" if state_key in self.data: - # Use get method, which will auto-handle None value computation + # get() auto-computes when value is None value = self.data.get(state_key) if value is None: - # If still None after computation, use zero vector + # If still None after compute, use zero vector agent_pose_data.append(np.zeros((1, dof_dims[key]))) else: agent_pose_data.append(value) @@ -281,11 +359,17 @@ def get_agent_pos(self, obs_action_keys=None): def get_agent_pos_mask(self, obs_action_keys=None): if obs_action_keys is None: - obs_action_keys = self.config.train_config["data"]["obs_action_keys"] + obs_action_keys = self.config.train_config["agent_pos_config"].keys() agent_pos_mask_data = [] for key in obs_action_keys: - # Remove follow_ or master_ prefix + # action_padding carries no information and should stay masked out. + if key == "action_padding": + dim = self.config.train_config["agent_pos_config"][key] + agent_pos_mask_data.append(np.zeros((1, dim))) + continue + + # Strip follow_ or master_ prefix if key.startswith("follow_"): key = key.replace("follow_", "") elif key.startswith("master_"): @@ -295,19 +379,19 @@ def get_agent_pos_mask(self, obs_action_keys=None): state_key = f"state_{key}" if state_key in self.data: - # Use get method, which will auto-handle None value computation + # get() auto-computes when value is None value = self.data.get(state_key) if value is None: agent_pos_mask_data.append(np.zeros((1, dof_dims[key]))) else: - agent_pos_mask_data.append(np.ones((1, dof_dims[key]))) + agent_pos_mask_data.append(np.ones((1, value.shape[1]))) else: raise ValueError(f"Key {state_key} not found in data") return np.concatenate(agent_pos_mask_data, axis=1)[None] # (1, 1, D) - def save_state_data_with_key(self, value, key): - # Remove follow_ or master_ prefix + def save_state_data_with_key(self, value, key, gt_dim=None): + # Strip follow_ or master_ prefix key = key.replace("follow_", "") key = key.replace("master_", "") @@ -315,16 +399,17 @@ def save_state_data_with_key(self, value, key): if isinstance(value, torch.Tensor): value = value.detach().cpu().numpy() - if f"state_{key}" not in self.data: # TODO: joint angle control - self.logger.warning(f"{key} is not a valid state key, not recorded") + if f"state_{key}" not in self.data: # TODO： joint angle control + self.logger.warning(f"{key} is not a valid state key; not recorded") return + gt_dim = dof_dims[key] if gt_dim is None else gt_dim - # Shape validation for value, expected shape is (1, D) - if value.shape == (1, dof_dims[key]): + # Shape check: expected (1, D) + if value.shape == (1, gt_dim): self.data[f"state_{key}"] = value - elif value.shape == (1, 1, dof_dims[key]): + elif value.shape == (1, 1, gt_dim): self.data[f"state_{key}"] = value[0] - elif value.shape == (dof_dims[key],): + elif value.shape == (gt_dim,): self.data[f"state_{key}"] = value[None] else: raise ValueError(f"Value shape {value.shape} is not legal") @@ -355,6 +440,12 @@ def save_action_data( dof_start = 0 for action_key in predict_action_keys: + # action_padding is a virtual key that only advances dof_start; no data is written. + if action_key == "action_padding": + dof_dim = self.config.train_config["dof_config"]["action_padding"] + dof_start += dof_dim + continue + action_key = action_key.replace("follow_", "") action_key = action_key.replace("master_", "") dof_dim = dof_dims[action_key] @@ -362,7 +453,7 @@ def save_action_data( self.data[action_key] = predict_action[:, dof_start : dof_start + dof_dim] dof_start += dof_dim - # For compatibility, provide convenient property access + # Convenience properties for compatibility @property def agent_pos(self): return self.get_agent_pos() diff --git a/wall_x/_vendor/harrix/serving-old/_wallx_infer/infer_config.py b/wall_x/_vendor/harrix/serving-old/_wallx_infer/infer_config.py new file mode 100644 index 00000000..73672546 --- /dev/null +++ b/wall_x/_vendor/harrix/serving-old/_wallx_infer/infer_config.py @@ -0,0 +1,343 @@ +import logging +import os +from dataclasses import dataclass, field +from typing import Any + +import yaml +from qwen_vl_utils.vision_process import IMAGE_FACTOR, MAX_PIXELS, MIN_PIXELS + +logger = logging.getLogger(__name__) + + +@dataclass +class InferenceDataConfig: + """Minimal data config needed by online inference. + + Serving does not build datasets, so it should not require closed-source + data backends just to read image resize settings from a checkpoint yaml. + """ + + resolution: dict[str, int] = field(default_factory=dict) + model_type: str = "qwen2_5" + max_pixels: int = MAX_PIXELS + min_pixels: int = MIN_PIXELS + image_factor: int = IMAGE_FACTOR + use_relative_action: bool = False + + def __getitem__(self, key: str) -> Any: + if not hasattr(self, key): + raise KeyError(key) + return getattr(self, key) + + +class InferConfig: + def __init__( + self, + checkpoint_path: str | None = None, + train_config_path: str | None = None, + robot_host: str = "0.0.0.0", + robot_port: int = 41776, + robot_type: str = "desktop", # ["desktop", "turtle", "ex001"] + robot_action_start_ratio: float = 0, # Start ratio when trimming executed action + robot_action_end_ratio: float = 0.8, # End ratio when trimming executed action + robot_action_interpolate_multiplier: int = 10, # Action interpolation multiplier + robot_use_joint_angle_control: bool = False, # Joint control (model must predict joints) + turtle_as_desktop: bool = False, # Turtle as desktop: fixed base/head cam/height + action_horizon: int = 32, # Must match model action horizon + action_dim: int | None = None, + ar_action_dim: int | None = None, + model_device: str = "cuda:0", + num_inference_timesteps: int = 10, + norm_key: str = "x2_normal", # ["x2_normal", "ex_normal"] + cam_names: list[str] = ["face_view", "left_wrist_view", "right_wrist_view"], + save_video_dir: str = "./videos", + robot_id: str = "10000", + model_type: str = "wallx", # ["wallx", "vga"] + ): + # Private path storage + assert checkpoint_path is not None + self._checkpoint_path = checkpoint_path + if os.path.exists(os.path.join(checkpoint_path, "normalizer_action.pth")): + self.normalizer_action_path = os.path.join( + checkpoint_path, "normalizer_action.pth" + ) + if os.path.exists(os.path.join(checkpoint_path, "normalizer_propri.pth")): + self.normalizer_propri_path = os.path.join( + checkpoint_path, "normalizer_propri.pth" + ) + + # Other config fields + self.robot_host = robot_host + self.robot_port = robot_port + self.robot_type = robot_type # ["desktop", "turtle", "ex001"] + self.robot_action_start_ratio = robot_action_start_ratio + self.robot_action_end_ratio = robot_action_end_ratio + self.robot_action_interpolate_multiplier = robot_action_interpolate_multiplier + self.robot_use_joint_angle_control = ( + robot_use_joint_angle_control # Joint-angle control + ) + self.turtle_as_desktop = turtle_as_desktop + self.robot_id = robot_id + self._action_horizon = ( + action_horizon # Default from train config flow action horizon + ) + self._action_dim = action_dim # Default from train config dof_config + self.ar_action_dim = ar_action_dim # Default from train config ar_dof_config + self.model_device = model_device + self.num_inference_timesteps = ( + num_inference_timesteps # flow matching related config + ) + self.save_video_dir = save_video_dir + self.model_type = model_type # ["wallx", "vga"] + # Config objects + self.train_config: dict = {} + self.model_config = None + self.data_config = None + self.norm_key = norm_key + self.cam_names = cam_names + # Load all configs + self._load_all_configs(train_config_path) + + @property + def checkpoint_path(self) -> str | None: + return self._checkpoint_path + + @checkpoint_path.setter + def checkpoint_path(self, value: str | None): + """Reload all configs when checkpoint_path changes.""" + if self._checkpoint_path != value: + self._checkpoint_path = value + self._load_all_configs() + + @property + def action_horizon(self) -> int: + return self._action_horizon + + @action_horizon.setter + def action_horizon(self, value: int): + self._action_horizon = value + + @property + def action_dim(self) -> int | None: + return self._action_dim + + @action_dim.setter + def action_dim(self, value: int | None): + self._action_dim = value + + @property + def ar_action_dim(self) -> int | None: + return self._ar_action_dim + + @ar_action_dim.setter + def ar_action_dim(self, value: int | None): + self._ar_action_dim = value + + def _load_all_configs(self, train_config_path=None): + """Load all configs.""" + self._load_train_config(train_config_path) + if self.model_type != "vga": + self._load_model_config_for_wallx() + self._load_data_config() + + # Update action_horizon/action_dim if needed + if self._action_horizon is None: + self._action_horizon = self.train_config.get("data", {}).get( + "action_horizon_flow", 32 + ) + assert self._action_horizon is not None and self._action_horizon > 0 + + if self._action_dim is None: + dof_config = self.train_config.get("dof_config") or self.train_config.get( + "data", {} + ).get("dof_config", {}) + self._action_dim = sum(dof_config.values()) + + if self._ar_action_dim is None: + self._ar_action_dim = sum( + self.train_config.get("ar_dof_config", {}).values() + ) + + def _load_train_config(self, train_config_path): + if train_config_path is None: + for fname in ("config.yml", "config.yaml"): + candidate = os.path.join(self._checkpoint_path, fname) + if os.path.exists(candidate): + train_config_path = candidate + break + else: + raise FileNotFoundError( + f"No config.yml/config.yaml found in {self._checkpoint_path}" + ) + with open(train_config_path, "r") as f: + self.train_config = yaml.load(f, Loader=yaml.FullLoader) + + ckpt_dir = self._checkpoint_path + preprocessor_file = os.path.join(ckpt_dir, "preprocessor_config.json") + if os.path.exists(preprocessor_file): + logger.info( + "[LoadConfig] Found %s, override processor_path.", + preprocessor_file, + ) + self.train_config["processor_path"] = ckpt_dir + + tokenizer_file = os.path.join(ckpt_dir, "tokenizer.json") + tokenizer_config_file = os.path.join(ckpt_dir, "tokenizer_config.json") + if self.train_config.get( + "action_tokenizer_path", None + ) is not None and not os.path.exists( + self.train_config.get("action_tokenizer_path", None) + ): + if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file): + logger.info( + "[LoadConfig] Found tokenizer files in %s, override action_tokenizer_path.", + ckpt_dir, + ) + self.train_config["action_tokenizer_path"] = ckpt_dir + else: + logger.warning("[LoadConfig] Cannot load action tokenizer! ") + + def _load_model_config_for_wallx(self): + model_type = self.train_config["model_type"] + if not os.path.isdir(self._checkpoint_path): + return + # For Qwen models + ckpt_config_path = os.path.join(self._checkpoint_path, "config.json") + resolved_cfg_path = None + + if os.path.exists(ckpt_config_path): + # Prefer checkpoint config + resolved_cfg_path = ckpt_config_path + logger.info( + "[LoadModelConfig] Using checkpoint config.json: %s", + ckpt_config_path, + ) + else: + # Fallback to original config path + fallback_cfg = self.train_config.get("qwen_vl_act_config_path", None) + if fallback_cfg is not None: + resolved_cfg_path = fallback_cfg + logger.info( + "[LoadModelConfig] Using fallback act config: %s", + fallback_cfg, + ) + + if resolved_cfg_path is None or (not os.path.exists(resolved_cfg_path)): + raise ValueError( + f"[LoadModelConfig] Cannot load model config! " + f"Checked:\n" + f" - Checkpoint config.json: {ckpt_config_path}\n" + f" - Fallback path: {self.train_config.get('qwen_vl_act_config_path', None)}" + ) + + # Save back to config for consistency + self.train_config["qwen_vl_act_config_path"] = resolved_cfg_path + + from wall_x.trainer.adapters import ADAPTER_REGISTRY, format_adapter_error + + adapter_cls = ADAPTER_REGISTRY.get(model_type) + if adapter_cls is None: + raise ValueError(f"[LoadModelConfig] {format_adapter_error(model_type)}") + ConfigClass = adapter_cls.config_class() + + logger.info( + "[LoadModelConfig] Loading model config from: %s", resolved_cfg_path + ) + if resolved_cfg_path.endswith(".json"): + self.model_config = ConfigClass.from_json_file(resolved_cfg_path) + else: + self.model_config = ConfigClass.from_pretrained(resolved_cfg_path) + + self.model_config.update_model_config(self.train_config) + + self.model_config._attn_implementation = "sdpa" + self.model_config.vision_config._attn_implementation = "flash_attention_2" + + logger.info("[LoadModelConfig] Model config loaded and updated successfully.") + + def _load_data_config(self): + # Prefer typed TrainConfig path (handles new 8-section schema where + # dof_config lives under task:). Fall back to legacy raw-dict path + # for old flat yamls. See trainer/adapters/base_adapter.py. + typed_dcfg = self._try_typed_data_config() + if typed_dcfg is not None: + self.data_config = typed_dcfg + return + + if self.model_type != "vga": + self.data_config = self._build_inference_data_config() + return + + # TEMPORARY: direct call to the private _set_data_backend. + # VGA inference still uses the legacy x2robot data config object. + # Wall-X online inference uses the lightweight InferenceDataConfig + # above so serving does not require excluded data backends. + from wall_x.data._registry import _set_data_backend + + _set_data_backend(self.train_config.get("dataset_type", "x2robot_v1")) + from wall_x.model.vga.openloop_visualization import get_data_configs + + dataload_config = get_data_configs(self.train_config.get("data", {})) + dataload_config["predict_action_keys"] = list( + dataload_config.get("dof_config", {}).keys() + ) + self.data_config = dataload_config + + def _build_inference_data_config(self) -> InferenceDataConfig: + data = self.train_config.get("data", {}) or {} + + def get(key: str, default: Any) -> Any: + return data.get(key, self.train_config.get(key, default)) + + resolution = get("resolution", {}) or {} + return InferenceDataConfig( + resolution=dict(resolution), + model_type=get( + "model_type", self.train_config.get("model_type", "qwen2_5") + ), + max_pixels=get("max_pixels", MAX_PIXELS), + min_pixels=get("min_pixels", MIN_PIXELS), + image_factor=get("image_factor", IMAGE_FACTOR), + use_relative_action=get("use_relative_action", False), + ) + + def _try_typed_data_config(self): + """Attempt to load TrainConfig and build X2RDataConfig via typed path. + + Returns the X2RDataConfig on success, or None if the yaml is not in + TrainConfig schema (legacy flat yaml). + """ + yml_path = None + for fname in ("config.yml", "config.yaml"): + candidate = os.path.join(self._checkpoint_path, fname) + if os.path.exists(candidate): + yml_path = candidate + break + if yml_path is None: + return None + try: + from wall_x.config.loader import load_config + from wall_x.trainer.adapters.base_adapter import load_x2r_data_config + + typed_cfg = load_config(yml_path) + except Exception as e: + logger.warning( + "[InferConfig] Typed config load failed for %s: %s", + yml_path, + e, + ) + return None + try: + return load_x2r_data_config(typed_cfg) + except Exception as e: + logger.warning( + "[InferConfig] Typed config load failed for %s: %s", + yml_path, + e, + ) + return None + + +if __name__ == "__main__": + config = InferConfig() + logger.info("%s", config.train_config) diff --git a/wall_x/infer/logger.py b/wall_x/_vendor/harrix/serving-old/_wallx_infer/logger.py similarity index 71% rename from wall_x/infer/logger.py rename to wall_x/_vendor/harrix/serving-old/_wallx_infer/logger.py index ec8bf3fa..11ef3e82 100644 --- a/wall_x/infer/logger.py +++ b/wall_x/_vendor/harrix/serving-old/_wallx_infer/logger.py @@ -1,28 +1,4 @@ -""" -Hierarchical Inference Logging System - -Level structure: -- ENV: Environment layer (RealRobotEnv) -- ROBOT: Robot layer (Robot subclasses) -- CONTROLLER: Controller layer (RobotController, RobotCommunication) -- MODEL: Model layer (WallxModelWrapper) -- UTILS: Utility layer (various utility classes) - -Usage examples: - # Method 1: Auto-detect level - from wall_x.infer.logger import get_logger - logger = get_logger(__name__) - logger.info("This is an info message") - - # Method 2: Manually specify level - logger = get_logger(__name__, "ROBOT") - logger.debug("Robot state updated") - - # Method 3: Use shortcut methods - from wall_x.infer.logger import InferLogger - logger = InferLogger.get_robot_logger("DesktopRobot") - logger.warning("Action out of bounds") -""" +"""Layered logging helpers for inference components.""" import logging import sys @@ -36,25 +12,27 @@ HAS_COLORLOG = True except ImportError: HAS_COLORLOG = False - print("[WARNING] colorlog not installed. Install with: pip install colorlog") + logging.getLogger(__name__).warning( + "colorlog not installed. Install with: pip install colorlog" + ) class InferLogger: """ - Hierarchical inference logging system + Layered inference logging """ _loggers = {} _initialized = False - # Level definitions + # Level identifiers LEVEL_ENV = "ENV" LEVEL_ROBOT = "ROBOT" LEVEL_CONTROLLER = "CONTROLLER" LEVEL_MODEL = "MODEL" LEVEL_UTILS = "UTILS" - # Level color mapping (for terminal output) + # Level colors for console LEVEL_COLORS = { LEVEL_ENV: "cyan", LEVEL_ROBOT: "green", @@ -73,14 +51,14 @@ def setup( colorful: bool = True, ): """ - Initialize the logging system + Initialize logging Args: - log_level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL) + log_level: DEBUG, INFO, WARNING, ERROR, CRITICAL log_dir: Log file directory - console_output: Whether to output to console - file_output: Whether to output to file - colorful: Whether to use colorful output (requires colorlog) + console_output: Enable console output + file_output: Enable file output + colorful: Colored console (requires colorlog) """ if cls._initialized: return @@ -104,11 +82,11 @@ def setup( @classmethod def get_logger(cls, name: str, level: str = None) -> logging.Logger: """ - Get logger for specified level + Get logger for a level Args: - name: Logger name (usually module name or class name) - level: Level identifier (ENV, ROBOT, CONTROLLER, MODEL, UTILS) + name: Logger name (module or class) + level: Level tag (ENV, ROBOT, CONTROLLER, MODEL, UTILS) Returns: Configured logger instance @@ -116,7 +94,7 @@ def get_logger(cls, name: str, level: str = None) -> logging.Logger: if not cls._initialized: cls.setup() - # Auto-detect level + # Auto-detect level from name if level is None: level = cls._detect_level(name) @@ -133,13 +111,13 @@ def get_logger(cls, name: str, level: str = None) -> logging.Logger: # Clear existing handlers logger.handlers.clear() - # Console output + # Console handler if cls.console_output: console_handler = logging.StreamHandler(sys.stdout) console_handler.setLevel(cls.log_level) if cls.colorful: - # Colorful formatting + # Colored formatter color = cls.LEVEL_COLORS.get(level, "white") console_format = ( f"%(log_color)s[%(asctime)s]%(reset)s " @@ -161,7 +139,7 @@ def get_logger(cls, name: str, level: str = None) -> logging.Logger: }, ) else: - # Plain formatting + # Plain formatter console_format = ( f"[%(asctime)s] [{level:^10}] [%(name)s] " f"%(levelname)-8s %(message)s" @@ -173,7 +151,7 @@ def get_logger(cls, name: str, level: str = None) -> logging.Logger: console_handler.setFormatter(console_formatter) logger.addHandler(console_handler) - # File output + # File handler if cls.file_output and cls.log_file: file_handler = logging.FileHandler(cls.log_file, encoding="utf-8") file_handler.setLevel(cls.log_level) @@ -191,7 +169,7 @@ def get_logger(cls, name: str, level: str = None) -> logging.Logger: @classmethod def _detect_level(cls, name: str) -> str: - """Auto-detect level based on name""" + """Infer level from logger name.""" name_lower = name.lower() if "env" in name_lower: @@ -211,32 +189,32 @@ def _detect_level(cls, name: str) -> str: @classmethod def get_env_logger(cls, name: str = "Environment") -> logging.Logger: - """Get environment layer logger""" + """Environment-level logger.""" return cls.get_logger(name, cls.LEVEL_ENV) @classmethod def get_robot_logger(cls, name: str = "Robot") -> logging.Logger: - """Get robot layer logger""" + """Robot-level logger.""" return cls.get_logger(name, cls.LEVEL_ROBOT) @classmethod def get_controller_logger(cls, name: str = "Controller") -> logging.Logger: - """Get controller layer logger""" + """Controller-level logger.""" return cls.get_logger(name, cls.LEVEL_CONTROLLER) @classmethod def get_model_logger(cls, name: str = "Model") -> logging.Logger: - """Get model layer logger""" + """Model-level logger.""" return cls.get_logger(name, cls.LEVEL_MODEL) @classmethod def get_utils_logger(cls, name: str = "Utils") -> logging.Logger: - """Get utility layer logger""" + """Utils-level logger.""" return cls.get_logger(name, cls.LEVEL_UTILS) @classmethod def set_level(cls, level: str): - """Dynamically modify log level for all loggers""" + """Change log level for all loggers.""" new_level = getattr(logging, level.upper()) cls.log_level = new_level for logger in cls._loggers.values(): @@ -246,7 +224,7 @@ def set_level(cls, level: str): @classmethod def close_all(cls): - """Close file handles for all loggers""" + """Close all logger file handles.""" for logger in cls._loggers.values(): for handler in logger.handlers[:]: handler.close() @@ -255,22 +233,22 @@ def close_all(cls): cls._initialized = False -# Convenience functions +# Convenience helpers def get_logger(name: str, level: str = None) -> logging.Logger: """ - Convenience function to get logger + Convenience wrapper to get a logger Args: - name: Logger name (usually use __name__) - level: Level identifier (optional, will auto-detect) + name: Logger name (usually __name__) + level: Optional level (auto-detected if omitted) Returns: Configured logger instance - Usage examples: - from wall_x.infer.logger import get_logger - logger = get_logger(__name__) # Auto-detect level - logger = get_logger(__name__, "ROBOT") # Manually specify level + Example: + from wall_x._vendor.harrix.serving._wallx_infer.logger import get_logger + logger = get_logger(__name__) # Auto-detect level from name + logger = get_logger(__name__, "ROBOT") # explicit level """ return InferLogger.get_logger(name, level) @@ -283,17 +261,17 @@ def setup_logger( colorful: bool = True, ): """ - Convenience function to setup logging system + Convenience wrapper to configure logging Args: - log_level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL) + log_level: DEBUG, INFO, WARNING, ERROR, CRITICAL log_dir: Log file directory - console_output: Whether to output to console - file_output: Whether to output to file - colorful: Whether to use colorful output + console_output: Enable console output + file_output: Enable file output + colorful: Enable colored console output - Usage examples: - from wall_x.infer.logger import setup_logger + Example: + from wall_x._vendor.harrix.serving._wallx_infer.logger import setup_logger setup_logger(log_level="DEBUG", log_dir="./logs") """ InferLogger.setup(log_level, log_dir, console_output, file_output, colorful) diff --git a/wall_x/_vendor/harrix/serving-old/_wallx_infer/model_wrapper.py b/wall_x/_vendor/harrix/serving-old/_wallx_infer/model_wrapper.py new file mode 100644 index 00000000..885366c2 --- /dev/null +++ b/wall_x/_vendor/harrix/serving-old/_wallx_infer/model_wrapper.py @@ -0,0 +1,827 @@ +import os +import torch +import json +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.model.core.action.normalizer import Normalizer +from wall_x.utils.constant import action_statistic_dof as default_action_statistic_dof +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger +from wall_x.utils.timers import timer, ScopeTimer +from wall_x.data.backends.lerobot.utils import NormStats + + +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 + + +def _load_norm_stats(norm_stats_path, action_key): + with open(norm_stats_path, "r") as f: + norm_stats = json.load(f) + + q01 = torch.tensor(norm_stats["norm_stats"][action_key]["q01"]) + q99 = torch.tensor(norm_stats["norm_stats"][action_key]["q99"]) + delta = q99 - q01 + + return NormStats(min=q01, max=q99, delta=delta) + + +class WallxModelWrapper: + def __init__(self, config: InferConfig): + self.config = 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() + self.logger.info(f"Model {self.config.checkpoint_path} loaded") + + self.norm_key = self.config.norm_key + self._register_normalizers() + self.logger.info(f"Normalizers {self.norm_key} registered") + + # 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 _resolve_norm_key(self): + available_action_keys = set(self.normalizer_action.min.keys()) + available_propri_keys = set(self.normalizer_propri.min.keys()) + available_keys = sorted(available_action_keys & available_propri_keys) + + if self.norm_key in available_keys: + return + + if not available_keys: + self.logger.warning( + "[NormDebug] no available norm keys found in normalizer" + ) + return + + fallback_key = None + prefix_matches = [ + k for k in available_keys if k.startswith(f"{self.norm_key}_") + ] + if len(prefix_matches) == 1: + fallback_key = prefix_matches[0] + elif len(available_keys) == 1: + fallback_key = available_keys[0] + + if fallback_key is not None: + self.logger.warning( + "[NormDebug] norm_key=%s not found, falling back to %s. available_keys=%s", + self.norm_key, + fallback_key, + available_keys, + ) + self.norm_key = fallback_key + self.config.norm_key = fallback_key + return + + self.logger.warning( + "[NormDebug] norm_key=%s not found in normalizer. available_keys=%s", + self.norm_key, + available_keys, + ) + + 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 ADAPTER_REGISTRY, format_adapter_error + + model_type = self.config.train_config["model_type"] + adapter_cls = ADAPTER_REGISTRY.get(model_type) + if adapter_cls is None: + raise ValueError(format_adapter_error(model_type)) + ModelClass = adapter_cls.inference_model_class() + self.ModelClass = ModelClass + + self.logger.info(f"Initializing model: {model_type}") + + self.model = ModelClass( + self.config.model_config, + self.processor, + self.tokenizer_mixin, + ) + + # log attention implementation — variant-specific layout dispatched + # via the adapter (qwen2_5 uses model.model._attn_implementation; + # qwen3 / qwen3_5 use model.model.language_model.config._attn_impl). + 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 complete") + self.logger.info("Casting selected params to bfloat16") + self.model.to_bfloat16_for_selected_params() + self.logger.info("bfloat16 cast complete") + + 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 if present (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"❌ 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.", + ) + + 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): + missing = object() + action_statistic_dof = missing + + def load_action_statistic_dof(): + nonlocal action_statistic_dof + if action_statistic_dof is not missing: + return action_statistic_dof + if self.config.train_config.get("customized_action_statistic_dof", None): + with open( + self.config.train_config["customized_action_statistic_dof"], "r" + ) as f: + action_statistic_dof = json.load(f) + else: + action_statistic_dof = default_action_statistic_dof + return action_statistic_dof + + # Prefer lerobot norm_stats.json from checkpoint when present + checkpoint_path = self.config.checkpoint_path + norm_key = self.config.norm_key + norm_stats_path = os.path.join(checkpoint_path, "norm_stats.json") + if os.path.exists(norm_stats_path): + self.logger.info( + f"Loading lerobot norm stats from checkpoint {checkpoint_path}: {norm_stats_path}" + ) + propri_stats = _load_norm_stats(norm_stats_path, "observation.state") + self.normalizer_propri = Normalizer.from_lerobot_norm_stats( + propri_stats, norm_key + ) + action_stats = _load_norm_stats(norm_stats_path, "action") + self.normalizer_action = Normalizer.from_lerobot_norm_stats( + action_stats, norm_key + ) + self._resolve_norm_key() + self._log_norm_debug(self.norm_key) + return + + # Otherwise load normalizers from checkpoint files when available + if hasattr(self.config, "normalizer_action_path"): + self.normalizer_action = Normalizer.from_ckpt( + self.config.normalizer_action_path + ) + else: + action_statistic_dof = load_action_statistic_dof() + self.normalizer_action = Normalizer( + action_statistic_dof, + self.config.train_config["dof_config"], + min_key=self.config.train_config.get("min_key", "min"), + delta_key=self.config.train_config.get("delta_key", "delta"), + ) + if hasattr(self.config, "normalizer_propri_path"): + self.normalizer_propri = Normalizer.from_ckpt( + self.config.normalizer_propri_path + ) + else: + action_statistic_dof = load_action_statistic_dof() + self.normalizer_propri = Normalizer( + action_statistic_dof, + self.config.train_config["agent_pos_config"], + min_key=self.config.train_config.get("min_key", "min"), + delta_key=self.config.train_config.get("delta_key", "delta"), + ) + self._resolve_norm_key() + 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, + truncation=True, + return_tensors="pt", + max_length=1000, + pad_to_128_multiple=False, + pad_prefix_to_same_length=pad_prefix, + norm_state=( + additional_inputs["proprioception"] + if "proprioception" in additional_inputs + and self.config.train_config.get( + "use_state_string_representation", False + ) + 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=self.config.train_config["data"].get("state_bins", 256), + 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:" + for cam_name in self.cam_names: + user_request += f" {cam_name}: {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:" + for cam_name in self.cam_names: + user_request += f" {cam_name}: {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:" + for cam_name in self.cam_names: + user_request += f" {cam_name}: {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 ["face_view", "left_wrist_view", "right_wrist_view"]: + 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): + self.logger.info("Starting flow action generation") + 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) + + 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, + **model_input, + ) + + self.logger.info("Flow action generation complete") + + 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) + + # Batch inputs in one preprocesser_call + 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("Starting AR action generation") + 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 complete") + + 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("Starting subtask generation") + 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 complete: {subtask}") + return subtask + + def infer_vqa(self, observation, instruction): + self.logger.info("Starting VQA generation") + + 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 complete: {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: # Retry up to 3 times + 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, # v3.1 delta decode + **model_input, + ) + if model_output["predict_action"] is not None: + break + cnt += 1 + self.logger.warning(f"DLLM action failed; retry {cnt}") + self.logger.info("DLLM action generation complete") + + model_output["robot_state_action_data"] = observation["robot_state_action_data"] + model_output["robot_state_action_data"].save_action_data( + model_output["predict_action"] + ) + self.logger.info("Saved DLLM action to robot_state_action_data") + return model_output diff --git a/wall_x/_vendor/harrix/serving-old/_wallx_infer/robot.py b/wall_x/_vendor/harrix/serving-old/_wallx_infer/robot.py new file mode 100644 index 00000000..19dc8f15 --- /dev/null +++ b/wall_x/_vendor/harrix/serving-old/_wallx_infer/robot.py @@ -0,0 +1,964 @@ +import numpy as np +from abc import ABC +import re + +from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig +from wall_x._vendor.harrix.serving._wallx_infer.utils import VehiclePoseHandler, UnifiedTrajectoryProcessor +from wall_x._vendor.harrix.serving._wallx_infer.base_dataclass import RobotStateActionData +from wall_x._vendor.harrix.serving._wallx_infer.socket_controller import RobotController + +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger + +robot_action_key_mapping = { + "follow_left_ee_cartesian_pos": "follow1_pos[:3]", + "follow_left_ee_rotation": "follow1_pos[3:6]", + "follow_left_gripper": "follow1_pos[6:7]", + "follow_right_ee_cartesian_pos": "follow2_pos[:3]", + "follow_right_ee_rotation": "follow2_pos[3:6]", + "follow_right_gripper": "follow2_pos[6:7]", + "head_actions": "head_pos", + "head_rotation": "head_pos", # match ex001 + "height": "lift", + "velocity_decomposed": "velocity_decomposed", + "velocity_decomposed_odom": "velocity_decomposed_odom", # match ex001 + "follow_left_arm_joint_cur": "follow1_joints_cur", # 1 -> 7 + "follow_right_arm_joint_cur": "follow2_joints_cur", # 1 -> 7 + "follow_left_gripper_cur": "follow1_joints_cur[-1:]", + "follow_right_gripper_cur": "follow2_joints_cur[-1:]", + "follow_left_arm_joint_pos": "follow1_joints", + "follow_right_arm_joint_pos": "follow2_joints", + "follow_left_wrench_ext_local_force": "follow1_end_effort_force", + "follow_left_wrench_ext_local_torque": "follow1_end_effort_torque", + "follow_right_wrench_ext_local_force": "follow2_end_effort_force", + "follow_right_wrench_ext_local_torque": "follow2_end_effort_torque", + "follow_left_wrench_ext_local_force_from_joint": "follow1_end_effort_force_from_joint", + "follow_left_wrench_ext_local_torque_from_joint": "follow1_end_effort_torque_from_joint", + "follow_right_wrench_ext_local_force_from_joint": "follow2_end_effort_force_from_joint", + "follow_right_wrench_ext_local_torque_from_joint": "follow2_end_effort_torque_from_joint", + "follow_left_wrench_ext_world_force": "follow1_wrench_ext_world_force", + "follow_left_wrench_ext_world_torque": "follow1_wrench_ext_world_torque", + "follow_left_wrench_ext_world_force_from_joint": "follow1_wrench_ext_world_force_from_joint", + "follow_left_wrench_ext_world_torque_from_joint": "follow1_wrench_ext_world_torque_from_joint", + "follow_right_wrench_ext_world_force": "follow2_wrench_ext_world_force", + "follow_right_wrench_ext_world_torque": "follow2_wrench_ext_world_torque", + "follow_right_wrench_ext_world_force_from_joint": "follow2_wrench_ext_world_force_from_joint", + "follow_right_wrench_ext_world_torque_from_joint": "follow2_wrench_ext_world_torque_from_joint", + "follow_left_arm_joint_dev": "follow1_joints_dev", + "follow_right_arm_joint_dev": "follow2_joints_dev", +} + + +class Robot(ABC): + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.config = config + + self.robot_controller = RobotController( + robot_id=robot_id, host=config.robot_host, port=config.robot_port + ) + + self.robot_controller.connect() + + self.is_received = False # Blocking flag + self.current_robot_state_action = None + + self.logger = InferLogger.get_robot_logger("Robot") + + def _get_views_and_state(self): + state = self.robot_controller.recv_action() + views = self.robot_controller.recv_image( + ["camera_left", "camera_front", "camera_right"] + ) + self.is_received = True + return state, views + + def _get_dof_mask(self): + dof_config = self.config.train_config["dof_config"] + total_dof = sum(dof_config.values()) + dof_mask = np.ones((1, self.config.action_horizon, total_dof)) + return dof_mask + + @staticmethod + def _try_compose_arm_joint_pos_with_gripper(state: dict, key: str): + """Build 7D arm joint position from runtime state when possible. + + Why this is needed: + - Joint-control training expects `*_arm_joint_pos` to be 7D + (6 arm joints + 1 gripper). + - Some online robot states expose arm joints as 6D (`follow*_joints`) + and gripper separately in `follow*_pos[6]`. + - If we pass raw 6D joints directly, `save_state_data_with_key` will fail + shape validation against dof_dims=7. + + Behavior: + - Only handles `follow_left_arm_joint_pos` / `follow_right_arm_joint_pos`. + - Returns a composed 7D vector when both required sources exist and shapes + match expected runtime format. + - Returns None for all other keys or incompatible state shapes, so caller + falls back to original mapping logic. + """ + + if key == "follow_left_arm_joint_pos": + joints_key = "follow1_joints" + ee_key = "follow1_pos" + elif key == "follow_right_arm_joint_pos": + joints_key = "follow2_joints" + ee_key = "follow2_pos" + else: + return None + + if joints_key not in state or ee_key not in state: + return None + + joints = np.asarray(state[joints_key]).reshape(-1) + ee = np.asarray(state[ee_key]).reshape(-1) + if joints.shape[0] == 6 and ee.shape[0] >= 7: + return np.concatenate([joints, ee[6:7]], axis=0) + + return None + + def get_observation(self): + state, views = self._get_views_and_state() + + robot_state_action_data = RobotStateActionData(config=self.config) + for key in robot_action_key_mapping.keys(): + state_key_str = robot_action_key_mapping[key] + value = None + # Prefer a robust 7D composition path for arm_joint_pos: + # online state often provides `follow*_joints` as 6D plus gripper + # in `follow*_pos[6]`. Compose first to match model/dof schema. + composed_value = self._try_compose_arm_joint_pos_with_gripper(state, key) + if composed_value is not None: + value = composed_value + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + continue + + match = re.match(r"(\w+)\[(.*)\]", state_key_str) + if match: + base_key = match.group(1) + if base_key in state: + slicing_str = match.group(2) + slice_parts = slicing_str.split(":") + slice_args = [(int(p) if p.strip() else None) for p in slice_parts] + s = slice(*slice_args) + value = state[base_key][s] + else: + if state_key_str in state: + value = state[state_key_str] + if value is not None: + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + + dof_mask = self._get_dof_mask() + robot_state_action_data.dof_mask = dof_mask + self.current_robot_state_action = robot_state_action_data + return { + "robot_state_action_data": robot_state_action_data, + "face_view": views["camera_front"][0], # (1, 480, 640, 3) + "left_wrist_view": views["camera_left"][0], # (1, 480, 640, 3) + "right_wrist_view": views["camera_right"][0], # (1, 480, 640, 3) + } + + def go_home(self): + if self.current_robot_state_action is None: + self.get_observation() + + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "left_gripper" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "right_gripper" + ) + + action_dict = {"robot_state_action_data": self.current_robot_state_action} + self.apply_action( + action_dict, + robot_action_interpolate_multiplier=150, + robot_action_start_ratio=0, + robot_action_end_ratio=1, + ) + self.logger.info("Robot returned to home position") + + def _get_left_arm_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + if not self.config.robot_use_joint_angle_control: + left_ee_cartesian_pos = np.concatenate( + [ + robot_state_action_data.data["state_left_ee_cartesian_pos"], + robot_state_action_data.data["action_left_ee_cartesian_pos"], + ], + axis=0, + ) + left_ee_rotation = np.concatenate( + [ + robot_state_action_data.data["state_left_ee_rotation"], + robot_state_action_data.data["action_left_ee_rotation"], + ], + axis=0, + ) + left_gripper = np.concatenate( + [ + robot_state_action_data.data["state_left_gripper"], + robot_state_action_data.data["action_left_gripper"], + ], + axis=0, + ) + return np.concatenate( + [left_ee_cartesian_pos, left_ee_rotation, left_gripper], axis=1 + ) + else: + left_arm_joint_pos = np.concatenate( + [ + robot_state_action_data.data["state_left_arm_joint_pos"], + robot_state_action_data.data["action_left_arm_joint_pos"], + ], + axis=0, + ) + return left_arm_joint_pos + + def _get_right_arm_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + if not self.config.robot_use_joint_angle_control: + right_ee_cartesian_pos = np.concatenate( + [ + robot_state_action_data.data["state_right_ee_cartesian_pos"], + robot_state_action_data.data["action_right_ee_cartesian_pos"], + ], + axis=0, + ) + right_ee_rotation = np.concatenate( + [ + robot_state_action_data.data["state_right_ee_rotation"], + robot_state_action_data.data["action_right_ee_rotation"], + ], + axis=0, + ) + right_gripper = np.concatenate( + [ + robot_state_action_data.data["state_right_gripper"], + robot_state_action_data.data["action_right_gripper"], + ], + axis=0, + ) + return np.concatenate( + [right_ee_cartesian_pos, right_ee_rotation, right_gripper], axis=1 + ) + else: + right_arm_joint_pos = np.concatenate( + [ + robot_state_action_data.data["state_right_arm_joint_pos"], + robot_state_action_data.data["action_right_arm_joint_pos"], + ], + axis=0, + ) + return right_arm_joint_pos + + def get_serialized_actions( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> dict: + """Build trimmed/interpolated serialized actions from model_output; does not send. + + Used by apply_action and serving policy; returns dict for robot or client. + """ + assert "robot_state_action_data" in input + + left_arm_action = self._get_left_arm_action(input["robot_state_action_data"]) + right_arm_action = self._get_right_arm_action(input["robot_state_action_data"]) + + # Trim action sequence + action_length = len(left_arm_action) + if robot_action_start_ratio is None: + robot_action_start_ratio = self.config.robot_action_start_ratio + if robot_action_end_ratio is None: + robot_action_end_ratio = self.config.robot_action_end_ratio + start_frame = int(robot_action_start_ratio * action_length) + end_frame = int(robot_action_end_ratio * action_length) + + left_arm_action = left_arm_action[start_frame:end_frame] + right_arm_action = right_arm_action[start_frame:end_frame] + + # Interpolate + if robot_action_interpolate_multiplier is None: + robot_action_interpolate_multiplier = ( + self.config.robot_action_interpolate_multiplier + ) + target_length = robot_action_interpolate_multiplier * len(left_arm_action) + left_arm_action, right_arm_action = ( + UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [left_arm_action, right_arm_action], target_length + ) + ) + + if not self.config.robot_use_joint_angle_control: + return { + "follow1_pos": left_arm_action.tolist(), + "follow2_pos": right_arm_action.tolist(), + } + return { + "follow1_joints": left_arm_action.tolist(), + "follow2_joints": right_arm_action.tolist(), + } + + def apply_action( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> None: + serialized_actions = self.get_serialized_actions( + input, + robot_action_interpolate_multiplier=robot_action_interpolate_multiplier, + robot_action_start_ratio=robot_action_start_ratio, + robot_action_end_ratio=robot_action_end_ratio, + ) + self._send_actions(serialized_actions) + + def _send_actions(self, serialized_actions: dict) -> None: + """Send actions to robot controller.""" + self.robot_controller.robot_comm.send_dict(serialized_actions) + self.is_received = False + self.current_robot_state_action = None + + +class DesktopRobot(Robot): + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.is_received = False # Blocking flag + self.current_robot_state_action = None + + self.logger = InferLogger.get_robot_logger("Robot") + + def _get_dof_mask(self): + dof_config = self.config.train_config["dof_config"] + total_dof = sum(dof_config.values()) + dof_mask = np.ones((1, self.config.action_horizon, total_dof)) + + # For desktop: mask head_actions, height, and velocity_decomposed + mask_keys = ["head_actions", "height", "velocity_decomposed"] + start_idx = 0 + for key, dof_size in dof_config.items(): + if key in mask_keys: + # Zero mask for these DOFs + dof_mask[:, :, start_idx : start_idx + dof_size] = 0 + start_idx += dof_size + + return dof_mask + + +class DesktopRobotVGA(Robot): + def __init__(self, config: InferConfig, robot_id=10) -> None: + super().__init__(config, robot_id) + self.is_received = False # Blocking flag + self.current_robot_state_action = None + + self.logger = InferLogger.get_robot_logger("Robot") + + def _get_dof_mask(self): + dof_config = self.config.train_config["data"]["dof_config"] + total_dof = sum(dof_config.values()) + dof_mask = np.ones((1, self.config.action_horizon, total_dof)) + + # For desktop: mask head_actions, height, and velocity_decomposed + mask_keys = ["head_actions", "height", "velocity_decomposed"] + start_idx = 0 + for key, dof_size in dof_config.items(): + if key in mask_keys: + # Zero mask for these DOFs + dof_mask[:, :, start_idx : start_idx + dof_size] = 0 + start_idx += dof_size + + return dof_mask + + +class DesktopRobotPreprocessor(DesktopRobot): + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.config = config + + self.is_received = False # Blocking flag + self.current_robot_state_action = None + self.logger = InferLogger.get_robot_logger("Robot") + + def _send_actions(self, serialized_actions: dict) -> None: + # bypass + return + + def get_observation(self, state, views): + robot_state_action_data = RobotStateActionData(config=self.config) + for key in robot_action_key_mapping.keys(): + state_key_str = robot_action_key_mapping[key] + value = None + match = re.match(r"(\w+)\[(.*)\]", state_key_str) + if match: + base_key = match.group(1) + if base_key in state: + slicing_str = match.group(2) + slice_parts = slicing_str.split(":") + slice_args = [(int(p) if p.strip() else None) for p in slice_parts] + s = slice(*slice_args) + value = state[base_key][s] + else: + if state_key_str in state: + value = state[state_key_str] + if value is not None: + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + + dof_mask = self._get_dof_mask() + robot_state_action_data.dof_mask = dof_mask + self.current_robot_state_action = robot_state_action_data + return { + "robot_state_action_data": robot_state_action_data, + "face_view": views["camera_front"][0], # (1, 480, 640, 3) + "left_wrist_view": views["camera_left"][0], # (1, 480, 640, 3) + "right_wrist_view": views["camera_right"][0], # (1, 480, 640, 3) + } + + +class TurtleRobot(Robot): + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.vehicle_pose_handler = VehiclePoseHandler() + self.last_speed = [0, 0, 0] + + def _get_state_from_controller(self): + state = super()._get_state_from_controller() + + self.vehicle_pose_handler.update_pose(state["car_pose"]) + state["velocity_decomposed"] = self.last_speed + + if self.config.turtle_as_desktop: + state["head_pos"] = [0, -1] + state["lift"] = [0.4] + + return state + + def _calculate_car_pose(self, base_velocity_pred): + # Integrate velocity to pose (vectorized) + dt = 1 / 20 + current_pose = ( + self.vehicle_pose_handler.current_pose.copy() + if self.vehicle_pose_handler.current_pose is not None + else np.array([0.0, 0.0, 0.0]) + ) + + # Batch integrate positions + poses_frames = [] + for i in range(len(base_velocity_pred)): + current_pose = self.vehicle_pose_handler.velocity_to_pose( + base_velocity_pred[i, 0], + base_velocity_pred[i, 1], + base_velocity_pred[i, 2], + dt, + current_pose, + ) + poses_frames.append(current_pose.copy()) + + return poses_frames + + def _get_car_velocity( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + car_velocity = robot_state_action_data.data["action_velocity_decomposed"] + self.last_speed = car_velocity[-1, :].copy().tolist() + return car_velocity + + def _get_head_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + head_actions = np.concatenate( + [ + robot_state_action_data.data["state_head_actions"], + robot_state_action_data.data["action_head_actions"], + ], + axis=0, + ) + return head_actions + + def _get_height_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + height = np.concatenate( + [ + robot_state_action_data.data["state_height"], + robot_state_action_data.data["action_height"], + ], + axis=0, + ) + return height + + def go_home(self): + if self.current_robot_state_action is None: + self.get_observation() + + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "left_gripper" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "right_gripper" + ) + + self.current_robot_state_action.save_action_data_with_key( + np.array([[0, -1]]), "head_actions" + ) + self.current_robot_state_action.save_action_data_with_key( + np.array([[0.4]]), "height" + ) + self.current_robot_state_action.save_action_data_with_key( + np.array([self.last_speed]), "velocity_decomposed" + ) + + action_dict = {"robot_state_action_data": self.current_robot_state_action} + self.apply_action( + action_dict, + robot_action_interpolate_multiplier=150, + robot_action_start_ratio=0, + robot_action_end_ratio=1, + ) + self.logger.info("Robot returned to home position") + + def get_serialized_actions( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> dict: + """Build Turtle serialized actions (dual arms + head/lift/car_pose); does not send.""" + assert "robot_state_action_data" in input + + left_arm_action = self._get_left_arm_action(input["robot_state_action_data"]) + right_arm_action = self._get_right_arm_action(input["robot_state_action_data"]) + action_length = len(left_arm_action) + if robot_action_start_ratio is None: + robot_action_start_ratio = self.config.robot_action_start_ratio + if robot_action_end_ratio is None: + robot_action_end_ratio = self.config.robot_action_end_ratio + start_frame = int(robot_action_start_ratio * action_length) + end_frame = int(robot_action_end_ratio * action_length) + + left_arm_action = left_arm_action[start_frame:end_frame] + right_arm_action = right_arm_action[start_frame:end_frame] + if robot_action_interpolate_multiplier is None: + robot_action_interpolate_multiplier = ( + self.config.robot_action_interpolate_multiplier + ) + target_length = robot_action_interpolate_multiplier * action_length + ( + left_arm_action, + right_arm_action, + ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [left_arm_action, right_arm_action], + target_length, + ) + serialized_actions = { + "follow1_pos": left_arm_action.tolist(), + "follow2_pos": right_arm_action.tolist(), + } + + if not self.config.turtle_as_desktop: + head_action = self._get_head_action(input["robot_state_action_data"]) + height_action = self._get_height_action(input["robot_state_action_data"]) + car_velocity_action = self._get_car_velocity( + input["robot_state_action_data"] + ) + head_action = head_action[start_frame:end_frame] + height_action = height_action[start_frame:end_frame] + car_velocity_action = car_velocity_action[start_frame:end_frame] + car_pose_action = np.array(self._calculate_car_pose(car_velocity_action)) + ( + head_action, + height_action, + car_pose_action, + ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [head_action, height_action, car_pose_action], + target_length, + ) + serialized_actions["head_pos"] = head_action.tolist() + serialized_actions["lift"] = height_action.tolist() + serialized_actions["car_pose"] = car_pose_action.tolist() + else: + serialized_actions["head_pos"] = [ + [0, -1] for _ in range(len(left_arm_action)) + ] + serialized_actions["lift"] = [0.4 for _ in range(len(left_arm_action))] + serialized_actions["car_pose"] = [ + [0.0, 0.0, 0.0] for _ in range(len(left_arm_action)) + ] + return serialized_actions + + def apply_action( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> None: + serialized_actions = self.get_serialized_actions( + input, + robot_action_interpolate_multiplier=robot_action_interpolate_multiplier, + robot_action_start_ratio=robot_action_start_ratio, + robot_action_end_ratio=robot_action_end_ratio, + ) + self._send_actions(serialized_actions) + + +class TurtleRobotPreprocessor(TurtleRobot): + """Turtle preprocessor: no real robot; builds observations and get_serialized_actions for serving.""" + + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.config = config + self.is_received = False + self.current_robot_state_action = None + self.logger = InferLogger.get_robot_logger("Robot") + self.vehicle_pose_handler = VehiclePoseHandler() + self.last_speed = [0, 0, 0] + + def _send_actions(self, serialized_actions: dict) -> None: + return + + def get_observation(self, state, views): + robot_state_action_data = RobotStateActionData(config=self.config) + for key in robot_action_key_mapping.keys(): + state_key_str = robot_action_key_mapping[key] + value = None + match = re.match(r"(\w+)\[(.*)\]", state_key_str) + if match: + base_key = match.group(1) + if base_key in state: + slicing_str = match.group(2) + slice_parts = slicing_str.split(":") + slice_args = [(int(p) if p.strip() else None) for p in slice_parts] + s = slice(*slice_args) + value = state[base_key][s] + else: + if state_key_str in state: + value = state[state_key_str] + if value is not None: + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + dof_mask = self._get_dof_mask() + robot_state_action_data.dof_mask = dof_mask + self.current_robot_state_action = robot_state_action_data + return { + "robot_state_action_data": robot_state_action_data, + "face_view": views["camera_front"][0], + "left_wrist_view": views["camera_left"][0], + "right_wrist_view": views["camera_right"][0], + } + + +class EX001Robot(Robot): + """EX001 robot: no pose integration; uses velocity_decomposed_odom state for prediction.""" + + def __init__(self, config: InferConfig, robot_id=10) -> None: + super().__init__(config, robot_id) + self.last_speed = [0, 0, 0] + + def _get_velocity_decomposed( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + """Concatenate state and action velocity_decomposed_odom.""" + velocity_decomposed = np.concatenate( + [ + robot_state_action_data.data["state_velocity_decomposed_odom"], + robot_state_action_data.data["action_velocity_decomposed_odom"], + ], + axis=0, + ) + self.last_speed = velocity_decomposed[-1, :].copy().tolist() + return velocity_decomposed + + def _get_head_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + """Concatenate head state/action; supports head_rotation / head_actions.""" + if robot_state_action_data.data.get("state_head_rotation") is not None: + state_key, action_key = "state_head_rotation", "action_head_rotation" + else: + state_key, action_key = "state_head_actions", "action_head_actions" + head_actions = np.concatenate( + [ + robot_state_action_data.data[state_key], + robot_state_action_data.data[action_key], + ], + axis=0, + ) + return head_actions + + def _get_height_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + """Concatenate state and action height sequences.""" + height = np.concatenate( + [ + robot_state_action_data.data["state_height"], + robot_state_action_data.data["action_height"], + ], + axis=0, + ) + return height + + def go_home(self): + """Move to home position.""" + if self.current_robot_state_action is None: + self.get_observation() + + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "left_gripper" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "right_gripper" + ) + + self.current_robot_state_action.save_action_data_with_key( + np.array([[0, -1]]), "head_actions" + ) + self.current_robot_state_action.save_action_data_with_key( + np.array([[0.4]]), "height" + ) + self.current_robot_state_action.save_action_data_with_key( + np.array([self.last_speed]), "velocity_decomposed_odom" + ) + + action_dict = {"robot_state_action_data": self.current_robot_state_action} + self.apply_action( + action_dict, + robot_action_interpolate_multiplier=150, + robot_action_start_ratio=0, + robot_action_end_ratio=1, + ) + self.logger.info("Robot returned to home position") + + def get_serialized_actions( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> dict: + """Build EX001 serialized actions (dual arms + head/lift/velocity_decomposed_odom); does not send. + + Uses predicted velocity_decomposed_odom directly; no pose integration. + """ + assert "robot_state_action_data" in input + + left_arm_action = self._get_left_arm_action(input["robot_state_action_data"]) + right_arm_action = self._get_right_arm_action(input["robot_state_action_data"]) + action_length = len(left_arm_action) + + # action_length is 1 state frame + action_horizon action frames + # Effective action length is action_length - 1 + actual_action_length = action_length - 1 + + if robot_action_start_ratio is None: + robot_action_start_ratio = self.config.robot_action_start_ratio + if robot_action_end_ratio is None: + robot_action_end_ratio = self.config.robot_action_end_ratio + + # Trim indices in action segment (skip frame 0 state) + start_frame = 1 + int(robot_action_start_ratio * actual_action_length) + end_frame = 1 + int(robot_action_end_ratio * actual_action_length) + + left_arm_action = left_arm_action[start_frame:end_frame] + right_arm_action = right_arm_action[start_frame:end_frame] + + if robot_action_interpolate_multiplier is None: + robot_action_interpolate_multiplier = ( + self.config.robot_action_interpolate_multiplier + ) + target_length = robot_action_interpolate_multiplier * len(left_arm_action) + + ( + left_arm_action, + right_arm_action, + ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [left_arm_action, right_arm_action], + target_length, + ) + + serialized_actions = { + "follow1_pos": left_arm_action.tolist(), + "follow2_pos": right_arm_action.tolist(), + } + + # Head, lift, and velocity actions + head_action = self._get_head_action(input["robot_state_action_data"]) + height_action = self._get_height_action(input["robot_state_action_data"]) + velocity_decomposed_action = self._get_velocity_decomposed( + input["robot_state_action_data"] + ) + + # Trim with same start_frame/end_frame + head_action = head_action[start_frame:end_frame] + height_action = height_action[start_frame:end_frame] + velocity_decomposed_action = velocity_decomposed_action[start_frame:end_frame] + + # Interpolate + ( + head_action, + height_action, + velocity_decomposed_action, + ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [head_action, height_action, velocity_decomposed_action], + target_length, + ) + + serialized_actions["head_pos"] = head_action.tolist() + serialized_actions["lift"] = height_action.tolist() + # Send velocity_decomposed_odom; no pose integration + serialized_actions["velocity_decomposed_odom"] = ( + velocity_decomposed_action.tolist() + ) + return serialized_actions + + def apply_action( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> None: + serialized_actions = self.get_serialized_actions( + input, + robot_action_interpolate_multiplier=robot_action_interpolate_multiplier, + robot_action_start_ratio=robot_action_start_ratio, + robot_action_end_ratio=robot_action_end_ratio, + ) + self._send_actions(serialized_actions) + + +class EX001RobotPreprocessor(EX001Robot): + """EX001 preprocessor: no real robot; builds observations for serving.""" + + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.config = config + self.is_received = False + self.current_robot_state_action = None + self.logger = InferLogger.get_robot_logger("Robot") + self.last_speed = [0, 0, 0] + + def _send_actions(self, serialized_actions: dict) -> None: + """No-op send; preprocessing only.""" + return + + def get_observation(self, state, views): + """Build observation from external state and views.""" + robot_state_action_data = RobotStateActionData(config=self.config) + for key in robot_action_key_mapping.keys(): + state_key_str = robot_action_key_mapping[key] + value = None + + # Keep the same 7D adaptation behavior as Robot.get_observation(), + # so offline/serving preprocessor path is consistent with runtime path. + composed_value = self._try_compose_arm_joint_pos_with_gripper(state, key) + if composed_value is not None: + value = composed_value + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + continue + + match = re.match(r"(\w+)\[(.*)\]", state_key_str) + if match: + base_key = match.group(1) + if base_key in state: + slicing_str = match.group(2) + slice_parts = slicing_str.split(":") + slice_args = [(int(p) if p.strip() else None) for p in slice_parts] + s = slice(*slice_args) + value = state[base_key][s] + else: + if state_key_str in state: + value = state[state_key_str] + if value is not None: + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + dof_mask = self._get_dof_mask() + robot_state_action_data.dof_mask = dof_mask + self.current_robot_state_action = robot_state_action_data + + # Build observation; handle key name aliases + observation = { + "robot_state_action_data": robot_state_action_data, + } + + # Try multiple possible view keys + camera_mappings = { + "face_view": ["camera_front", "face_view", "front_view"], + "left_wrist_view": ["camera_left", "left_wrist_view", "left_view"], + "right_wrist_view": ["camera_right", "right_wrist_view", "right_view"], + } + + for obs_key, possible_view_keys in camera_mappings.items(): + found = False + for view_key in possible_view_keys: + if view_key in views and views[view_key] is not None: + view_data = views[view_key] + # If view_data is (1,H,W,3), take [0] + if ( + isinstance(view_data, np.ndarray) + and len(view_data.shape) == 4 + and view_data.shape[0] == 1 + ): + observation[obs_key] = view_data[0] + else: + observation[obs_key] = view_data + found = True + self.logger.info( + f"EX001RobotPreprocessor: using '{view_key}' as '{obs_key}'" + ) + break + if not found: + self.logger.error( + f"EX001RobotPreprocessor: no view for '{obs_key}'; tried keys: {possible_view_keys}" + ) + # Black image placeholder + observation[obs_key] = np.zeros((480, 640, 3), dtype=np.uint8) + + return observation diff --git a/wall_x/_vendor/harrix/serving-old/_wallx_infer/socket_controller.py b/wall_x/_vendor/harrix/serving-old/_wallx_infer/socket_controller.py new file mode 100644 index 00000000..b40eace8 --- /dev/null +++ b/wall_x/_vendor/harrix/serving-old/_wallx_infer/socket_controller.py @@ -0,0 +1,423 @@ +import socket +import struct +import json +import cv2 +import numpy as np +import time +import select +import errno +import os +import pickle + +import threading + +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger + +_DEFAULT_ROBOT_CONFIG = { + 0: {"host": "127.0.0.1", "action_port": 57749, "keyboard_port": 58849}, + 1: {"host": "127.0.0.1", "action_port": 57750, "keyboard_port": 58850}, + 2: {"host": "127.0.0.1", "action_port": 57751, "keyboard_port": 58851}, + 3: {"host": "127.0.0.1", "action_port": 57761, "keyboard_port": 58861}, +} + + +class RobotRegistry: + def __init__(self): + self.registry = _DEFAULT_ROBOT_CONFIG + self.logger = InferLogger.get_controller_logger("RobotRegistry") + self.logger.debug(f"Initialized robot registry with {len(self.registry)} robots") + + def register_robot(self, robot_id, host, action_port, keyboard_port): + if robot_id in self.registry: + self.logger.warning(f"Robot ID {robot_id} is already registered") + return False + else: + self.registry[robot_id] = { + "host": host, + "action_port": action_port, + "keyboard_port": keyboard_port, + } + self.logger.info( + f"Registered robot ID={robot_id}: host={host}, action_port={action_port}, keyboard_port={keyboard_port}" + ) + return True + + def get_robot_info(self, robot_id): + if robot_id in self.registry: + info = self.registry[robot_id] + self.logger.debug(f"Robot ID={robot_id} info: {info}") + return info + else: + self.logger.error(f"Robot ID={robot_id} not found") + return None + + def exist(self, robot_id): + return robot_id in self.registry + + +class RobotCommunication: + def __init__(self, robot_id, host=None, action_port=None, keyboard_port=None): + self.logger = InferLogger.get_controller_logger("RobotCommunication") + self.logger.debug(f"Initializing robot communication: robot_id={robot_id}") + + self.robot_register = RobotRegistry() + if self.robot_register.exist(robot_id): + self.robot_info = self.robot_register.get_robot_info(robot_id) + self.logger.debug(f"Using registered robot config: {robot_id}") + else: + self.robot_register.register_robot( + robot_id, host, action_port, keyboard_port + ) + self.robot_info = self.robot_register.get_robot_info(robot_id) + self.logger.debug(f"Registering new robot config: {robot_id}") + + self.action_sock = None + self.action_conn = None + + self.keyboard_sock = None + self.keyboard_conn = None + + self.client_socks = [] + + def connect(self): + self.logger.info("Establishing socket connection...") + + self.action_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.action_sock.setblocking(True) + self.action_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + + host = self.robot_info["host"] + port = self.robot_info["action_port"] + self.action_sock.bind((host, port)) + self.action_sock.listen(1) + self.logger.info(f"Listening on action port: {host}:{self.action_sock.getsockname()[1]}") + + action_thread = threading.Thread(target=self.handle_action_client) + + self.logger.debug("Starting connection handler thread...") + action_thread.start() + + action_thread.join() + self.logger.info("Socket connection established") + + def handle_action_client(self): + self.logger.debug("Waiting for client connection...") + self.action_conn, addr = self.action_sock.accept() + self.logger.info(f"Accepted connection from {addr}") + + def recv_image(self, index): + self.logger.debug(f"Receiving image index={index}...") + image_size = struct.unpack(" dict: + if not self.views_path: + raise RuntimeError( + "DummyRobotController requires views_path or " + "WALL_X_DUMMY_VIEWS_PATH." + ) + self.logger.debug(f"Loading view data from pickle: {cam_names}") + with open(self.views_path, "rb") as f: + all_views = pickle.load(f) + self.logger.debug(f"View data loaded: {len(all_views)} views") + return all_views + + def recv_action(self): + if not self.state_path: + raise RuntimeError( + "DummyRobotController requires state_path or " + "WALL_X_DUMMY_STATE_PATH." + ) + self.logger.debug("Loading action data from pickle") + with open(self.state_path, "rb") as f: + all_actions = pickle.load(f) + self.logger.debug(f"Action data loaded: {len(all_actions)} fields") + return all_actions + + +class RobotController: + def __init__( + self, + robot_id: int, + host: str = None, + port: str = None, + max_time_step: int = 10000, + ): + self.logger = InferLogger.get_controller_logger("RobotController") + self.logger.info( + f"Initialized robot controller: robot_id={robot_id}, max_time_step={max_time_step}" + ) + + self.robot_comm = RobotCommunication(robot_id, host, port) + self.max_time_step = max_time_step + self.global_step = 0 + + def connect(self): + self.logger.info("Connecting robot controller...") + self.robot_comm.connect() + self.logger.info("Robot controller connected") + + def close(self): + self.logger.info("Closing robot controller...") + self.robot_comm.close() + self.logger.info("Robot controller closed") + + def prediction(self, views, actions) -> dict: + raise NotImplementedError + + def recv_image( + self, cam_names: list = ["camera_left", "camera_front", "camera_right"] + ) -> dict: + self.logger.debug(f"Receiving images: {cam_names}") + views = {} + for i, name in enumerate(cam_names): + image = np.array(self.robot_comm.recv_image(i)) + views[name] = image[None, :] + self.logger.debug(f"Received {name}: shape={image.shape}") + self.logger.debug(f"All images received: {len(views)} views") + return views + + def recv_action(self): + self.logger.debug("Receiving action data...") + action_data = self.robot_comm.recv_action_data() + return action_data + + def recv_keyboard_input(self): + self.robot_comm.accept_connections() + json_data = self.robot_comm.recv_keyboard_input() + if json_data is not None: + self.logger.debug( + f"Keyboard input: motionStatus={json_data.get('motionStatus')}, armMode={json_data.get('armMode')}" + ) + return json_data["motionStatus"], json_data["armMode"] + + return None, None + + def reset(self): + self.logger.info(f"Reset controller: global_step {self.global_step} -> 0") + self.global_step = 0 + + def record_start(self): + self.logger.info("Starting recording...") + record_signal = {"cmd": "RECORD_START"} + self.robot_comm.send_dict(record_signal) + + def record_continue(self): + self.logger.debug("Continuing recording...") + record_signal = {"cmd": "RECORD_CONTINUE"} + self.robot_comm.send_dict(record_signal) + + def record_start_process(self): + if self.global_step == 0: + self.logger.info("Waiting for START signal...") + action = None + while action != "START": + action, _ = self.recv_keyboard_input() + self.logger.info("Received START signal") + self.record_start() + + self.record_continue() + + def set_zero(self): + self.logger.info("Setting zero position...") + record_signal = {"cmd": "INIT_ZERO", "gripper": [0.0, 0.0]} + self.robot_comm.send_dict(record_signal) + self.reset() + self.logger.info("Zero position set") + + def record_stop(self): + self.logger.info("Stopping recording...") + record_signal = {"cmd": "RECORD_STOP"} + self.robot_comm.send_dict(record_signal) + time.sleep(0.01) + + def recover_from_failure(self): + self.logger.warning("Recovering from failure...") + record_signal = {"cmd": "TO_MASTER_SLAVE"} + self.robot_comm.send_dict(record_signal) + time.sleep(0.01) + + self.logger.info("Waiting for START signal...") + action = None + while action != "START": + action, _ = self.recv_keyboard_input() + + self.logger.info("Recovery complete; resuming recording") + self.record_start() + + def reset_to_zero(self): + self.logger.info("Resetting to zero position...") + self.record_stop() + self.set_zero() + + def to_slave(self): + self.logger.info("Switching to slave mode...") + record_signal = {"cmd": "TO_SLAVE"} + self.robot_comm.send_dict(record_signal) + time.sleep(0.01) + + def run( + self, + record_mode=False, + cam_names: list = ["camera_left", "camera_front", "camera_right"], + ): + + while self.global_step <= self.max_time_step: + if record_mode: + self.record_start_process() + + self.global_step += 1 + + action = self.recv_action() + view = self.recv_image(cam_names) + + pred = self.prediction(view, action) + self.robot_comm.send_dict(pred) + + if record_mode: + action, arm_mode = self.recv_keyboard_input() + + if action is not None and arm_mode is not None: + self.logger.info("action: %s, arm_mode: %s", action, arm_mode) + if action == "STOP" and arm_mode == "ARM_TEST_MODE_MS": + self.record_stop() + self.recover_from_failure() + action, arm_mode = None, None + while action != "STOP" or arm_mode != "ARM_TEST_MODE_S": + action, arm_mode = self.recv_keyboard_input() + self.record_stop() + self.to_slave() + self.reset() + elif arm_mode == "ARM_TEST_MODE_S" and action == "INIT": + self.reset_to_zero() diff --git a/wall_x/infer/utils.py b/wall_x/_vendor/harrix/serving-old/_wallx_infer/utils.py similarity index 76% rename from wall_x/infer/utils.py rename to wall_x/_vendor/harrix/serving-old/_wallx_infer/utils.py index e4fba55c..9b49aebe 100644 --- a/wall_x/infer/utils.py +++ b/wall_x/_vendor/harrix/serving-old/_wallx_infer/utils.py @@ -1,21 +1,21 @@ import numpy as np from scipy.signal import savgol_filter -from scipy.spatial.transform import Rotation as R # TODO: Convert to numba functions +from scipy.spatial.transform import Rotation as R # TODO: convert to numba from collections import deque import threading -from wall_x.infer.logger import InferLogger +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger class KeyboardThread(threading.Thread): """ - Simple keyboard listening thread that provides stop and reset functionality + Simple keyboard listener thread with stop and reset """ def __init__(self): self.should_reset = False self.should_stop = False - self.new_instruction_index = None # Used to store new instruction index + self.new_instruction_index = None # Stores new instruction index self.logger = InferLogger.get_utils_logger("KeyboardThread") super(KeyboardThread, self).__init__(name="keyboard-thread", daemon=True) @@ -23,7 +23,7 @@ def __init__(self): self.start() def run(self): - """Listen to keyboard input""" + """Listen for keyboard input.""" while True: try: user_input = input().strip().lower() @@ -33,22 +33,18 @@ def run(self): self.logger.info("[Keyboard] Stop signal sent") elif user_input in ["r", "reset"]: - self.logger.info("[Keyboard] Executing reset...") + self.logger.info("[Keyboard] Resetting...") self.should_reset = True self.logger.info("[Keyboard] Reset signal sent") elif user_input.isdigit(): - # Handle digit input, switch to corresponding instruction index + # Numeric input: switch instruction index index = int(user_input) self.new_instruction_index = index - self.logger.info( - f"[Keyboard] Switched to instruction index: {index}" - ) + self.logger.info(f"[Keyboard] Switched to instruction index: {index}") else: - self.logger.info( - f"[Keyboard] Received input: {user_input}. No action taken." - ) + self.logger.info(f"[Keyboard] Input {user_input}; no action taken.") except EOFError: break @@ -57,11 +53,11 @@ def run(self): def show_help(self): self.logger.info( - "[Keyboard] Keyboard control: Enter 's' to stop, 'r' to reset, 'number' to switch instruction index" + "[Keyboard] Controls: 's' stop, 'r' reset, digit switches instruction index" ) -# Robot arm trajectory parameters +# Arm trajectory parameters ARM_MAX_VELOCITY = 0.02 ARM_EXECUTION_HZ = 20 ARM_MIN_EXECUTION_TIME = 5.0 @@ -69,18 +65,18 @@ def show_help(self): class UnifiedTrajectoryProcessor: - """Unified trajectory processor""" + """Unified trajectory processor.""" @staticmethod def interpolate_trajectory_batch(trajectories, target_length, smooth=True): """ - Batch interpolate multiple trajectories to unified length + Interpolate multiple trajectories to a common length Args: - trajectories: list of np.array, each array with shape (N, D) - target_length: int, target length - smooth: bool, whether to smooth + trajectories: list of np.ndarray, each shape (N, D) + target_length: int, target sequence length + smooth: bool, apply smoothing Returns: - list of np.array, interpolated trajectories + list of np.ndarray, interpolated trajectories """ if not trajectories: return [] @@ -99,8 +95,8 @@ def interpolate_trajectory_batch(trajectories, target_length, smooth=True): original_indices = np.linspace(0, len(traj) - 1, len(traj)) target_indices = np.linspace(0, len(traj) - 1, target_length) - # Handle different types of data - if traj.shape[1] == 7: # Robot arm data [x,y,z,rx,ry,rz,gripper] + # Handle different data types + if traj.shape[1] == 7: # Arm data [x,y,z,rx,ry,rz,gripper] interpolated = UnifiedTrajectoryProcessor._interpolate_arm_trajectory( traj, original_indices, target_indices, target_length ) @@ -111,7 +107,7 @@ def interpolate_trajectory_batch(trajectories, target_length, smooth=True): target_indices, original_indices, traj[:, i] ) - # Smooth processing + # Smoothing if smooth and len(interpolated) >= 5: interpolated = UnifiedTrajectoryProcessor._smooth_trajectory( interpolated @@ -125,7 +121,7 @@ def interpolate_trajectory_batch(trajectories, target_length, smooth=True): def _interpolate_arm_trajectory( traj, original_indices, target_indices, target_length ): - """Optimized robot arm trajectory interpolation""" + """Optimized arm trajectory interpolation.""" interpolated = np.zeros((target_length, 7)) # Vectorized interpolation for position and gripper @@ -140,7 +136,7 @@ def _interpolate_arm_trajectory( target_indices, original_indices, quaternions[:, i] ) - # Batch normalization + # Batch normalize norms = np.linalg.norm(interpolated_quats, axis=1, keepdims=True) interpolated_quats = interpolated_quats / norms @@ -153,7 +149,7 @@ def _interpolate_arm_trajectory( def _interpolate_position_trajectory( traj, original_indices, target_indices, target_length ): - """Optimized position trajectory interpolation""" + """Optimized position trajectory interpolation.""" interpolated = np.zeros((target_length, 3)) for i in range(3): interpolated[:, i] = np.interp(target_indices, original_indices, traj[:, i]) @@ -161,12 +157,12 @@ def _interpolate_position_trajectory( @staticmethod def _smooth_trajectory(trajectory): - """Vectorized smooth processing""" + """Vectorized trajectory smoothing.""" if len(trajectory) < 5: return trajectory try: - # Batch smooth all dimensions + # Smooth all dimensions smoothed = np.zeros_like(trajectory) for dim in range(trajectory.shape[1]): smoothed[:, dim] = savgol_filter( @@ -188,9 +184,9 @@ def _smooth_trajectory(trajectory): @staticmethod def calculate_optimal_trajectory_length(left_traj, right_traj): - """Calculate optimal trajectory length""" + """Compute optimal trajectory length.""" - # Vectorized distance calculation + # Vectorized distance computation def calc_distance(traj): if len(traj) < 2: return 0.0 @@ -213,24 +209,25 @@ def calc_distance(traj): class VehiclePoseHandler: - """Vehicle pose and velocity calculation""" + """Vehicle pose and velocity utilities.""" def __init__(self): self.current_pose = None self.previous_pose = None self.pose_history = deque(maxlen=10) + self.logger = InferLogger.get_utils_logger("VehiclePoseHandler") def update_pose(self, new_pose): - """Update vehicle pose""" + """Update vehicle pose.""" if new_pose is not None: self.previous_pose = self.current_pose self.current_pose = np.array(new_pose) self.pose_history.append(self.current_pose.copy()) - print("current_pose", self.current_pose, flush=True) + self.logger.info("current_pose %s", self.current_pose) return self.current_pose def velocity_to_pose(self, vx_body, vy_body, vyaw, dt, start_pose=None): - """Convert body frame velocity to global frame position""" + """Integrate body-frame velocity to global pose.""" if start_pose is None: if self.current_pose is not None: start_pose = self.current_pose.copy() @@ -239,21 +236,21 @@ def velocity_to_pose(self, vx_body, vy_body, vyaw, dt, start_pose=None): x, y, theta = start_pose - # Convert body frame velocity to global frame displacement + # Body-frame velocity to global displacement cos_theta = np.cos(theta) sin_theta = np.sin(theta) - # Coordinate transformation: body frame -> global frame + # Transform: body frame -> global frame dx_global = (vx_body * cos_theta - vy_body * sin_theta) * dt dy_global = (vx_body * sin_theta + vy_body * cos_theta) * dt dtheta = vyaw * dt - # Calculate new position + # Compute new pose x_new = x + dx_global y_new = y + dy_global theta_new = theta + dtheta - # Constrain angle to [-pi, pi] range + # Wrap angle to [-pi, pi] theta_new = (theta_new + np.pi) % (2 * np.pi) - np.pi return np.array([x_new, y_new, theta_new]) @@ -261,21 +258,21 @@ def velocity_to_pose(self, vx_body, vy_body, vyaw, dt, start_pose=None): def compute_body_velocities_from_poses( self, current_pose, previous_pose, dt=1 / 20 ): - """Compute body frame velocity from pose changes""" + """Compute body-frame velocity from pose delta.""" if current_pose is None or previous_pose is None: return np.array([0.0, 0.0, 0.0]) - # Calculate displacement in global frame + # Displacement in global frame dx_global = current_pose[0] - previous_pose[0] dy_global = current_pose[1] - previous_pose[1] dtheta = current_pose[2] - previous_pose[2] - # Use previous frame's angle for coordinate transformation + # Use previous frame angle for transform theta = previous_pose[2] cos_theta = np.cos(theta) sin_theta = np.sin(theta) - # Convert global frame displacement to body frame velocity + # Global displacement to body-frame velocity vx_body = (dx_global * cos_theta + dy_global * sin_theta) / dt vy_body = (-dx_global * sin_theta + dy_global * cos_theta) / dt vyaw = dtheta / dt diff --git a/wall_x/_vendor/harrix/serving-old/client.py b/wall_x/_vendor/harrix/serving-old/client.py new file mode 100644 index 00000000..b2fc70c1 --- /dev/null +++ b/wall_x/_vendor/harrix/serving-old/client.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +""" +Example client for Wall-X model server with sync support. + +This script demonstrates how to connect to a Wall-X server and request +action predictions from observations in both sync and async contexts. +""" + +import asyncio +import logging +from typing import Dict +import threading +import time +import numpy as np +import torch +import os + + +try: + import msgpack + import msgpack_numpy as m + + m.patch() +except ImportError: + print("Please install msgpack-numpy: pip install msgpack-numpy") + exit(1) + +try: + import websockets +except ImportError: + print("Please install websockets: pip install websockets") + exit(1) + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +class RobotClient: + """Client for connecting to Wall-X model server.""" + + def __init__(self, config_path: str, uri: str = "ws://localhost:8000"): + """Initialize client. + + Args: + uri: WebSocket URI of the server (e.g., ws://localhost:8000) + """ + self.uri = uri + self.websocket = None + self.metadata = None + self._loop = None + self._thread = None + + async def connect(self): + """Connect to the server and receive metadata.""" + logger.info(f"Connecting to {self.uri}...") + self.websocket = await websockets.connect( + self.uri, + ping_interval=None, + ping_timeout=None, + max_size=None, + ) + + self.metadata = msgpack.unpackb(await self.websocket.recv()) + logger.info(f"Connected! Server metadata: {self.metadata}") + + async def predict(self, obs: Dict) -> Dict: + """Get action prediction from observation. + + Args: + obs: Observation dictionary containing: + - 'image': Image array (H, W, C) + - 'prompt': Optional text prompt + - 'state': Optional robot state + + Returns: + Dictionary with: + - 'action': Predicted action array + - 'server_timing': Timing information + """ + if self.websocket is None: + raise RuntimeError("Not connected. Call connect() first.") + + ziped = msgpack.packb(obs) + size_bytes = len(ziped) + print(f"{size_bytes / 1024:.2f} KB") + print(f"{size_bytes / 1024 / 1024:.2f} MB") + await self.websocket.send(ziped) + response = msgpack.unpackb(await self.websocket.recv()) + return response + + async def close(self): + """Close the connection.""" + if self.websocket: + await self.websocket.close() + logger.info("Connection closed") + + async def reset(self): + """Reset the policy (if supported).""" + pass + + # ============ Synchronous methods (using independent thread event loop) ============ + + def _start_background_loop(self): + """Start event loop in background thread.""" + self._loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + def _ensure_loop(self): + """Ensure background event loop is running.""" + if self._loop is None or not self._loop.is_running(): + self._thread = threading.Thread( + target=self._start_background_loop, daemon=True + ) + self._thread.start() + # Wait for loop to start + import time + + while self._loop is None: + time.sleep(0.01) + + def _run_async(self, coro): + """Run coroutine in background event loop.""" + self._ensure_loop() + future = asyncio.run_coroutine_threadsafe(coro, self._loop) + return future.result() + + def connect_sync(self): + """Synchronously connect to server.""" + return self._run_async(self.connect()) + + def predict_sync(self, obs: Dict) -> Dict: + """Synchronous prediction method. + + Args: + obs: Observation dictionary + + Returns: + Prediction result dictionary + """ + return self._run_async(self.predict(obs)) + + def close_sync(self): + """Synchronously close connection.""" + result = self._run_async(self.close()) + # Stop event loop + if self._loop: + self._loop.call_soon_threadsafe(self._loop.stop) + return result + + +# ============ Synchronous version of main function ============ + + +def main_sync(args): + """Synchronous version of main function.""" + + # Create client and connect + client = RobotClient(args.config_path, uri=args.uri) + client.connect_sync() + import pickle + + with open(args.config_path, "rb") as f: + all_views = pickle.load(f) + + with open(args.config_path, "rb") as f: + all_actions = pickle.load(f) + print(all_actions.keys()) + print(all_views.keys()) + + pos = { + "follow1_pos": all_actions["follow1_pos"], + "follow2_pos": all_actions["follow2_pos"], + } + use_base64 = False + if use_base64: + import cv2 + import base64 + + for k, v in all_views.items(): + v = v.squeeze(0) + encoded_img = cv2.imencode(".jpg", v)[1] + print(encoded_img.shape) + + base64_str = base64.b64encode(encoded_img).decode("utf-8") + all_views[k] = base64_str + obs = { + "state": pos, + "views": all_views, + "instruction": "Pick up the cup.", + } + else: + path = args.config_path + follow_pos_1 = torch.load(os.path.join(path, "state_follow1_pos.pt")).numpy() + follow_pos_2 = torch.load(os.path.join(path, "state_follow2_pos.pt")).numpy() + camera_front = torch.load(os.path.join(path, "views_camera_front.pt")).numpy() + camera_left = torch.load(os.path.join(path, "views_camera_left.pt")).numpy() + camera_right = torch.load(os.path.join(path, "views_camera_right.pt")).numpy() + obs = { + "state": { + "follow1_pos": follow_pos_1, + "follow2_pos": follow_pos_2, + }, + "views": { + "camera_front": camera_front, + "camera_left": camera_left, + "camera_right": camera_right, + }, + "instruction": "Pick up the cup.", + } + for _ in range(1): + start_time = time.perf_counter() + response = client.predict_sync(obs) + print(response["action"].shape) + print(response["action"]) + end_time = time.perf_counter() + print(f"Time taken: {(end_time - start_time) * 1000} ms") + action_label = torch.load(args.config_path).numpy() + np.testing.assert_array_equal(response["action"], action_label) + print("compare success") + client.close_sync() + + +if __name__ == "__main__": + """Asynchronous version of main function.""" + import argparse + + parser = argparse.ArgumentParser(description="Wall-X client examples") + parser.add_argument( + "--example", + choices=["single", "multiple", "benchmark"], + default="single", + help="Example to run", + ) + parser.add_argument( + "--uri", + default="ws://localhost:8000", + help="Server URI", + ) + parser.add_argument( + "--pred_horizon", type=int, default=32, help="Prediction horizon" + ) + parser.add_argument("--action_dim", type=int, default=7, help="Action dimension") + parser.add_argument( + "--config_path", + default="/path/to/example_data", + help="Train config path", + ) + parser.add_argument( + "--save_dir", + default="/path/to/example_data", + help="Save directory", + ) + args = parser.parse_args() + + # Synchronous mode + main_sync(args) + + # Asynchronous mode + # asyncio.run(main(args)) diff --git a/wall_x/_vendor/harrix/serving-old/launch_serving.py b/wall_x/_vendor/harrix/serving-old/launch_serving.py new file mode 100644 index 00000000..b5c0916a --- /dev/null +++ b/wall_x/_vendor/harrix/serving-old/launch_serving.py @@ -0,0 +1,268 @@ +#!/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. + +Based on the OpenPI serve_policy.py script structure. +""" + +import dataclasses +import enum +import inspect +import logging +import os +import socket +import sys +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): + from wall_x._vendor.harrix.serving.policy.wall_x_policy 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, + ) + + +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 + # 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 + + # ── 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, + ) +} + + +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'." + ) + policy = get_wallx_policy( + model_config, args.image_passing_mode, args.serialize_actions + ) + 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}") + + # 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 + + # Get network info + hostname = socket.gethostname() + try: + local_ip = socket.gethostbyname(hostname) + except Exception: + local_ip = "unknown" + + logger.info(f"Server hostname: {hostname}") + logger.info(f"Server IP: {local_ip}") + logger.info(f"Server will be available at: ws://{args.host}:{args.port}") + logger.info(f"Health check endpoint: http://{args.host}:{args.port}/healthz") + + batching_str = ( + f"batch_size={args.max_batch_size}, wait={args.max_wait_time_ms}ms" + if args.max_batch_size + else "disabled" + ) + logger.info(f"Batching: {batching_str}") + + # Create and start server + server = WebsocketPolicyServer( + policy=policy, + host=args.host, + port=args.port, + metadata=policy_metadata, + max_batch_size=args.max_batch_size, + max_wait_time_ms=args.max_wait_time_ms, + max_queue_size=args.max_queue_size, + timeout_ms=args.timeout_ms, + ) + + logger.info("Starting server...") + try: + server.serve_forever() + except KeyboardInterrupt: + logger.info("Server stopped by user") + except Exception as e: + logger.error(f"Server error: {e}") + sys.exit(1) + + +if __name__ == "__main__": + main(tyro.cli(Args)) diff --git a/wall_x/serving/policy/__init__.py b/wall_x/_vendor/harrix/serving-old/policy/__init__.py similarity index 100% rename from wall_x/serving/policy/__init__.py rename to wall_x/_vendor/harrix/serving-old/policy/__init__.py diff --git a/wall_x/_vendor/harrix/serving-old/policy/wall_x_policy.py b/wall_x/_vendor/harrix/serving-old/policy/wall_x_policy.py new file mode 100644 index 00000000..8a2eb6b8 --- /dev/null +++ b/wall_x/_vendor/harrix/serving-old/policy/wall_x_policy.py @@ -0,0 +1,250 @@ +import base64 +import logging +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 import WallxModelWrapper +from wall_x._vendor.harrix.serving._wallx_infer.robot import ( + DesktopRobotPreprocessor, + TurtleRobotPreprocessor, + EX001RobotPreprocessor, +) +from wall_x.utils.timers import ScopeTimer + +logger = logging.getLogger(__name__) + +# Inference modes aligned with model_wrapper (extensible for VQA/batch) +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, + ): + """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.model_wrapper = WallxModelWrapper(config) + self.robot_preprocessor = self._register_robot_preprocessor() + 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", + self.image_passing_mode, + config.robot_type, + self.default_infer_mode, + ) + + 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}") + + def reset(self) -> None: + """Reset the policy state.""" + self.action_buffer = [] + self.buffer_index = 0 + logger.debug("Policy reset") + + def _run_action_infer( + self, observation: Dict, instruction: str, mode: str + ) -> Dict[str, Any]: + """Run model_wrapper action inference by mode; returns model_output with robot_state_action_data. + + Modes: flow | ar | flow_with_subtask. VQA can be added here later. + """ + if mode == INFER_MODE_FLOW: + return self.model_wrapper.infer_flow_action(observation, instruction) + 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: + """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["instruction"] + if self.image_passing_mode == "base64": + for k, v in views.items(): + img_bytes = base64.b64decode(v) + img_array = np.frombuffer(img_bytes, dtype=np.uint8) + decoded_img = cv2.imdecode(img_array, cv2.IMREAD_COLOR) + decoded_img = cv2.cvtColor(decoded_img, cv2.COLOR_BGR2RGB) + views[k] = np.expand_dims(decoded_img, axis=0) + + with ScopeTimer("get_observation"): + observation = self.robot_preprocessor.get_observation(state, views) + + mode = obs.get("infer_mode", self.default_infer_mode) + with ScopeTimer(f"infer_{mode}"): + model_output = self._run_action_infer(observation, instruction, mode) + + if not self.serialize_actions: + return model_output + + return self.robot_preprocessor.get_serialized_actions( + model_output, robot_action_interpolate_multiplier=1 + ) # Client interpolates over websocket + + # ── 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["instruction"] + + if self.image_passing_mode == "base64": + import cv2 + + for k, v in views.items(): + img_bytes = base64.b64decode(v) + img_array = np.frombuffer(img_bytes, dtype=np.uint8) + decoded_img = cv2.imdecode(img_array, cv2.IMREAD_COLOR) + decoded_img = cv2.cvtColor(decoded_img, cv2.COLOR_BGR2RGB) + views[k] = np.expand_dims(decoded_img, 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: + if self.serialize_actions: + action = self.robot_preprocessor.get_serialized_actions( + model_output, robot_action_interpolate_multiplier=1 + ) + results.append(action) + else: + predict_action = model_output["predict_action"] + if isinstance(predict_action, torch.Tensor): + predict_action = predict_action.cpu().numpy() + results.append({"action": predict_action}) + + return results + + except Exception as e: + logger.error(f"Batch inference failed: {e}", exc_info=True) + return [{"action": {}, "error": str(e)} for _ in obs_list] + + @property + def metadata(self) -> Dict[str, Any]: + return {"batch_enabled": True, "model": "wall-x"} diff --git a/wall_x/_vendor/harrix/serving-old/scheduler.py b/wall_x/_vendor/harrix/serving-old/scheduler.py new file mode 100644 index 00000000..8d36fdb3 --- /dev/null +++ b/wall_x/_vendor/harrix/serving-old/scheduler.py @@ -0,0 +1,160 @@ +import asyncio +import logging +import time +import uuid +from collections import deque +from dataclasses import dataclass +from typing import Any, Dict, List + +logger = logging.getLogger(__name__) + + +@dataclass +class Request: + request_id: str + obs: Dict[str, Any] + future: asyncio.Future + timestamp: float + + +class RequestScheduler: + def __init__( + self, + policy, + max_batch_size: int = 8, + max_wait_time_ms: float = 100, + max_queue_size: int = 128, + timeout_ms: float = 5000, + ): + self.policy = policy + self.max_batch_size = max_batch_size + self.max_wait_time = max_wait_time_ms / 1000.0 + self.max_queue_size = max_queue_size + self.timeout = timeout_ms / 1000.0 + + self.queue = deque() + self.queue_lock = asyncio.Lock() + self.queue_not_empty = asyncio.Condition(self.queue_lock) + + self.running = False + self.batch_task = None + + async def start(self): + self.running = True + self.batch_task = asyncio.create_task(self._batch_loop()) + logger.info( + f"RequestScheduler started: max_batch={self.max_batch_size}, max_wait={self.max_wait_time*1000}ms" + ) + + async def stop(self): + self.running = False + # Process all remaining requests before shutting down + async with self.queue_lock: + remaining = len(self.queue) + if remaining: + logger.info( + f"Graceful shutdown: processing {remaining} remaining request(s)" + ) + while True: + async with self.queue_lock: + if len(self.queue) == 0: + break + batch = [] + while self.queue and len(batch) < self.max_batch_size: + batch.append(self.queue.popleft()) + if batch: + await self._process_batch(batch) + # Now stop the batch loop + async with self.queue_lock: + self.queue_not_empty.notify_all() + if self.batch_task: + await self.batch_task + logger.info("RequestScheduler stopped") + + async def add_request(self, obs: Dict[str, Any]) -> Dict[str, Any]: + request_id = str(uuid.uuid4()) + future = asyncio.Future() + request = Request( + request_id=request_id, obs=obs, future=future, timestamp=time.monotonic() + ) + + async with self.queue_lock: + if len(self.queue) >= self.max_queue_size: + raise RuntimeError( + f"Queue full: {len(self.queue)}/{self.max_queue_size}" + ) + self.queue.append(request) + self.queue_not_empty.notify() + + try: + result = await asyncio.wait_for(future, timeout=self.timeout) + return result + except asyncio.TimeoutError: + logger.error(f"Request {request_id} timeout after {self.timeout}s") + raise + + async def _batch_loop(self): + while self.running: + try: + batch = await self._collect_batch() + if batch: + await self._process_batch(batch) + except Exception as e: + logger.error(f"Unexpected error in batch loop: {e}", exc_info=True) + + async def _collect_batch(self) -> List[Request]: + async with self.queue_not_empty: + while self.running and len(self.queue) == 0: + await self.queue_not_empty.wait() + + if not self.running: + return [] + + batch = [] + deadline = time.monotonic() + self.max_wait_time + + while len(batch) < self.max_batch_size: + if len(self.queue) > 0: + batch.append(self.queue.popleft()) + + if len(batch) >= self.max_batch_size: + break + + if len(self.queue) == 0: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + await asyncio.wait_for( + self.queue_not_empty.wait(), timeout=remaining + ) + except asyncio.TimeoutError: + break + + return batch + + async def _process_batch(self, batch: List[Request]): + if not batch: + return + + start_time = time.monotonic() + logger.info(f"Processing batch of {len(batch)} requests") + + try: + obs_list = [req.obs for req in batch] + results = await asyncio.to_thread(self.policy.infer_batch, obs_list) + + for req, result in zip(batch, results): + if not req.future.done(): + req.future.set_result(result) + + infer_time = time.monotonic() - start_time + logger.info( + f"Batch processed in {infer_time*1000:.1f}ms, throughput: {len(batch)/infer_time:.1f} req/s" + ) + + except Exception as e: + logger.error(f"Batch processing failed: {e}", exc_info=True) + for req in batch: + if not req.future.done(): + req.future.set_exception(e) diff --git a/wall_x/_vendor/harrix/serving-old/websocket_policy_server.py b/wall_x/_vendor/harrix/serving-old/websocket_policy_server.py new file mode 100644 index 00000000..54ae810e --- /dev/null +++ b/wall_x/_vendor/harrix/serving-old/websocket_policy_server.py @@ -0,0 +1,184 @@ +import asyncio +import http +import logging +import time +import traceback +from typing import Dict, Any, List + +try: + import msgpack + import msgpack_numpy as m + + m.patch() +except ImportError: + logging.warning( + "msgpack-numpy not installed. Install with: pip install msgpack-numpy" + ) + msgpack = None + +import websockets.asyncio.server as _server +import websockets.frames + +logger = logging.getLogger(__name__) + + +class BasePolicy: + """Base class for policies that can be served.""" + + def infer(self, obs: Dict) -> Dict: + """Infer actions from observations.""" + raise NotImplementedError + + def infer_batch(self, obs_list: List[Dict]) -> List[Dict]: + """Batch inference. Default implementation calls infer() per sample.""" + return [self.infer(obs) for obs in obs_list] + + def reset(self) -> None: + """Reset the policy to its initial state.""" + pass + + @property + def metadata(self) -> Dict[str, Any]: + """Return metadata about the policy.""" + return {} + + +class WebsocketPolicyServer: + """Serves a policy using the websocket protocol. + + Implements a websocket server that: + 1. Sends policy metadata on connection + 2. Receives observations + 3. Returns predicted actions (single or batched) + 4. Tracks timing information + + When batching parameters are provided, dynamic batching is enabled: + requests from multiple clients are queued and processed in batches. + """ + + def __init__( + self, + policy: BasePolicy, + host: str = "0.0.0.0", + port: int = 8000, + metadata: Dict | None = None, + # Dynamic batching parameters (None = disabled) + max_batch_size: int | None = None, + max_wait_time_ms: float | None = None, + max_queue_size: int = 100, + timeout_ms: float = 30000, + ) -> None: + self._policy = policy + self._host = host + self._port = port + self._metadata = metadata or {} + + # Dynamic batching + self._scheduler = None + if max_batch_size is not None: + from .scheduler import RequestScheduler + + self._scheduler = RequestScheduler( + policy=policy, + max_batch_size=max_batch_size, + max_wait_time_ms=( + max_wait_time_ms if max_wait_time_ms is not None else 0 + ), + max_queue_size=max_queue_size, + timeout_ms=timeout_ms, + ) + + logging.getLogger("websockets.server").setLevel(logging.INFO) + + @property + def batching_enabled(self) -> bool: + return self._scheduler is not None + + def serve_forever(self) -> None: + asyncio.run(self.run()) + + async def run(self): + # Start the scheduler if batching is enabled + if self._scheduler is not None: + await self._scheduler.start() + + try: + async with _server.serve( + self._handler, + self._host, + self._port, + compression=None, + max_size=None, + ping_interval=None, + ping_timeout=None, + process_request=_health_check, + ) as server: + mode_str = "batched" if self.batching_enabled else "single" + logger.info( + f"Server started on {self._host}:{self._port} (mode={mode_str})" + ) + await server.serve_forever() + finally: + if self._scheduler is not None: + await self._scheduler.stop() + + async def _handler(self, websocket: _server.ServerConnection): + logger.info(f"Connection from {websocket.remote_address} opened") + + if msgpack is None: + await websocket.close( + code=websockets.frames.CloseCode.INTERNAL_ERROR, + reason="msgpack-numpy not installed on server", + ) + return + + # Send metadata to client + await websocket.send(msgpack.packb(self._metadata)) + + prev_total_time = None + while True: + try: + start_time = time.monotonic() + obs = msgpack.unpackb(await websocket.recv()) + + infer_time = time.monotonic() + + if self._scheduler is not None: + # Dynamic batching path + action = await self._scheduler.add_request(obs) + else: + # Single inference path + action = await asyncio.to_thread(self._policy.infer, obs) + + infer_time = time.monotonic() - infer_time + + action["server_timing"] = { + "infer_ms": infer_time * 1000, + } + if prev_total_time is not None: + action["server_timing"]["prev_total_ms"] = prev_total_time * 1000 + + await websocket.send(msgpack.packb(action)) + prev_total_time = time.monotonic() - start_time + + except websockets.ConnectionClosed: + logger.info(f"Connection from {websocket.remote_address} closed") + break + except Exception as e: + logger.error(f"Error handling request: {e}") + await websocket.send(traceback.format_exc()) + await websocket.close( + code=websockets.frames.CloseCode.INTERNAL_ERROR, + reason="Internal server error. Traceback included in previous frame.", + ) + raise + + +def _health_check( + connection: _server.ServerConnection, request: _server.Request +) -> _server.Response | None: + if request.path == "/healthz": + return connection.respond(http.HTTPStatus.OK, "OK\n") + if request.path == "/v2/health/ready": + return connection.respond(http.HTTPStatus.OK, "OK\n") + return None diff --git a/wall_x/_vendor/harrix/serving/__init__.py b/wall_x/_vendor/harrix/serving/__init__.py new file mode 100644 index 00000000..87f42b29 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/__init__.py @@ -0,0 +1,3 @@ +from .websocket_policy_server import WebsocketPolicyServer, BasePolicy + +__all__ = ["WebsocketPolicyServer", "BasePolicy"] diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/__init__.py b/wall_x/_vendor/harrix/serving/_wallx_infer/__init__.py new file mode 100644 index 00000000..195ab41d --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/__init__.py @@ -0,0 +1 @@ +"""Vendored infer subset for harrix serving.""" diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/base_dataclass.py b/wall_x/_vendor/harrix/serving/_wallx_infer/base_dataclass.py new file mode 100644 index 00000000..50e6ba8b --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/base_dataclass.py @@ -0,0 +1,476 @@ +from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig +from typing import Optional, List +from dataclasses import dataclass, field +import numpy as np +import torch +import wall_x._vendor.x2robot_utils.geometry as data_utils +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger + +dof_dims = { + "left_ee_cartesian_pos": 3, + "left_ee_cartesian_pos_relative": 3, + "left_ee_rotation": 3, + "left_ee_rotation_relative": 3, + "left_ee_rotation_6D": 6, + "left_ee_rotation_6D_relative": 6, + "left_arm_joint_pos": 7, # 6 joint + 1 gripper + "left_gripper": 1, + "left_gripper_cur": 1, + "left_arm_joint_cur": 7, # 1 -> 7 + "right_ee_cartesian_pos": 3, + "right_ee_cartesian_pos_relative": 3, + "right_ee_rotation": 3, + "right_ee_rotation_relative": 3, + "right_ee_rotation_6D": 6, + "right_ee_rotation_6D_relative": 6, + "right_arm_joint_pos": 7, # 6 joint + 1 gripper + "right_gripper": 1, + "right_gripper_cur": 1, + "right_arm_joint_cur": 7, # 1 -> 7 + "head_actions": 2, + "height": 1, + "car_pose": 3, + "velocity_decomposed": 3, + "velocity_decomposed_odom": 3, + "head_rotation": 2, # match ex001 + "left_joint": 6, + "left_joint_gripper": 1, + "right_joint": 6, + "right_joint_gripper": 1, + "left_rotation_quat": 4, + "right_rotation_quat": 4, + "left_quaternion": 4, + "right_quaternion": 4, + "left_wrench_ext_local_force": 3, + "left_wrench_ext_local_torque": 3, + "right_wrench_ext_local_force": 3, + "right_wrench_ext_local_torque": 3, + "left_wrench_ext_local_force_from_joint": 3, + "left_wrench_ext_local_torque_from_joint": 3, + "right_wrench_ext_local_force_from_joint": 3, + "right_wrench_ext_local_torque_from_joint": 3, + "left_wrench_ext_world_force": 3, + "left_wrench_ext_world_torque": 3, + "left_wrench_ext_world_force_from_joint": 3, + "left_wrench_ext_world_torque_from_joint": 3, + "right_wrench_ext_world_force": 3, + "right_wrench_ext_world_torque": 3, + "right_wrench_ext_world_force_from_joint": 3, + "right_wrench_ext_world_torque_from_joint": 3, + "left_arm_joint_dev": 7, + "right_arm_joint_dev": 7, +} + + +class ComputedDict(dict): + """Dict that registers compute rules and auto-fills None values on get.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._compute_rules = {} # key -> compute_function + + def register_compute_rule(self, key, compute_func): + """ + Register a compute rule + + Args: + key: Key to compute + compute_func: Callable taking self, returns computed value + """ + self._compute_rules[key] = compute_func + + def get(self, key, default=None): + """Override get to support lazy computation.""" + value = super().get(key, default) + + # If value is None and a rule exists, try to compute + if value is None and key in self._compute_rules: + try: + computed_value = self._compute_rules[key](self) + if computed_value is not None: + # Cache computed value + self[key] = computed_value + return computed_value + except Exception: + pass # On failure return None or default + + return value if value is not None else default + + def __getitem__(self, key): + """Override [] to support lazy computation.""" + value = super().get(key, None) + + # If value is None and a rule exists, try to compute + if value is None and key in self._compute_rules: + try: + computed_value = self._compute_rules[key](self) + if computed_value is not None: + # Cache computed value + self[key] = computed_value + return computed_value + except Exception: + pass # On failure raise KeyError or return None + + if key in self: + return super().__getitem__(key) + raise KeyError(key) + + +@dataclass +class RobotStateActionData: + config: InferConfig = None + data: ComputedDict = field( + default_factory=lambda: ComputedDict( + { + # State (formerly pose) - state_ prefix + "state_left_ee_cartesian_pos": None, # (1, 3) + "state_left_ee_rotation": None, # (1, 3) + "state_left_ee_rotation_6D": None, + "state_left_arm_joint_pos": None, + "state_left_gripper": None, # (1, 1) + "state_left_arm_joint_cur": None, + "state_left_gripper_cur": None, + "state_right_ee_cartesian_pos": None, # (1, 3) + "state_right_ee_rotation": None, + "state_right_ee_rotation_6D": None, # (1, 6) + "state_right_arm_joint_pos": None, + "state_right_gripper": None, + "state_right_gripper_cur": None, + "state_right_arm_joint_cur": None, # (1, 1) + "state_head_actions": None, + "state_head_rotation": None, # match ex001 + "state_height": None, + "state_car_pose": None, + "state_velocity_decomposed": None, + "state_velocity_decomposed_odom": None, + # support joint control + "state_left_joint": None, + "state_left_joint_gripper": None, + "state_right_joint": None, + "state_right_joint_gripper": None, + # support quaternion control + "state_left_quaternion": None, + "state_right_quaternion": None, + "state_left_rotation_quat": None, + "state_right_rotation_quat": None, + # support wrench/force-torque observation + "state_left_wrench_ext_local_force": None, + "state_left_wrench_ext_local_torque": None, + "state_right_wrench_ext_local_force": None, + "state_right_wrench_ext_local_torque": None, + "state_left_wrench_ext_local_force_from_joint": None, + "state_left_wrench_ext_local_torque_from_joint": None, + "state_right_wrench_ext_local_force_from_joint": None, + "state_right_wrench_ext_local_torque_from_joint": None, + "state_left_wrench_ext_world_force": None, + "state_left_wrench_ext_world_torque": None, + "state_left_wrench_ext_world_force_from_joint": None, + "state_left_wrench_ext_world_torque_from_joint": None, + "state_right_wrench_ext_world_force": None, + "state_right_wrench_ext_world_torque": None, + "state_right_wrench_ext_world_force_from_joint": None, + "state_right_wrench_ext_world_torque_from_joint": None, + # support joint deviation + "state_left_arm_joint_dev": None, + "state_right_arm_joint_dev": None, + # Action - action_ prefix + "action_left_ee_cartesian_pos": None, + "action_left_ee_cartesian_pos_relative": None, + "action_left_ee_rotation": None, + "action_left_ee_rotation_relative": None, + "action_left_ee_rotation_6D": None, + "action_left_ee_rotation_6D_relative": None, + "action_left_gripper": None, + "action_left_arm_joint_pos": None, + "action_right_ee_cartesian_pos": None, + "action_right_ee_cartesian_pos_relative": None, + "action_right_ee_rotation": None, + "action_right_ee_rotation_relative": None, + "action_right_ee_rotation_6D": None, + "action_right_ee_rotation_6D_relative": None, + "action_right_gripper": None, + "action_right_arm_joint_pos": None, + "action_head_actions": None, + "action_head_rotation": None, # match ex001 + "action_height": None, + "action_car_pose": None, + "action_velocity_decomposed": None, + "action_velocity_decomposed_odom": None, + # support joint control + "action_left_joint": None, + "action_left_joint_gripper": None, + "action_right_joint": None, + "action_right_joint_gripper": None, + # support quaternion control + "action_left_quaternion": None, + "action_right_quaternion": None, + "action_left_rotation_quat": None, + "action_right_rotation_quat": None, + } + ) + ) + dof_mask: np.ndarray = None + logger = InferLogger.get_robot_logger("RobotStateActionData") + + def __post_init__(self): + """Register a compute rule""" + # State rules: euler angles -> 6D rotation + self.data.register_compute_rule( + "state_left_ee_rotation_6D", + lambda d: ( + data_utils.euler_to_matrix_zyx_6d_nb(d["state_left_ee_rotation"]) + if d.get("state_left_ee_rotation") is not None + else None + ), + ) + self.data.register_compute_rule( + "state_right_ee_rotation_6D", + lambda d: ( + data_utils.euler_to_matrix_zyx_6d_nb(d["state_right_ee_rotation"]) + if d.get("state_right_ee_rotation") is not None + else None + ), + ) + + # Action rules: absolute position from relative + state + self.data.register_compute_rule( + "action_left_ee_cartesian_pos", + lambda d: ( + d.get("state_left_ee_cartesian_pos") + + d.get("action_left_ee_cartesian_pos_relative") + if d.get("state_left_ee_cartesian_pos") is not None + and d.get("action_left_ee_cartesian_pos_relative") is not None + else None + ), + ) + self.data.register_compute_rule( + "action_right_ee_cartesian_pos", + lambda d: ( + d.get("state_right_ee_cartesian_pos") + + d.get("action_right_ee_cartesian_pos_relative") + if d.get("state_right_ee_cartesian_pos") is not None + and d.get("action_right_ee_cartesian_pos_relative") is not None + else None + ), + ) + + # Action rules: derive absolute RPY + self.data.register_compute_rule( # delta rpy -> abs rpy + "action_left_ee_rotation", + lambda d: ( + data_utils.compose_state_and_delta_to_abs_rpy( + d["action_left_ee_rotation_relative"], + d["state_left_ee_rotation"][0], + ) + if d.get("action_left_ee_rotation_relative") is not None + and d.get("state_left_ee_rotation") is not None + else None + ), + ) + self.data.register_compute_rule( # abs 6D -> abs rpy + "action_left_ee_rotation", + lambda d: ( + data_utils.so3_to_euler_zyx_batch_nb(d["action_left_ee_rotation_6D"]) + if d.get("action_left_ee_rotation_6D") is not None + else None + ), + ) + self.data.register_compute_rule( # delta 6D -> abs 6D + "action_left_ee_rotation_6D", + lambda d: ( + data_utils.compose_state_and_delta_to_abs_6d( + d["action_left_ee_rotation_6D_relative"], + d["state_left_ee_rotation_6D"][0], + ) + if d.get("action_left_ee_rotation_6D_relative") is not None + and d.get("state_left_ee_rotation_6D") is not None + else None + ), + ) + + self.data.register_compute_rule( # delta rpy -> abs rpy + "action_right_ee_rotation", + lambda d: ( + data_utils.compose_state_and_delta_to_abs_rpy( + d["action_right_ee_rotation_relative"], + d["state_right_ee_rotation"][0], + ) + if d.get("action_right_ee_rotation_relative") is not None + and d.get("state_right_ee_rotation") is not None + else None + ), + ) + self.data.register_compute_rule( # abs 6D -> abs rpy + "action_right_ee_rotation", + lambda d: ( + data_utils.so3_to_euler_zyx_batch_nb(d["action_right_ee_rotation_6D"]) + if d.get("action_right_ee_rotation_6D") is not None + else None + ), + ) + self.data.register_compute_rule( # delta 6D -> abs 6D + "action_right_ee_rotation_6D", + lambda d: ( + data_utils.compose_state_and_delta_to_abs_6d( + d["action_right_ee_rotation_6D_relative"], + d["state_right_ee_rotation_6D"][0], + ) + if d.get("action_right_ee_rotation_6D_relative") is not None + and d.get("state_right_ee_rotation_6D") is not None + else None + ), + ) + + def get_agent_pos(self, obs_action_keys=None): + if obs_action_keys is None: + obs_action_keys = self.config.train_config["agent_pos_config"].keys() + + agent_pose_data = [] + for key in obs_action_keys: + # action_padding is a virtual key used to pad agent_pos width. + if key == "action_padding": + dim = self.config.train_config["agent_pos_config"][key] + agent_pose_data.append(np.zeros((1, dim))) + continue + + # Strip follow_ or master_ prefix + if key.startswith("follow_"): + key = key.replace("follow_", "") + elif key.startswith("master_"): + key = key.replace("master_", "") + + # Add state_ prefix to access state data + state_key = f"state_{key}" + + if state_key in self.data: + # get() auto-computes when value is None + value = self.data.get(state_key) + if value is None: + # If still None after compute, use zero vector + agent_pose_data.append(np.zeros((1, dof_dims[key]))) + else: + agent_pose_data.append(value) + else: + raise ValueError(f"Key {state_key} not found in data") + + agent_pose_data = np.concatenate(agent_pose_data, axis=1)[None] # (1, 1, D) + + return agent_pose_data + + def get_agent_pos_mask(self, obs_action_keys=None): + if obs_action_keys is None: + obs_action_keys = self.config.train_config["agent_pos_config"].keys() + + agent_pos_mask_data = [] + for key in obs_action_keys: + # action_padding carries no information and should stay masked out. + if key == "action_padding": + dim = self.config.train_config["agent_pos_config"][key] + agent_pos_mask_data.append(np.zeros((1, dim))) + continue + + # Strip follow_ or master_ prefix + if key.startswith("follow_"): + key = key.replace("follow_", "") + elif key.startswith("master_"): + key = key.replace("master_", "") + + # Add state_ prefix to access state data + state_key = f"state_{key}" + + if state_key in self.data: + # get() auto-computes when value is None + value = self.data.get(state_key) + if value is None: + agent_pos_mask_data.append(np.zeros((1, dof_dims[key]))) + else: + agent_pos_mask_data.append(np.ones((1, value.shape[1]))) + else: + raise ValueError(f"Key {state_key} not found in data") + + return np.concatenate(agent_pos_mask_data, axis=1)[None] # (1, 1, D) + + def save_state_data_with_key(self, value, key, gt_dim=None): + # Strip follow_ or master_ prefix + key = key.replace("follow_", "") + key = key.replace("master_", "") + + # if torch, convert to numpy + if isinstance(value, torch.Tensor): + value = value.detach().cpu().numpy() + + if f"state_{key}" not in self.data: # TODO： joint angle control + self.logger.warning(f"{key} is not a valid state key; not recorded") + return + gt_dim = dof_dims[key] if gt_dim is None else gt_dim + + # Shape check: expected (1, D) + if value.shape == (1, gt_dim): + self.data[f"state_{key}"] = value + elif value.shape == (1, 1, gt_dim): + self.data[f"state_{key}"] = value[0] + elif value.shape == (gt_dim,): + self.data[f"state_{key}"] = value[None] + else: + raise ValueError(f"Value shape {value.shape} is not legal") + + def save_action_data_with_key(self, value, key): + key = key.replace("follow_", "") + key = key.replace("master_", "") + + if isinstance(value, torch.Tensor): + value = value.detach().cpu().numpy() + + if value.shape == (dof_dims[key],): + self.data[f"action_{key}"] = value[None] + else: + self.data[f"action_{key}"] = value + + def save_action_data( + self, predict_action, predict_action_keys: Optional[List[str]] = None + ): + if predict_action_keys is None: + predict_action_keys = getattr( + self.config.data_config, "predict_action_keys", None + ) + if predict_action_keys is None: + try: + predict_action_keys = self.config.data_config["predict_action_keys"] + except (KeyError, AttributeError): + predict_action_keys = list( + self.config.train_config["dof_config"].keys() + ) + + if isinstance(predict_action, torch.Tensor): + predict_action = predict_action.detach().cpu().numpy() + + if predict_action.ndim == 3: + predict_action = predict_action[0] + + dof_start = 0 + for action_key in predict_action_keys: + # action_padding is a virtual key that only advances dof_start; no data is written. + if action_key == "action_padding": + dof_dim = self.config.train_config["dof_config"]["action_padding"] + dof_start += dof_dim + continue + + action_key = action_key.replace("follow_", "") + action_key = action_key.replace("master_", "") + dof_dim = dof_dims[action_key] + action_key = f"action_{action_key}" + self.data[action_key] = predict_action[:, dof_start : dof_start + dof_dim] + dof_start += dof_dim + + # Convenience properties for compatibility + @property + def agent_pos(self): + return self.get_agent_pos() + + @property + def agent_pos_mask(self): + return self.get_agent_pos_mask() + + @property + def action(self): + pass # TODO: support action access diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py b/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py new file mode 100644 index 00000000..7ed332cd --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py @@ -0,0 +1,434 @@ +import logging +import os +from dataclasses import dataclass, field +from typing import Any + +import yaml +from qwen_vl_utils.vision_process import IMAGE_FACTOR, MAX_PIXELS, MIN_PIXELS + +logger = logging.getLogger(__name__) + + +def _env_bool(name: str, default: bool) -> bool: + v = os.environ.get(name) + if v is None: + return default + return v.strip().lower() in ("1", "true", "yes", "on") + + +@dataclass +class InferenceDataConfig: + """Minimal data config needed by online inference. + + Serving does not build datasets, so it should not require closed-source + data backends just to read image resize settings from a checkpoint yaml. + """ + + resolution: dict[str, int] = field(default_factory=dict) + model_type: str = "qwen2_5" + max_pixels: int = MAX_PIXELS + min_pixels: int = MIN_PIXELS + image_factor: int = IMAGE_FACTOR + use_relative_action: bool = False + predict_action_keys: list[str] = field(default_factory=list) + + def __getitem__(self, key: str) -> Any: + if not hasattr(self, key): + raise KeyError(key) + return getattr(self, key) + + +class InferConfig: + def __init__( + self, + checkpoint_path: str | None = None, + train_config_path: str | None = None, + robot_host: str = "0.0.0.0", + robot_port: int = 41776, + robot_type: str = "desktop", # ["desktop", "turtle", "ex001"] + robot_action_start_ratio: float = 0, # Start ratio when trimming executed action + robot_action_end_ratio: float = 0.8, # End ratio when trimming executed action + robot_action_interpolate_multiplier: int = 10, # Action interpolation multiplier + robot_use_joint_angle_control: bool = False, # Joint control (model must predict joints) + turtle_as_desktop: bool = False, # Turtle as desktop: fixed base/head cam/height + action_horizon: int = 32, # Must match model action horizon + action_dim: int | None = None, + ar_action_dim: int | None = None, + model_device: str = "cuda:0", + num_inference_timesteps: int = 10, + num_inference_steps: int | None = None, + cfg_scale: float | None = None, + seed: int | None = None, + norm_key: str = "x2_normal", # ["x2_normal", "ex_normal"] + cam_names: list[str] | None = None, + camera_front_key: str = "camera_front", + camera_left_key: str = "camera_left", + camera_right_key: str = "camera_right", + default_instruction: str | None = None, + prompt_template: str | None = None, + qwen25_prompt_template: str | None = None, + prompt_priority_order: str | None = None, + save_video_dir: str = "./videos", + robot_id: str = "10000", + model_type: str = "wallx", # ["wallx", "vga"] + smooth_action: bool = False, + smooth_gripper: bool = True, + ): + # Private path storage + assert checkpoint_path is not None + from wall_x._vendor.harrix.utils.ckpt_load import resolve_checkpoint_dir + + if not os.path.isdir(checkpoint_path) and not os.path.isfile(checkpoint_path): + raise FileNotFoundError( + f"Checkpoint path not found: {checkpoint_path!r}. " + "Serving requires a directory containing model.safetensors and " + "normalizer_action.pth / normalizer_propri.pth (or norm_stats.json)." + ) + self._checkpoint_path = resolve_checkpoint_dir(checkpoint_path) + action_pth = os.path.join(self._checkpoint_path, "normalizer_action.pth") + propri_pth = os.path.join(self._checkpoint_path, "normalizer_propri.pth") + if os.path.exists(action_pth): + self.normalizer_action_path = action_pth + if os.path.exists(propri_pth): + self.normalizer_propri_path = propri_pth + + # Other config fields + self.robot_host = robot_host + self.robot_port = robot_port + self.robot_type = robot_type # ["desktop", "turtle", "ex001"] + self.robot_action_start_ratio = robot_action_start_ratio + self.robot_action_end_ratio = robot_action_end_ratio + self.robot_action_interpolate_multiplier = robot_action_interpolate_multiplier + self.robot_use_joint_angle_control = ( + robot_use_joint_angle_control # Joint-angle control + ) + self.turtle_as_desktop = turtle_as_desktop + self.robot_id = robot_id + self._action_horizon = ( + action_horizon # Default from train config flow action horizon + ) + self._action_dim = action_dim # Default from train config dof_config + self.ar_action_dim = ar_action_dim # Default from train config ar_dof_config + self.model_device = model_device + self.num_inference_timesteps = ( + num_inference_timesteps # flow matching related config + ) + self.num_inference_steps = num_inference_steps + self.cfg_scale = cfg_scale + self.seed = seed + self.save_video_dir = save_video_dir + self.model_type = model_type # ["wallx", "vga"] + # Config objects + self.train_config: dict = {} + self.model_config = None + self.data_config = None + self.norm_key = norm_key + self.cam_names = cam_names or [ + "face_view", + "left_wrist_view", + "right_wrist_view", + ] + self.camera_front_key = camera_front_key + self.camera_left_key = camera_left_key + self.camera_right_key = camera_right_key + self.default_instruction = default_instruction + self.prompt_template = prompt_template + self.qwen25_prompt_template = qwen25_prompt_template + self.prompt_priority_order = prompt_priority_order + self.smooth_action = _env_bool("WALLX_SMOOTH_ACTION", smooth_action) + self.smooth_gripper = _env_bool("WALLX_SMOOTH_GRIPPER", smooth_gripper) + # Load all configs + self._load_all_configs(train_config_path) + self._apply_cam_names_from_train_config(cam_names) + + def _apply_cam_names_from_train_config( + self, cam_names: list[str] | None + ) -> None: + """Use train YAML camera mapping when CLI did not override cam_names.""" + if cam_names is not None: + return + from wall_x._vendor.harrix.utils.train_config import ( + resolve_cam_names_from_train_config, + ) + + resolved = resolve_cam_names_from_train_config(self.train_config) + if resolved: + self.cam_names = resolved + logger.info( + "[InferConfig] cam_names from train key_mappings: %s", + self.cam_names, + ) + + @property + def checkpoint_path(self) -> str | None: + return self._checkpoint_path + + @checkpoint_path.setter + def checkpoint_path(self, value: str | None): + """Reload all configs when checkpoint_path changes.""" + if self._checkpoint_path != value: + self._checkpoint_path = value + self._load_all_configs() + + @property + def action_horizon(self) -> int: + return self._action_horizon + + @action_horizon.setter + def action_horizon(self, value: int): + self._action_horizon = value + + @property + def action_dim(self) -> int | None: + return self._action_dim + + @action_dim.setter + def action_dim(self, value: int | None): + self._action_dim = value + + @property + def ar_action_dim(self) -> int | None: + return self._ar_action_dim + + @ar_action_dim.setter + def ar_action_dim(self, value: int | None): + self._ar_action_dim = value + + def _load_all_configs(self, train_config_path=None): + """Load all configs.""" + self._load_train_config(train_config_path) + if self.model_type != "vga": + self._load_model_config_for_wallx() + self._load_data_config() + + # Update action_horizon/action_dim if needed + if self._action_horizon is None: + self._action_horizon = self.train_config.get("data", {}).get( + "action_horizon_flow", 32 + ) + assert self._action_horizon is not None and self._action_horizon > 0 + + if self._action_dim is None: + dof_config = ( + self.train_config.get("dof_config") + or self.train_config.get("task", {}).get("dof_config") + or self.train_config.get("data", {}).get("dof_config", {}) + ) + self._action_dim = sum(dof_config.values()) + + if self._ar_action_dim is None: + ar_dof_config = self.train_config.get("ar_dof_config") or self.train_config.get( + "task", {} + ).get("ar_dof_config", {}) + self._ar_action_dim = sum(ar_dof_config.values()) + + def _load_train_config(self, train_config_path): + if train_config_path is None: + for fname in ("config.yml", "config.yaml"): + candidate = os.path.join(self._checkpoint_path, fname) + if os.path.exists(candidate): + train_config_path = candidate + break + else: + raise FileNotFoundError( + f"No config.yml/config.yaml found in {self._checkpoint_path}" + ) + with open(train_config_path, "r") as f: + self.train_config = yaml.load(f, Loader=yaml.FullLoader) + + ckpt_dir = self._checkpoint_path + preprocessor_file = os.path.join(ckpt_dir, "preprocessor_config.json") + if os.path.exists(preprocessor_file): + logger.info( + "[LoadConfig] Found %s, override processor_path.", + preprocessor_file, + ) + self.train_config["processor_path"] = ckpt_dir + + tokenizer_file = os.path.join(ckpt_dir, "tokenizer.json") + tokenizer_config_file = os.path.join(ckpt_dir, "tokenizer_config.json") + if self.train_config.get( + "action_tokenizer_path", None + ) is not None and not os.path.exists( + self.train_config.get("action_tokenizer_path", None) + ): + if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file): + logger.info( + "[LoadConfig] Found tokenizer files in %s, override action_tokenizer_path.", + ckpt_dir, + ) + self.train_config["action_tokenizer_path"] = ckpt_dir + else: + logger.warning("[LoadConfig] Cannot load action tokenizer! ") + + from wall_x._vendor.harrix.utils.train_config import ( + normalize_train_config_for_inference, + strip_action_tokenizer_fields, + ) + + self.train_config = normalize_train_config_for_inference( + self.train_config, train_config_path + ) + # Flow serving does not load AR action tokenizers; crop embed_tokens instead. + strip_action_tokenizer_fields(self.train_config) + self._train_config_path = train_config_path + + def _load_model_config_for_wallx(self): + model_type = self.train_config["model_type"] + if not os.path.isdir(self._checkpoint_path): + return + # For Qwen models + ckpt_config_path = os.path.join(self._checkpoint_path, "config.json") + resolved_cfg_path = None + + if os.path.exists(ckpt_config_path): + # Prefer checkpoint config + resolved_cfg_path = ckpt_config_path + logger.info( + "[LoadModelConfig] Using checkpoint config.json: %s", + ckpt_config_path, + ) + else: + # Fallback to original config path + fallback_cfg = self.train_config.get("qwen_vl_act_config_path", None) + if fallback_cfg is not None: + resolved_cfg_path = fallback_cfg + logger.info( + "[LoadModelConfig] Using fallback act config: %s", + fallback_cfg, + ) + + if resolved_cfg_path is None or (not os.path.exists(resolved_cfg_path)): + raise ValueError( + f"[LoadModelConfig] Cannot load model config! " + f"Checked:\n" + f" - Checkpoint config.json: {ckpt_config_path}\n" + f" - Fallback path: {self.train_config.get('qwen_vl_act_config_path', None)}" + ) + + # Save back to config for consistency + self.train_config["qwen_vl_act_config_path"] = resolved_cfg_path + + from wall_x.trainer.adapters import resolve_adapter + + adapter_cls = resolve_adapter(model_type) + ConfigClass = adapter_cls.config_class() + + logger.info( + "[LoadModelConfig] Loading model config from: %s", resolved_cfg_path + ) + if resolved_cfg_path.endswith(".json"): + self.model_config = ConfigClass.from_json_file(resolved_cfg_path) + else: + self.model_config = ConfigClass.from_pretrained(resolved_cfg_path) + + self.model_config.update_model_config(self.train_config) + + self.model_config._attn_implementation = "sdpa" + vision_attn = os.environ.get("WALLX_VISION_ATTN_IMPLEMENTATION", "flash_attention_2") + self.model_config.vision_config._attn_implementation = vision_attn + logger.info( + "[LoadModelConfig] vision _attn_implementation=%s (override via WALLX_VISION_ATTN_IMPLEMENTATION)", + vision_attn, + ) + + logger.info("[LoadModelConfig] Model config loaded and updated successfully.") + + def _load_data_config(self): + # Prefer typed TrainConfig path (handles new 8-section schema where + # dof_config lives under task:). Fall back to legacy raw-dict path + # for old flat yamls. See trainer/adapters/base_adapter.py. + typed_dcfg = self._try_typed_data_config() + if typed_dcfg is not None: + self.data_config = typed_dcfg + elif self.model_type != "vga": + self.data_config = self._build_inference_data_config() + else: + # TEMPORARY: direct call to the private _set_data_backend. + # VGA inference still uses the legacy x2robot data config object. + # Wall-X online inference uses the lightweight InferenceDataConfig + # above so serving does not require excluded data backends. + from wall_x.data._registry import _set_data_backend + + _set_data_backend(self.train_config.get("dataset_type", "x2robot_v1")) + from wall_x.model.vga.openloop_visualization import get_data_configs + + dataload_config = get_data_configs(self.train_config.get("data", {})) + dataload_config["predict_action_keys"] = list( + dataload_config.get("dof_config", {}).keys() + ) + self.data_config = dataload_config + + self._ensure_predict_action_keys() + + def _ensure_predict_action_keys(self) -> None: + dof_keys = list((self.train_config.get("dof_config") or {}).keys()) + if not dof_keys: + return + if isinstance(self.data_config, InferenceDataConfig): + if not self.data_config.predict_action_keys: + self.data_config.predict_action_keys = dof_keys + elif isinstance(self.data_config, dict): + self.data_config.setdefault("predict_action_keys", dof_keys) + + def _build_inference_data_config(self) -> InferenceDataConfig: + data = self.train_config.get("data", {}) or {} + + def get(key: str, default: Any) -> Any: + return data.get(key, self.train_config.get(key, default)) + + resolution = get("resolution", {}) or {} + dof_config = self.train_config.get("dof_config") or {} + return InferenceDataConfig( + resolution=dict(resolution), + model_type=get( + "model_type", self.train_config.get("model_type", "qwen2_5") + ), + max_pixels=get("max_pixels", MAX_PIXELS), + min_pixels=get("min_pixels", MIN_PIXELS), + image_factor=get("image_factor", IMAGE_FACTOR), + use_relative_action=get("use_relative_action", False), + predict_action_keys=list(dof_config.keys()), + ) + + def _try_typed_data_config(self): + """Attempt to load TrainConfig and build X2RDataConfig via typed path. + + Returns the X2RDataConfig on success, or None if the yaml is not in + TrainConfig schema (legacy flat yaml). + """ + yml_path = getattr(self, "_train_config_path", None) + if yml_path is None: + for fname in ("config.yml", "config.yaml"): + candidate = os.path.join(self._checkpoint_path, fname) + if os.path.exists(candidate): + yml_path = candidate + break + if yml_path is None: + return None + try: + from wall_x.config.loader import load_config + from wall_x.trainer.adapters.base_adapter import load_trainer_data_config + + typed_cfg = load_config(yml_path) + except Exception as e: + logger.warning( + "[InferConfig] Typed config load failed for %s: %s", + yml_path, + e, + ) + return None + try: + return load_trainer_data_config(typed_cfg) + except Exception as e: + logger.warning( + "[InferConfig] Typed config load failed for %s: %s", + yml_path, + e, + ) + return None + + +if __name__ == "__main__": + config = InferConfig() + logger.info("%s", config.train_config) diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/logger.py b/wall_x/_vendor/harrix/serving/_wallx_infer/logger.py new file mode 100644 index 00000000..11ef3e82 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/logger.py @@ -0,0 +1,277 @@ +"""Layered logging helpers for inference components.""" + +import logging +import sys +from typing import Optional +from pathlib import Path +from datetime import datetime + +try: + import colorlog + + HAS_COLORLOG = True +except ImportError: + HAS_COLORLOG = False + logging.getLogger(__name__).warning( + "colorlog not installed. Install with: pip install colorlog" + ) + + +class InferLogger: + """ + Layered inference logging + """ + + _loggers = {} + _initialized = False + + # Level identifiers + LEVEL_ENV = "ENV" + LEVEL_ROBOT = "ROBOT" + LEVEL_CONTROLLER = "CONTROLLER" + LEVEL_MODEL = "MODEL" + LEVEL_UTILS = "UTILS" + + # Level colors for console + LEVEL_COLORS = { + LEVEL_ENV: "cyan", + LEVEL_ROBOT: "green", + LEVEL_CONTROLLER: "yellow", + LEVEL_MODEL: "purple", # colorlog uses 'purple' not 'magenta' + LEVEL_UTILS: "blue", + } + + @classmethod + def setup( + cls, + log_level: str = "INFO", + log_dir: Optional[str] = None, + console_output: bool = True, + file_output: bool = True, + colorful: bool = True, + ): + """ + Initialize logging + + Args: + log_level: DEBUG, INFO, WARNING, ERROR, CRITICAL + log_dir: Log file directory + console_output: Enable console output + file_output: Enable file output + colorful: Colored console (requires colorlog) + """ + if cls._initialized: + return + + cls.log_level = getattr(logging, log_level.upper()) + cls.console_output = console_output + cls.file_output = file_output + cls.colorful = colorful and HAS_COLORLOG + + # Create log directory + if file_output and log_dir: + cls.log_dir = Path(log_dir) + cls.log_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + cls.log_file = cls.log_dir / f"infer_{timestamp}.log" + else: + cls.log_file = None + + cls._initialized = True + + @classmethod + def get_logger(cls, name: str, level: str = None) -> logging.Logger: + """ + Get logger for a level + + Args: + name: Logger name (module or class) + level: Level tag (ENV, ROBOT, CONTROLLER, MODEL, UTILS) + + Returns: + Configured logger instance + """ + if not cls._initialized: + cls.setup() + + # Auto-detect level from name + if level is None: + level = cls._detect_level(name) + + logger_key = f"{level}.{name}" + + if logger_key in cls._loggers: + return cls._loggers[logger_key] + + # Create new logger + logger = logging.getLogger(logger_key) + logger.setLevel(cls.log_level) + logger.propagate = False + + # Clear existing handlers + logger.handlers.clear() + + # Console handler + if cls.console_output: + console_handler = logging.StreamHandler(sys.stdout) + console_handler.setLevel(cls.log_level) + + if cls.colorful: + # Colored formatter + color = cls.LEVEL_COLORS.get(level, "white") + console_format = ( + f"%(log_color)s[%(asctime)s]%(reset)s " + f"%(bold_{color})s[{level:^10}]%(reset)s " + f"%(bold_white)s[%(name)s]%(reset)s " + f"%(log_color)s%(levelname)-8s%(reset)s " + f"%(message)s" + ) + + console_formatter = colorlog.ColoredFormatter( + console_format, + datefmt="%H:%M:%S", + log_colors={ + "DEBUG": "cyan", + "INFO": "green", + "WARNING": "yellow", + "ERROR": "red", + "CRITICAL": "bold_red", + }, + ) + else: + # Plain formatter + console_format = ( + f"[%(asctime)s] [{level:^10}] [%(name)s] " + f"%(levelname)-8s %(message)s" + ) + console_formatter = logging.Formatter( + console_format, datefmt="%H:%M:%S" + ) + + console_handler.setFormatter(console_formatter) + logger.addHandler(console_handler) + + # File handler + if cls.file_output and cls.log_file: + file_handler = logging.FileHandler(cls.log_file, encoding="utf-8") + file_handler.setLevel(cls.log_level) + + file_format = ( + f"[%(asctime)s] [{level:^10}] [%(name)s] " + f"%(levelname)-8s %(message)s" + ) + file_formatter = logging.Formatter(file_format, datefmt="%Y-%m-%d %H:%M:%S") + file_handler.setFormatter(file_formatter) + logger.addHandler(file_handler) + + cls._loggers[logger_key] = logger + return logger + + @classmethod + def _detect_level(cls, name: str) -> str: + """Infer level from logger name.""" + name_lower = name.lower() + + if "env" in name_lower: + return cls.LEVEL_ENV + elif "robot" in name_lower and "controller" not in name_lower: + return cls.LEVEL_ROBOT + elif ( + "controller" in name_lower + or "communication" in name_lower + or "socket" in name_lower + ): + return cls.LEVEL_CONTROLLER + elif "model" in name_lower or "wrapper" in name_lower: + return cls.LEVEL_MODEL + else: + return cls.LEVEL_UTILS + + @classmethod + def get_env_logger(cls, name: str = "Environment") -> logging.Logger: + """Environment-level logger.""" + return cls.get_logger(name, cls.LEVEL_ENV) + + @classmethod + def get_robot_logger(cls, name: str = "Robot") -> logging.Logger: + """Robot-level logger.""" + return cls.get_logger(name, cls.LEVEL_ROBOT) + + @classmethod + def get_controller_logger(cls, name: str = "Controller") -> logging.Logger: + """Controller-level logger.""" + return cls.get_logger(name, cls.LEVEL_CONTROLLER) + + @classmethod + def get_model_logger(cls, name: str = "Model") -> logging.Logger: + """Model-level logger.""" + return cls.get_logger(name, cls.LEVEL_MODEL) + + @classmethod + def get_utils_logger(cls, name: str = "Utils") -> logging.Logger: + """Utils-level logger.""" + return cls.get_logger(name, cls.LEVEL_UTILS) + + @classmethod + def set_level(cls, level: str): + """Change log level for all loggers.""" + new_level = getattr(logging, level.upper()) + cls.log_level = new_level + for logger in cls._loggers.values(): + logger.setLevel(new_level) + for handler in logger.handlers: + handler.setLevel(new_level) + + @classmethod + def close_all(cls): + """Close all logger file handles.""" + for logger in cls._loggers.values(): + for handler in logger.handlers[:]: + handler.close() + logger.removeHandler(handler) + cls._loggers.clear() + cls._initialized = False + + +# Convenience helpers +def get_logger(name: str, level: str = None) -> logging.Logger: + """ + Convenience wrapper to get a logger + + Args: + name: Logger name (usually __name__) + level: Optional level (auto-detected if omitted) + + Returns: + Configured logger instance + + Example: + from wall_x._vendor.harrix.serving._wallx_infer.logger import get_logger + logger = get_logger(__name__) # Auto-detect level from name + logger = get_logger(__name__, "ROBOT") # explicit level + """ + return InferLogger.get_logger(name, level) + + +def setup_logger( + log_level: str = "INFO", + log_dir: Optional[str] = None, + console_output: bool = True, + file_output: bool = True, + colorful: bool = True, +): + """ + Convenience wrapper to configure logging + + Args: + log_level: DEBUG, INFO, WARNING, ERROR, CRITICAL + log_dir: Log file directory + console_output: Enable console output + file_output: Enable file output + colorful: Enable colored console output + + Example: + from wall_x._vendor.harrix.serving._wallx_infer.logger import setup_logger + setup_logger(log_level="DEBUG", log_dir="./logs") + """ + InferLogger.setup(log_level, log_dir, console_output, file_output, colorful) diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py b/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py new file mode 100644 index 00000000..85ae6f0f --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py @@ -0,0 +1,764 @@ +import os +import torch +import json +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 +from wall_x._vendor.harrix.utils.train_config import ( + resolve_camera_label, + resolve_max_length, + resolve_state_bins, + resolve_use_state_string_representation, +) +from wall_x.model.core.action.normalizer import Normalizer +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): + self.config = 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() + self.logger.info(f"Model {self.config.checkpoint_path} loaded") + + self.norm_key = self.config.norm_key + self._register_normalizers() + self.logger.info(f"Normalizers {self.norm_key} registered") + + # 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 + + model_type = self.config.train_config["model_type"] + adapter_cls = resolve_adapter(model_type) + ModelClass = adapter_cls.inference_model_class() + self.ModelClass = ModelClass + + self.logger.info(f"Initializing model: {model_type} ({ModelClass.__name__})") + + self.model = ModelClass( + self.config.model_config, + self.processor, + self.tokenizer_mixin, + ) + + # log attention implementation — variant-specific layout dispatched + # via the adapter (qwen2_5 uses model.model._attn_implementation; + # qwen3 / qwen3_5 use model.model.language_model.config._attn_impl). + 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 complete") + self.logger.info("Casting selected params to bfloat16") + self.model.to_bfloat16_for_selected_params() + self.logger.info("bfloat16 cast complete") + + 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 if present (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"❌ 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.", + ) + + state_dict = reshape_compatible_state_dict( + state_dict, self.model.state_dict(), log_fn=self.logger.info + ) + 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, + truncation=True, + return_tensors="pt", + max_length=resolve_max_length(self.config.train_config), + 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): + self.logger.info("Starting flow action generation") + 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, + **model_input, + ) + + self.logger.info("Flow action generation complete") + + 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) + + # Batch inputs in one preprocesser_call + 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("Starting AR action generation") + 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 complete") + + 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("Starting subtask generation") + 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 complete: {subtask}") + return subtask + + def infer_vqa(self, observation, instruction): + self.logger.info("Starting VQA generation") + + 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 complete: {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: # Retry up to 3 times + 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, # v3.1 delta decode + **model_input, + ) + if model_output["predict_action"] is not None: + break + cnt += 1 + self.logger.warning(f"DLLM action failed; retry {cnt}") + self.logger.info("DLLM action generation complete") + + model_output["robot_state_action_data"] = observation["robot_state_action_data"] + model_output["robot_state_action_data"].save_action_data( + model_output["predict_action"] + ) + self.logger.info("Saved DLLM action to robot_state_action_data") + return model_output diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py b/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py new file mode 100644 index 00000000..d86e4e7a --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py @@ -0,0 +1,1183 @@ +import numpy as np +from abc import ABC +import re + +from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig +from wall_x._vendor.harrix.serving._wallx_infer.utils import VehiclePoseHandler, UnifiedTrajectoryProcessor +from wall_x._vendor.harrix.serving._wallx_infer.base_dataclass import ( + RobotStateActionData, + dof_dims, +) +from wall_x._vendor.harrix.serving._wallx_infer.socket_controller import RobotController + +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger + +robot_action_key_mapping = { + "follow_left_ee_cartesian_pos": "follow1_pos[:3]", + "follow_left_ee_rotation": "follow1_pos[3:6]", + "follow_left_gripper": "follow1_pos[6:7]", + "follow_right_ee_cartesian_pos": "follow2_pos[:3]", + "follow_right_ee_rotation": "follow2_pos[3:6]", + "follow_right_gripper": "follow2_pos[6:7]", + "head_actions": "head_pos", + "head_rotation": "head_pos", # match ex001 + "height": "lift", + "velocity_decomposed": "velocity_decomposed", + "velocity_decomposed_odom": "velocity_decomposed_odom", # match ex001 + "follow_left_arm_joint_cur": "follow1_joints_cur", # 1 -> 7 + "follow_right_arm_joint_cur": "follow2_joints_cur", # 1 -> 7 + "follow_left_gripper_cur": "follow1_joints_cur[-1:]", + "follow_right_gripper_cur": "follow2_joints_cur[-1:]", + "follow_left_arm_joint_pos": "follow1_joints", + "follow_right_arm_joint_pos": "follow2_joints", + "follow_left_wrench_ext_local_force": "follow1_end_effort_force", + "follow_left_wrench_ext_local_torque": "follow1_end_effort_torque", + "follow_right_wrench_ext_local_force": "follow2_end_effort_force", + "follow_right_wrench_ext_local_torque": "follow2_end_effort_torque", + "follow_left_wrench_ext_local_force_from_joint": "follow1_end_effort_force_from_joint", + "follow_left_wrench_ext_local_torque_from_joint": "follow1_end_effort_torque_from_joint", + "follow_right_wrench_ext_local_force_from_joint": "follow2_end_effort_force_from_joint", + "follow_right_wrench_ext_local_torque_from_joint": "follow2_end_effort_torque_from_joint", + "follow_left_wrench_ext_world_force": "follow1_wrench_ext_world_force", + "follow_left_wrench_ext_world_torque": "follow1_wrench_ext_world_torque", + "follow_left_wrench_ext_world_force_from_joint": "follow1_wrench_ext_world_force_from_joint", + "follow_left_wrench_ext_world_torque_from_joint": "follow1_wrench_ext_world_torque_from_joint", + "follow_right_wrench_ext_world_force": "follow2_wrench_ext_world_force", + "follow_right_wrench_ext_world_torque": "follow2_wrench_ext_world_torque", + "follow_right_wrench_ext_world_force_from_joint": "follow2_wrench_ext_world_force_from_joint", + "follow_right_wrench_ext_world_torque_from_joint": "follow2_wrench_ext_world_torque_from_joint", + "follow_left_arm_joint_dev": "follow1_joints_dev", + "follow_right_arm_joint_dev": "follow2_joints_dev", +} + + +def _parse_follow_pos(follow_pos) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Parse websocket ``follow{1,2}_pos``: pos3 + euler3 + grip1.""" + arr = np.asarray(follow_pos, dtype=np.float64).reshape(-1) + if arr.shape[0] < 7: + raise ValueError(f"follow_pos expects 7 dims, got shape {arr.shape}") + return arr[:3], arr[3:6], arr[6:7] + + +def ingest_websocket_follow_pos( + state: dict, + config: InferConfig, + robot_state_action_data: RobotStateActionData, +) -> None: + """Convert client ``follow1_pos``/``follow2_pos`` (14D) into train-config proprio.""" + from wall_x._vendor.harrix.utils.train_config import resolve_agent_pos_config + from wall_x._vendor.x2robot_utils import geometry as geom + + agent_cfg = resolve_agent_pos_config(config.train_config) + arm_ws_keys = ( + ("left", "follow1_pos", "follow_left_"), + ("right", "follow2_pos", "follow_right_"), + ) + for _side, ws_key, cfg_prefix in arm_ws_keys: + if ws_key not in state or state[ws_key] is None: + continue + pos, euler, grip = _parse_follow_pos(state[ws_key]) + rot6d = geom.euler_to_matrix_zyx_6d_nb(euler.reshape(1, 3)) + for key in agent_cfg: + if key == "action_padding" or not key.startswith(cfg_prefix): + continue + if "cartesian_pos" in key: + robot_state_action_data.save_state_data_with_key(pos, key) + elif "rotation_6d" in key.lower(): + robot_state_action_data.save_state_data_with_key( + rot6d.reshape(1, -1), key + ) + elif "rotation" in key: + robot_state_action_data.save_state_data_with_key(euler, key) + elif "gripper" in key: + robot_state_action_data.save_state_data_with_key(grip, key) + + for key, dim in agent_cfg.items(): + if key == "action_padding" or key.startswith( + ("follow_left_", "follow_right_") + ): + continue + norm_key = key.replace("follow_", "").replace("master_", "") + if robot_state_action_data.data.get(f"state_{norm_key}") is None: + robot_state_action_data.save_state_data_with_key( + np.zeros(int(dim), dtype=np.float64), key + ) + + +def export_websocket_follow_pos( + robot_state_action_data: RobotStateActionData, +) -> dict[str, list]: + """Convert model ``robot_state_action_data`` back to ``follow1_pos``/``follow2_pos``.""" + left_arm_action = _stack_state_action_series(robot_state_action_data, side="left") + right_arm_action = _stack_state_action_series( + robot_state_action_data, side="right" + ) + return { + "follow1_pos": left_arm_action.tolist(), + "follow2_pos": right_arm_action.tolist(), + } + + +def _zero_state_series(key: str) -> np.ndarray: + """Fallback proprio when websocket state omits an arm (e.g. only follow2_pos).""" + dim = dof_dims.get(key, 1) + return np.zeros((1, dim), dtype=np.float64) + + +def _identity_state_rotation_6d() -> np.ndarray: + from wall_x._vendor.x2robot_utils import geometry as geom + + return geom.euler_to_matrix_zyx_6d_nb(np.zeros((1, 3), dtype=np.float64)) + + +def _as_2d_series(value, name: str) -> np.ndarray: + """Normalize state/action arrays to (T, D) for follow*_pos serialization.""" + arr = np.asarray(value, dtype=np.float64) + if arr.ndim == 0: + arr = arr.reshape(1, 1) + elif arr.ndim == 1: + arr = arr.reshape(1, -1) if arr.size <= 7 else arr.reshape(-1, 1) + elif arr.ndim > 2: + arr = arr.reshape(arr.shape[0], -1) + if arr.ndim != 2: + raise ValueError(f"{name}: expected 2D series after normalize, got {arr.shape}") + return arr + + +def _stack_state_action_series( + robot_state_action_data: RobotStateActionData, + *, + side: str, + pos_key: str = "ee_cartesian_pos", + rot_key: str = "ee_rotation", + grip_key: str = "gripper", +) -> np.ndarray: + """Build (1+horizon, 7) follow arm pose: pos3 + rot3 + grip1.""" + from wall_x._vendor.x2robot_utils import geometry as geom + + data = robot_state_action_data.data + prefix = "left" if side == "left" else "right" + + state_pos = data.get(f"state_{prefix}_{pos_key}") + action_pos = data.get(f"action_{prefix}_{pos_key}") + rel_pos = data.get(f"action_{prefix}_{pos_key}_relative") + if action_pos is None and rel_pos is not None: + base_pos = ( + _as_2d_series(state_pos, f"state_{prefix}_{pos_key}") + if state_pos is not None + else _zero_state_series(f"{prefix}_{pos_key}") + ) + action_pos = base_pos + np.asarray(rel_pos, dtype=np.float64) + if state_pos is None: + state_pos = _zero_state_series(f"{prefix}_{pos_key}") + + state_rot = data.get(f"state_{prefix}_{rot_key}") + action_rot = data.get(f"action_{prefix}_{rot_key}") + + rot6d = data.get(f"action_{prefix}_{rot_key}_6D") + if action_rot is None and rot6d is not None: + action_rot = geom.so3_to_euler_zyx_batch_nb(np.asarray(rot6d, dtype=np.float64)) + if action_rot is None: + rel6d = data.get(f"action_{prefix}_{rot_key}_6D_relative") + state6d = data.get(f"state_{prefix}_{rot_key}_6D") + if state6d is None and state_rot is not None: + state6d = geom.euler_to_matrix_zyx_6d_nb( + _as_2d_series(state_rot, f"state_{prefix}_{rot_key}") + ) + if state6d is None: + state6d = _identity_state_rotation_6d() + if rel6d is not None: + rot6d = geom.compose_state_and_delta_to_abs_6d( + np.asarray(rel6d, dtype=np.float64), + np.asarray(state6d, dtype=np.float64).reshape(-1), + ) + action_rot = geom.so3_to_euler_zyx_batch_nb(rot6d) + + if state_rot is None: + state6d = data.get(f"state_{prefix}_{rot_key}_6D") + if state6d is not None: + state_rot = geom.so3_to_euler_zyx_batch_nb( + np.asarray(state6d, dtype=np.float64) + ) + else: + state_rot = _zero_state_series(f"{prefix}_{rot_key}") + + state_grip = data.get(f"state_{prefix}_{grip_key}") + if state_grip is None: + state_grip = _zero_state_series(f"{prefix}_{grip_key}") + action_grip = data.get(f"action_{prefix}_{grip_key}") + + if action_pos is None or action_rot is None or action_grip is None: + missing = [ + name + for name, value in ( + (f"action_{prefix}_{pos_key}", action_pos), + (f"action_{prefix}_{rot_key}", action_rot), + (f"action_{prefix}_{grip_key}", action_grip), + ) + if value is None + ] + raise ValueError( + f"Cannot serialize {prefix} arm action; missing action fields: {missing}" + ) + + pos = np.concatenate( + [ + _as_2d_series(state_pos, f"state_{prefix}_{pos_key}"), + _as_2d_series(action_pos, f"action_{prefix}_{pos_key}"), + ], + axis=0, + ) + rot = np.concatenate( + [ + _as_2d_series(state_rot, f"state_{prefix}_{rot_key}"), + _as_2d_series(action_rot, f"action_{prefix}_{rot_key}"), + ], + axis=0, + ) + grip = np.concatenate( + [ + _as_2d_series(state_grip, f"state_{prefix}_{grip_key}"), + _as_2d_series(action_grip, f"action_{prefix}_{grip_key}"), + ], + axis=0, + ) + return np.concatenate([pos, rot, grip], axis=1) + + +def _views_to_camera_observation(config: InferConfig, views: dict) -> dict[str, np.ndarray]: + """Map websocket camera keys to model observation keys. + + Only cameras listed in ``config.cam_names`` are required (e.g. LIBERO uses + face_view + right_wrist_view without left_wrist_view). + """ + camera_mappings = { + "face_view": [config.camera_front_key, "face_view", "front_view"], + "left_wrist_view": [ + config.camera_left_key, + "left_wrist_view", + "left_view", + ], + "right_wrist_view": [ + config.camera_right_key, + "right_wrist_view", + "right_view", + ], + } + observation: dict[str, np.ndarray] = {} + for obs_key in config.cam_names: + possible_keys = camera_mappings.get(obs_key, [obs_key]) + for view_key in possible_keys: + if view_key in views and views[view_key] is not None: + view_data = views[view_key] + if ( + isinstance(view_data, np.ndarray) + and view_data.ndim == 4 + and view_data.shape[0] == 1 + ): + observation[obs_key] = view_data[0] + else: + observation[obs_key] = view_data + break + else: + raise KeyError( + f"Missing view for {obs_key!r}, tried keys: {possible_keys}" + ) + return observation + + +class Robot(ABC): + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.config = config + + self.robot_controller = RobotController( + robot_id=robot_id, host=config.robot_host, port=config.robot_port + ) + + self.robot_controller.connect() + + self.is_received = False # Blocking flag + self.current_robot_state_action = None + + self.logger = InferLogger.get_robot_logger("Robot") + + def _get_views_and_state(self): + state = self.robot_controller.recv_action() + views = self.robot_controller.recv_image( + [ + self.config.camera_left_key, + self.config.camera_front_key, + self.config.camera_right_key, + ] + ) + self.is_received = True + return state, views + + def _get_dof_mask(self): + dof_config = self.config.train_config["dof_config"] + total_dof = sum(dof_config.values()) + dof_mask = np.ones((1, self.config.action_horizon, total_dof)) + return dof_mask + + @staticmethod + def _try_compose_arm_joint_pos_with_gripper(state: dict, key: str): + """Build 7D arm joint position from runtime state when possible. + + Why this is needed: + - Joint-control training expects `*_arm_joint_pos` to be 7D + (6 arm joints + 1 gripper). + - Some online robot states expose arm joints as 6D (`follow*_joints`) + and gripper separately in `follow*_pos[6]`. + - If we pass raw 6D joints directly, `save_state_data_with_key` will fail + shape validation against dof_dims=7. + + Behavior: + - Only handles `follow_left_arm_joint_pos` / `follow_right_arm_joint_pos`. + - Returns a composed 7D vector when both required sources exist and shapes + match expected runtime format. + - Returns None for all other keys or incompatible state shapes, so caller + falls back to original mapping logic. + """ + + if key == "follow_left_arm_joint_pos": + joints_key = "follow1_joints" + ee_key = "follow1_pos" + elif key == "follow_right_arm_joint_pos": + joints_key = "follow2_joints" + ee_key = "follow2_pos" + else: + return None + + if joints_key not in state or ee_key not in state: + return None + + joints = np.asarray(state[joints_key]).reshape(-1) + ee = np.asarray(state[ee_key]).reshape(-1) + if joints.shape[0] == 6 and ee.shape[0] >= 7: + return np.concatenate([joints, ee[6:7]], axis=0) + + return None + + def get_observation(self): + state, views = self._get_views_and_state() + + robot_state_action_data = RobotStateActionData(config=self.config) + for key in robot_action_key_mapping.keys(): + state_key_str = robot_action_key_mapping[key] + value = None + # Prefer a robust 7D composition path for arm_joint_pos: + # online state often provides `follow*_joints` as 6D plus gripper + # in `follow*_pos[6]`. Compose first to match model/dof schema. + composed_value = self._try_compose_arm_joint_pos_with_gripper(state, key) + if composed_value is not None: + value = composed_value + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + continue + + match = re.match(r"(\w+)\[(.*)\]", state_key_str) + if match: + base_key = match.group(1) + if base_key in state: + slicing_str = match.group(2) + slice_parts = slicing_str.split(":") + slice_args = [(int(p) if p.strip() else None) for p in slice_parts] + s = slice(*slice_args) + value = state[base_key][s] + else: + if state_key_str in state: + value = state[state_key_str] + if value is not None: + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + + dof_mask = self._get_dof_mask() + robot_state_action_data.dof_mask = dof_mask + self.current_robot_state_action = robot_state_action_data + return { + "robot_state_action_data": robot_state_action_data, + **_views_to_camera_observation(self.config, views), + } + + def go_home(self): + if self.current_robot_state_action is None: + self.get_observation() + + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "left_gripper" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "right_gripper" + ) + + action_dict = {"robot_state_action_data": self.current_robot_state_action} + self.apply_action( + action_dict, + robot_action_interpolate_multiplier=150, + robot_action_start_ratio=0, + robot_action_end_ratio=1, + ) + self.logger.info("Robot returned to home position") + + def _get_left_arm_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + if not self.config.robot_use_joint_angle_control: + return _stack_state_action_series(robot_state_action_data, side="left") + left_arm_joint_pos = np.concatenate( + [ + _as_2d_series( + robot_state_action_data.data["state_left_arm_joint_pos"], + "state_left_arm_joint_pos", + ), + _as_2d_series( + robot_state_action_data.data["action_left_arm_joint_pos"], + "action_left_arm_joint_pos", + ), + ], + axis=0, + ) + return left_arm_joint_pos + + def _get_right_arm_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + if not self.config.robot_use_joint_angle_control: + return _stack_state_action_series(robot_state_action_data, side="right") + right_arm_joint_pos = np.concatenate( + [ + _as_2d_series( + robot_state_action_data.data["state_right_arm_joint_pos"], + "state_right_arm_joint_pos", + ), + _as_2d_series( + robot_state_action_data.data["action_right_arm_joint_pos"], + "action_right_arm_joint_pos", + ), + ], + axis=0, + ) + return right_arm_joint_pos + + def get_serialized_actions( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> dict: + """Build trimmed/interpolated serialized actions from model_output; does not send. + + Used by apply_action and serving policy; returns dict for robot or client. + """ + assert "robot_state_action_data" in input + + left_arm_action = self._get_left_arm_action(input["robot_state_action_data"]) + right_arm_action = self._get_right_arm_action(input["robot_state_action_data"]) + + # Trim action sequence + action_length = len(left_arm_action) + if robot_action_start_ratio is None: + robot_action_start_ratio = self.config.robot_action_start_ratio + if robot_action_end_ratio is None: + robot_action_end_ratio = self.config.robot_action_end_ratio + start_frame = int(robot_action_start_ratio * action_length) + end_frame = int(robot_action_end_ratio * action_length) + + left_arm_action = left_arm_action[start_frame:end_frame] + right_arm_action = right_arm_action[start_frame:end_frame] + + # Interpolate + if robot_action_interpolate_multiplier is None: + robot_action_interpolate_multiplier = ( + self.config.robot_action_interpolate_multiplier + ) + target_length = robot_action_interpolate_multiplier * len(left_arm_action) + left_arm_action, right_arm_action = ( + UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [left_arm_action, right_arm_action], target_length + ) + ) + + if not self.config.robot_use_joint_angle_control: + return { + "follow1_pos": left_arm_action.tolist(), + "follow2_pos": right_arm_action.tolist(), + } + return { + "follow1_joints": left_arm_action.tolist(), + "follow2_joints": right_arm_action.tolist(), + } + + def apply_action( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> None: + serialized_actions = self.get_serialized_actions( + input, + robot_action_interpolate_multiplier=robot_action_interpolate_multiplier, + robot_action_start_ratio=robot_action_start_ratio, + robot_action_end_ratio=robot_action_end_ratio, + ) + self._send_actions(serialized_actions) + + def _send_actions(self, serialized_actions: dict) -> None: + """Send actions to robot controller.""" + self.robot_controller.robot_comm.send_dict(serialized_actions) + self.is_received = False + self.current_robot_state_action = None + + +class DesktopRobot(Robot): + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.is_received = False # Blocking flag + self.current_robot_state_action = None + + self.logger = InferLogger.get_robot_logger("Robot") + + def _get_dof_mask(self): + 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)) + + # For desktop: mask head_actions, height, and velocity_decomposed + mask_keys = ["head_actions", "height", "velocity_decomposed"] + start_idx = 0 + for key, dof_size in dof_config.items(): + if key in mask_keys: + # Zero mask for these DOFs + dof_mask[:, :, start_idx : start_idx + dof_size] = 0 + start_idx += dof_size + + return dof_mask + + +class DesktopRobotVGA(Robot): + def __init__(self, config: InferConfig, robot_id=10) -> None: + super().__init__(config, robot_id) + self.is_received = False # Blocking flag + self.current_robot_state_action = None + + self.logger = InferLogger.get_robot_logger("Robot") + + def _get_dof_mask(self): + dof_config = self.config.train_config["data"]["dof_config"] + total_dof = sum(dof_config.values()) + dof_mask = np.ones((1, self.config.action_horizon, total_dof)) + + # For desktop: mask head_actions, height, and velocity_decomposed + mask_keys = ["head_actions", "height", "velocity_decomposed"] + start_idx = 0 + for key, dof_size in dof_config.items(): + if key in mask_keys: + # Zero mask for these DOFs + dof_mask[:, :, start_idx : start_idx + dof_size] = 0 + start_idx += dof_size + + return dof_mask + + +class DesktopRobotPreprocessor(DesktopRobot): + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.config = config + + self.is_received = False # Blocking flag + self.current_robot_state_action = None + self.logger = InferLogger.get_robot_logger("Robot") + + def _send_actions(self, serialized_actions: dict) -> None: + # bypass + return + + def get_observation(self, state, views): + robot_state_action_data = RobotStateActionData(config=self.config) + if "follow1_pos" in state or "follow2_pos" in state: + ingest_websocket_follow_pos(state, self.config, robot_state_action_data) + for key in robot_action_key_mapping.keys(): + if key.startswith(("follow_left_", "follow_right_")): + continue + state_key_str = robot_action_key_mapping[key] + value = None + match = re.match(r"(\w+)\[(.*)\]", state_key_str) + if match: + base_key = match.group(1) + if base_key in state: + slicing_str = match.group(2) + slice_parts = slicing_str.split(":") + slice_args = [(int(p) if p.strip() else None) for p in slice_parts] + s = slice(*slice_args) + value = state[base_key][s] + else: + if state_key_str in state: + value = state[state_key_str] + if value is not None: + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + + dof_mask = self._get_dof_mask() + robot_state_action_data.dof_mask = dof_mask + self.current_robot_state_action = robot_state_action_data + return { + "robot_state_action_data": robot_state_action_data, + **_views_to_camera_observation(self.config, views), + } + + +class TurtleRobot(Robot): + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.vehicle_pose_handler = VehiclePoseHandler() + self.last_speed = [0, 0, 0] + + def _get_state_from_controller(self): + state = super()._get_state_from_controller() + + self.vehicle_pose_handler.update_pose(state["car_pose"]) + state["velocity_decomposed"] = self.last_speed + + if self.config.turtle_as_desktop: + state["head_pos"] = [0, -1] + state["lift"] = [0.4] + + return state + + def _calculate_car_pose(self, base_velocity_pred): + # Integrate velocity to pose (vectorized) + dt = 1 / 20 + current_pose = ( + self.vehicle_pose_handler.current_pose.copy() + if self.vehicle_pose_handler.current_pose is not None + else np.array([0.0, 0.0, 0.0]) + ) + + # Batch integrate positions + poses_frames = [] + for i in range(len(base_velocity_pred)): + current_pose = self.vehicle_pose_handler.velocity_to_pose( + base_velocity_pred[i, 0], + base_velocity_pred[i, 1], + base_velocity_pred[i, 2], + dt, + current_pose, + ) + poses_frames.append(current_pose.copy()) + + return poses_frames + + def _get_car_velocity( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + car_velocity = robot_state_action_data.data["action_velocity_decomposed"] + self.last_speed = car_velocity[-1, :].copy().tolist() + return car_velocity + + def _get_head_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + head_actions = np.concatenate( + [ + robot_state_action_data.data["state_head_actions"], + robot_state_action_data.data["action_head_actions"], + ], + axis=0, + ) + return head_actions + + def _get_height_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + height = np.concatenate( + [ + robot_state_action_data.data["state_height"], + robot_state_action_data.data["action_height"], + ], + axis=0, + ) + return height + + def go_home(self): + if self.current_robot_state_action is None: + self.get_observation() + + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "left_gripper" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "right_gripper" + ) + + self.current_robot_state_action.save_action_data_with_key( + np.array([[0, -1]]), "head_actions" + ) + self.current_robot_state_action.save_action_data_with_key( + np.array([[0.4]]), "height" + ) + self.current_robot_state_action.save_action_data_with_key( + np.array([self.last_speed]), "velocity_decomposed" + ) + + action_dict = {"robot_state_action_data": self.current_robot_state_action} + self.apply_action( + action_dict, + robot_action_interpolate_multiplier=150, + robot_action_start_ratio=0, + robot_action_end_ratio=1, + ) + self.logger.info("Robot returned to home position") + + def get_serialized_actions( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> dict: + """Build Turtle serialized actions (dual arms + head/lift/car_pose); does not send.""" + assert "robot_state_action_data" in input + + left_arm_action = self._get_left_arm_action(input["robot_state_action_data"]) + right_arm_action = self._get_right_arm_action(input["robot_state_action_data"]) + action_length = len(left_arm_action) + if robot_action_start_ratio is None: + robot_action_start_ratio = self.config.robot_action_start_ratio + if robot_action_end_ratio is None: + robot_action_end_ratio = self.config.robot_action_end_ratio + start_frame = int(robot_action_start_ratio * action_length) + end_frame = int(robot_action_end_ratio * action_length) + + left_arm_action = left_arm_action[start_frame:end_frame] + right_arm_action = right_arm_action[start_frame:end_frame] + if robot_action_interpolate_multiplier is None: + robot_action_interpolate_multiplier = ( + self.config.robot_action_interpolate_multiplier + ) + target_length = robot_action_interpolate_multiplier * action_length + ( + left_arm_action, + right_arm_action, + ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [left_arm_action, right_arm_action], + target_length, + ) + serialized_actions = { + "follow1_pos": left_arm_action.tolist(), + "follow2_pos": right_arm_action.tolist(), + } + + if not self.config.turtle_as_desktop: + head_action = self._get_head_action(input["robot_state_action_data"]) + height_action = self._get_height_action(input["robot_state_action_data"]) + car_velocity_action = self._get_car_velocity( + input["robot_state_action_data"] + ) + head_action = head_action[start_frame:end_frame] + height_action = height_action[start_frame:end_frame] + car_velocity_action = car_velocity_action[start_frame:end_frame] + car_pose_action = np.array(self._calculate_car_pose(car_velocity_action)) + ( + head_action, + height_action, + car_pose_action, + ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [head_action, height_action, car_pose_action], + target_length, + ) + serialized_actions["head_pos"] = head_action.tolist() + serialized_actions["lift"] = height_action.tolist() + serialized_actions["car_pose"] = car_pose_action.tolist() + else: + serialized_actions["head_pos"] = [ + [0, -1] for _ in range(len(left_arm_action)) + ] + serialized_actions["lift"] = [0.4 for _ in range(len(left_arm_action))] + serialized_actions["car_pose"] = [ + [0.0, 0.0, 0.0] for _ in range(len(left_arm_action)) + ] + return serialized_actions + + def apply_action( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> None: + serialized_actions = self.get_serialized_actions( + input, + robot_action_interpolate_multiplier=robot_action_interpolate_multiplier, + robot_action_start_ratio=robot_action_start_ratio, + robot_action_end_ratio=robot_action_end_ratio, + ) + self._send_actions(serialized_actions) + + +class TurtleRobotPreprocessor(TurtleRobot): + """Turtle preprocessor: no real robot; builds observations and get_serialized_actions for serving.""" + + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.config = config + self.is_received = False + self.current_robot_state_action = None + self.logger = InferLogger.get_robot_logger("Robot") + self.vehicle_pose_handler = VehiclePoseHandler() + self.last_speed = [0, 0, 0] + + def _send_actions(self, serialized_actions: dict) -> None: + return + + def get_observation(self, state, views): + robot_state_action_data = RobotStateActionData(config=self.config) + for key in robot_action_key_mapping.keys(): + state_key_str = robot_action_key_mapping[key] + value = None + match = re.match(r"(\w+)\[(.*)\]", state_key_str) + if match: + base_key = match.group(1) + if base_key in state: + slicing_str = match.group(2) + slice_parts = slicing_str.split(":") + slice_args = [(int(p) if p.strip() else None) for p in slice_parts] + s = slice(*slice_args) + value = state[base_key][s] + else: + if state_key_str in state: + value = state[state_key_str] + if value is not None: + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + dof_mask = self._get_dof_mask() + robot_state_action_data.dof_mask = dof_mask + self.current_robot_state_action = robot_state_action_data + return { + "robot_state_action_data": robot_state_action_data, + **_views_to_camera_observation(self.config, views), + } + + +class EX001Robot(Robot): + """EX001 robot: no pose integration; uses velocity_decomposed_odom state for prediction.""" + + def __init__(self, config: InferConfig, robot_id=10) -> None: + super().__init__(config, robot_id) + self.last_speed = [0, 0, 0] + + def _get_velocity_decomposed( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + """Concatenate state and action velocity_decomposed_odom.""" + velocity_decomposed = np.concatenate( + [ + robot_state_action_data.data["state_velocity_decomposed_odom"], + robot_state_action_data.data["action_velocity_decomposed_odom"], + ], + axis=0, + ) + self.last_speed = velocity_decomposed[-1, :].copy().tolist() + return velocity_decomposed + + def _get_head_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + """Concatenate head state/action; supports head_rotation / head_actions.""" + if robot_state_action_data.data.get("state_head_rotation") is not None: + state_key, action_key = "state_head_rotation", "action_head_rotation" + else: + state_key, action_key = "state_head_actions", "action_head_actions" + head_actions = np.concatenate( + [ + robot_state_action_data.data[state_key], + robot_state_action_data.data[action_key], + ], + axis=0, + ) + return head_actions + + def _get_height_action( + self, robot_state_action_data: RobotStateActionData + ) -> np.ndarray: + """Concatenate state and action height sequences.""" + height = np.concatenate( + [ + robot_state_action_data.data["state_height"], + robot_state_action_data.data["action_height"], + ], + axis=0, + ) + return height + + def go_home(self): + """Move to home position.""" + if self.current_robot_state_action is None: + self.get_observation() + + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "left_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "left_gripper" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_cartesian_pos" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 3)), "right_ee_rotation" + ) + self.current_robot_state_action.save_action_data_with_key( + np.zeros((1, 1)), "right_gripper" + ) + + self.current_robot_state_action.save_action_data_with_key( + np.array([[0, -1]]), "head_actions" + ) + self.current_robot_state_action.save_action_data_with_key( + np.array([[0.4]]), "height" + ) + self.current_robot_state_action.save_action_data_with_key( + np.array([self.last_speed]), "velocity_decomposed_odom" + ) + + action_dict = {"robot_state_action_data": self.current_robot_state_action} + self.apply_action( + action_dict, + robot_action_interpolate_multiplier=150, + robot_action_start_ratio=0, + robot_action_end_ratio=1, + ) + self.logger.info("Robot returned to home position") + + def get_serialized_actions( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> dict: + """Build EX001 serialized actions (dual arms + head/lift/velocity_decomposed_odom); does not send. + + Uses predicted velocity_decomposed_odom directly; no pose integration. + """ + assert "robot_state_action_data" in input + + left_arm_action = self._get_left_arm_action(input["robot_state_action_data"]) + right_arm_action = self._get_right_arm_action(input["robot_state_action_data"]) + action_length = len(left_arm_action) + + # action_length is 1 state frame + action_horizon action frames + # Effective action length is action_length - 1 + actual_action_length = action_length - 1 + + if robot_action_start_ratio is None: + robot_action_start_ratio = self.config.robot_action_start_ratio + if robot_action_end_ratio is None: + robot_action_end_ratio = self.config.robot_action_end_ratio + + # Trim indices in action segment (skip frame 0 state) + start_frame = 1 + int(robot_action_start_ratio * actual_action_length) + end_frame = 1 + int(robot_action_end_ratio * actual_action_length) + + left_arm_action = left_arm_action[start_frame:end_frame] + right_arm_action = right_arm_action[start_frame:end_frame] + + if robot_action_interpolate_multiplier is None: + robot_action_interpolate_multiplier = ( + self.config.robot_action_interpolate_multiplier + ) + target_length = robot_action_interpolate_multiplier * len(left_arm_action) + + ( + left_arm_action, + right_arm_action, + ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [left_arm_action, right_arm_action], + target_length, + ) + + serialized_actions = { + "follow1_pos": left_arm_action.tolist(), + "follow2_pos": right_arm_action.tolist(), + } + + # Head, lift, and velocity actions + head_action = self._get_head_action(input["robot_state_action_data"]) + height_action = self._get_height_action(input["robot_state_action_data"]) + velocity_decomposed_action = self._get_velocity_decomposed( + input["robot_state_action_data"] + ) + + # Trim with same start_frame/end_frame + head_action = head_action[start_frame:end_frame] + height_action = height_action[start_frame:end_frame] + velocity_decomposed_action = velocity_decomposed_action[start_frame:end_frame] + + # Interpolate + ( + head_action, + height_action, + velocity_decomposed_action, + ) = UnifiedTrajectoryProcessor.interpolate_trajectory_batch( + [head_action, height_action, velocity_decomposed_action], + target_length, + ) + + serialized_actions["head_pos"] = head_action.tolist() + serialized_actions["lift"] = height_action.tolist() + # Send velocity_decomposed_odom; no pose integration + serialized_actions["velocity_decomposed_odom"] = ( + velocity_decomposed_action.tolist() + ) + return serialized_actions + + def apply_action( + self, + input: dict, + robot_action_interpolate_multiplier=None, + robot_action_start_ratio=None, + robot_action_end_ratio=None, + ) -> None: + serialized_actions = self.get_serialized_actions( + input, + robot_action_interpolate_multiplier=robot_action_interpolate_multiplier, + robot_action_start_ratio=robot_action_start_ratio, + robot_action_end_ratio=robot_action_end_ratio, + ) + self._send_actions(serialized_actions) + + +class EX001RobotPreprocessor(EX001Robot): + """EX001 preprocessor: no real robot; builds observations for serving.""" + + def __init__(self, config: InferConfig, robot_id=10) -> None: + self.config = config + self.is_received = False + self.current_robot_state_action = None + self.logger = InferLogger.get_robot_logger("Robot") + self.last_speed = [0, 0, 0] + + def _send_actions(self, serialized_actions: dict) -> None: + """No-op send; preprocessing only.""" + return + + def get_observation(self, state, views): + """Build observation from external state and views.""" + robot_state_action_data = RobotStateActionData(config=self.config) + for key in robot_action_key_mapping.keys(): + state_key_str = robot_action_key_mapping[key] + value = None + + # Keep the same 7D adaptation behavior as Robot.get_observation(), + # so offline/serving preprocessor path is consistent with runtime path. + composed_value = self._try_compose_arm_joint_pos_with_gripper(state, key) + if composed_value is not None: + value = composed_value + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + continue + + match = re.match(r"(\w+)\[(.*)\]", state_key_str) + if match: + base_key = match.group(1) + if base_key in state: + slicing_str = match.group(2) + slice_parts = slicing_str.split(":") + slice_args = [(int(p) if p.strip() else None) for p in slice_parts] + s = slice(*slice_args) + value = state[base_key][s] + else: + if state_key_str in state: + value = state[state_key_str] + if value is not None: + robot_state_action_data.save_state_data_with_key( + np.asarray(value)[None], key + ) + dof_mask = self._get_dof_mask() + robot_state_action_data.dof_mask = dof_mask + self.current_robot_state_action = robot_state_action_data + + # Build observation; handle key name aliases + observation = { + "robot_state_action_data": robot_state_action_data, + } + + # Try multiple possible view keys + camera_mappings = { + "face_view": [ + self.config.camera_front_key, + "face_view", + "front_view", + ], + "left_wrist_view": [ + self.config.camera_left_key, + "left_wrist_view", + "left_view", + ], + "right_wrist_view": [ + self.config.camera_right_key, + "right_wrist_view", + "right_view", + ], + } + + for obs_key in self.config.cam_names: + possible_view_keys = camera_mappings.get(obs_key, [obs_key]) + found = False + for view_key in possible_view_keys: + if view_key in views and views[view_key] is not None: + view_data = views[view_key] + # If view_data is (1,H,W,3), take [0] + if ( + isinstance(view_data, np.ndarray) + and len(view_data.shape) == 4 + and view_data.shape[0] == 1 + ): + observation[obs_key] = view_data[0] + else: + observation[obs_key] = view_data + found = True + self.logger.info( + f"EX001RobotPreprocessor: using '{view_key}' as '{obs_key}'" + ) + break + if not found: + self.logger.error( + f"EX001RobotPreprocessor: no view for '{obs_key}'; tried keys: {possible_view_keys}" + ) + # Black image placeholder + observation[obs_key] = np.zeros((480, 640, 3), dtype=np.uint8) + + return observation diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py b/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py new file mode 100644 index 00000000..b40eace8 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py @@ -0,0 +1,423 @@ +import socket +import struct +import json +import cv2 +import numpy as np +import time +import select +import errno +import os +import pickle + +import threading + +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger + +_DEFAULT_ROBOT_CONFIG = { + 0: {"host": "127.0.0.1", "action_port": 57749, "keyboard_port": 58849}, + 1: {"host": "127.0.0.1", "action_port": 57750, "keyboard_port": 58850}, + 2: {"host": "127.0.0.1", "action_port": 57751, "keyboard_port": 58851}, + 3: {"host": "127.0.0.1", "action_port": 57761, "keyboard_port": 58861}, +} + + +class RobotRegistry: + def __init__(self): + self.registry = _DEFAULT_ROBOT_CONFIG + self.logger = InferLogger.get_controller_logger("RobotRegistry") + self.logger.debug(f"Initialized robot registry with {len(self.registry)} robots") + + def register_robot(self, robot_id, host, action_port, keyboard_port): + if robot_id in self.registry: + self.logger.warning(f"Robot ID {robot_id} is already registered") + return False + else: + self.registry[robot_id] = { + "host": host, + "action_port": action_port, + "keyboard_port": keyboard_port, + } + self.logger.info( + f"Registered robot ID={robot_id}: host={host}, action_port={action_port}, keyboard_port={keyboard_port}" + ) + return True + + def get_robot_info(self, robot_id): + if robot_id in self.registry: + info = self.registry[robot_id] + self.logger.debug(f"Robot ID={robot_id} info: {info}") + return info + else: + self.logger.error(f"Robot ID={robot_id} not found") + return None + + def exist(self, robot_id): + return robot_id in self.registry + + +class RobotCommunication: + def __init__(self, robot_id, host=None, action_port=None, keyboard_port=None): + self.logger = InferLogger.get_controller_logger("RobotCommunication") + self.logger.debug(f"Initializing robot communication: robot_id={robot_id}") + + self.robot_register = RobotRegistry() + if self.robot_register.exist(robot_id): + self.robot_info = self.robot_register.get_robot_info(robot_id) + self.logger.debug(f"Using registered robot config: {robot_id}") + else: + self.robot_register.register_robot( + robot_id, host, action_port, keyboard_port + ) + self.robot_info = self.robot_register.get_robot_info(robot_id) + self.logger.debug(f"Registering new robot config: {robot_id}") + + self.action_sock = None + self.action_conn = None + + self.keyboard_sock = None + self.keyboard_conn = None + + self.client_socks = [] + + def connect(self): + self.logger.info("Establishing socket connection...") + + self.action_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.action_sock.setblocking(True) + self.action_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + + host = self.robot_info["host"] + port = self.robot_info["action_port"] + self.action_sock.bind((host, port)) + self.action_sock.listen(1) + self.logger.info(f"Listening on action port: {host}:{self.action_sock.getsockname()[1]}") + + action_thread = threading.Thread(target=self.handle_action_client) + + self.logger.debug("Starting connection handler thread...") + action_thread.start() + + action_thread.join() + self.logger.info("Socket connection established") + + def handle_action_client(self): + self.logger.debug("Waiting for client connection...") + self.action_conn, addr = self.action_sock.accept() + self.logger.info(f"Accepted connection from {addr}") + + def recv_image(self, index): + self.logger.debug(f"Receiving image index={index}...") + image_size = struct.unpack(" dict: + if not self.views_path: + raise RuntimeError( + "DummyRobotController requires views_path or " + "WALL_X_DUMMY_VIEWS_PATH." + ) + self.logger.debug(f"Loading view data from pickle: {cam_names}") + with open(self.views_path, "rb") as f: + all_views = pickle.load(f) + self.logger.debug(f"View data loaded: {len(all_views)} views") + return all_views + + def recv_action(self): + if not self.state_path: + raise RuntimeError( + "DummyRobotController requires state_path or " + "WALL_X_DUMMY_STATE_PATH." + ) + self.logger.debug("Loading action data from pickle") + with open(self.state_path, "rb") as f: + all_actions = pickle.load(f) + self.logger.debug(f"Action data loaded: {len(all_actions)} fields") + return all_actions + + +class RobotController: + def __init__( + self, + robot_id: int, + host: str = None, + port: str = None, + max_time_step: int = 10000, + ): + self.logger = InferLogger.get_controller_logger("RobotController") + self.logger.info( + f"Initialized robot controller: robot_id={robot_id}, max_time_step={max_time_step}" + ) + + self.robot_comm = RobotCommunication(robot_id, host, port) + self.max_time_step = max_time_step + self.global_step = 0 + + def connect(self): + self.logger.info("Connecting robot controller...") + self.robot_comm.connect() + self.logger.info("Robot controller connected") + + def close(self): + self.logger.info("Closing robot controller...") + self.robot_comm.close() + self.logger.info("Robot controller closed") + + def prediction(self, views, actions) -> dict: + raise NotImplementedError + + def recv_image( + self, cam_names: list = ["camera_left", "camera_front", "camera_right"] + ) -> dict: + self.logger.debug(f"Receiving images: {cam_names}") + views = {} + for i, name in enumerate(cam_names): + image = np.array(self.robot_comm.recv_image(i)) + views[name] = image[None, :] + self.logger.debug(f"Received {name}: shape={image.shape}") + self.logger.debug(f"All images received: {len(views)} views") + return views + + def recv_action(self): + self.logger.debug("Receiving action data...") + action_data = self.robot_comm.recv_action_data() + return action_data + + def recv_keyboard_input(self): + self.robot_comm.accept_connections() + json_data = self.robot_comm.recv_keyboard_input() + if json_data is not None: + self.logger.debug( + f"Keyboard input: motionStatus={json_data.get('motionStatus')}, armMode={json_data.get('armMode')}" + ) + return json_data["motionStatus"], json_data["armMode"] + + return None, None + + def reset(self): + self.logger.info(f"Reset controller: global_step {self.global_step} -> 0") + self.global_step = 0 + + def record_start(self): + self.logger.info("Starting recording...") + record_signal = {"cmd": "RECORD_START"} + self.robot_comm.send_dict(record_signal) + + def record_continue(self): + self.logger.debug("Continuing recording...") + record_signal = {"cmd": "RECORD_CONTINUE"} + self.robot_comm.send_dict(record_signal) + + def record_start_process(self): + if self.global_step == 0: + self.logger.info("Waiting for START signal...") + action = None + while action != "START": + action, _ = self.recv_keyboard_input() + self.logger.info("Received START signal") + self.record_start() + + self.record_continue() + + def set_zero(self): + self.logger.info("Setting zero position...") + record_signal = {"cmd": "INIT_ZERO", "gripper": [0.0, 0.0]} + self.robot_comm.send_dict(record_signal) + self.reset() + self.logger.info("Zero position set") + + def record_stop(self): + self.logger.info("Stopping recording...") + record_signal = {"cmd": "RECORD_STOP"} + self.robot_comm.send_dict(record_signal) + time.sleep(0.01) + + def recover_from_failure(self): + self.logger.warning("Recovering from failure...") + record_signal = {"cmd": "TO_MASTER_SLAVE"} + self.robot_comm.send_dict(record_signal) + time.sleep(0.01) + + self.logger.info("Waiting for START signal...") + action = None + while action != "START": + action, _ = self.recv_keyboard_input() + + self.logger.info("Recovery complete; resuming recording") + self.record_start() + + def reset_to_zero(self): + self.logger.info("Resetting to zero position...") + self.record_stop() + self.set_zero() + + def to_slave(self): + self.logger.info("Switching to slave mode...") + record_signal = {"cmd": "TO_SLAVE"} + self.robot_comm.send_dict(record_signal) + time.sleep(0.01) + + def run( + self, + record_mode=False, + cam_names: list = ["camera_left", "camera_front", "camera_right"], + ): + + while self.global_step <= self.max_time_step: + if record_mode: + self.record_start_process() + + self.global_step += 1 + + action = self.recv_action() + view = self.recv_image(cam_names) + + pred = self.prediction(view, action) + self.robot_comm.send_dict(pred) + + if record_mode: + action, arm_mode = self.recv_keyboard_input() + + if action is not None and arm_mode is not None: + self.logger.info("action: %s, arm_mode: %s", action, arm_mode) + if action == "STOP" and arm_mode == "ARM_TEST_MODE_MS": + self.record_stop() + self.recover_from_failure() + action, arm_mode = None, None + while action != "STOP" or arm_mode != "ARM_TEST_MODE_S": + action, arm_mode = self.recv_keyboard_input() + self.record_stop() + self.to_slave() + self.reset() + elif arm_mode == "ARM_TEST_MODE_S" and action == "INIT": + self.reset_to_zero() diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/utils.py b/wall_x/_vendor/harrix/serving/_wallx_infer/utils.py new file mode 100644 index 00000000..9b49aebe --- /dev/null +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/utils.py @@ -0,0 +1,280 @@ +import numpy as np +from scipy.signal import savgol_filter +from scipy.spatial.transform import Rotation as R # TODO: convert to numba +from collections import deque +import threading + +from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger + + +class KeyboardThread(threading.Thread): + """ + Simple keyboard listener thread with stop and reset + """ + + def __init__(self): + self.should_reset = False + self.should_stop = False + self.new_instruction_index = None # Stores new instruction index + self.logger = InferLogger.get_utils_logger("KeyboardThread") + + super(KeyboardThread, self).__init__(name="keyboard-thread", daemon=True) + self.show_help() + self.start() + + def run(self): + """Listen for keyboard input.""" + while True: + try: + user_input = input().strip().lower() + + if user_input in ["s", "stop"]: + self.should_stop = not self.should_stop + self.logger.info("[Keyboard] Stop signal sent") + + elif user_input in ["r", "reset"]: + self.logger.info("[Keyboard] Resetting...") + self.should_reset = True + self.logger.info("[Keyboard] Reset signal sent") + + elif user_input.isdigit(): + # Numeric input: switch instruction index + index = int(user_input) + self.new_instruction_index = index + self.logger.info(f"[Keyboard] Switched to instruction index: {index}") + + else: + self.logger.info(f"[Keyboard] Input {user_input}; no action taken.") + + except EOFError: + break + except Exception as e: + self.logger.error(f"[Keyboard] Error: {e}") + + def show_help(self): + self.logger.info( + "[Keyboard] Controls: 's' stop, 'r' reset, digit switches instruction index" + ) + + +# Arm trajectory parameters +ARM_MAX_VELOCITY = 0.02 +ARM_EXECUTION_HZ = 20 +ARM_MIN_EXECUTION_TIME = 5.0 +ARM_MAX_EXECUTION_TIME = 15.0 + + +class UnifiedTrajectoryProcessor: + """Unified trajectory processor.""" + + @staticmethod + def interpolate_trajectory_batch(trajectories, target_length, smooth=True): + """ + Interpolate multiple trajectories to a common length + Args: + trajectories: list of np.ndarray, each shape (N, D) + target_length: int, target sequence length + smooth: bool, apply smoothing + Returns: + list of np.ndarray, interpolated trajectories + """ + if not trajectories: + return [] + + results = [] + for traj in trajectories: + if len(traj) == 0: + results.append(np.zeros((target_length, traj.shape[1]))) + continue + + if len(traj) == target_length: + results.append(traj) + continue + + # Vectorized interpolation + original_indices = np.linspace(0, len(traj) - 1, len(traj)) + target_indices = np.linspace(0, len(traj) - 1, target_length) + + # Handle different data types + if traj.shape[1] == 7: # Arm data [x,y,z,rx,ry,rz,gripper] + interpolated = UnifiedTrajectoryProcessor._interpolate_arm_trajectory( + traj, original_indices, target_indices, target_length + ) + else: # Other data (height, current, etc.) + interpolated = np.zeros((target_length, traj.shape[1])) + for i in range(traj.shape[1]): + interpolated[:, i] = np.interp( + target_indices, original_indices, traj[:, i] + ) + + # Smoothing + if smooth and len(interpolated) >= 5: + interpolated = UnifiedTrajectoryProcessor._smooth_trajectory( + interpolated + ) + + results.append(interpolated) + + return results + + @staticmethod + def _interpolate_arm_trajectory( + traj, original_indices, target_indices, target_length + ): + """Optimized arm trajectory interpolation.""" + interpolated = np.zeros((target_length, 7)) + + # Vectorized interpolation for position and gripper + for i in [0, 1, 2, 6]: # x, y, z, gripper + interpolated[:, i] = np.interp(target_indices, original_indices, traj[:, i]) + + # Quaternion interpolation (vectorized) + quaternions = R.from_euler("xyz", traj[:, 3:6]).as_quat() + interpolated_quats = np.zeros((target_length, 4)) + for i in range(4): + interpolated_quats[:, i] = np.interp( + target_indices, original_indices, quaternions[:, i] + ) + + # Batch normalize + norms = np.linalg.norm(interpolated_quats, axis=1, keepdims=True) + interpolated_quats = interpolated_quats / norms + + # Batch convert back to Euler angles + interpolated[:, 3:6] = R.from_quat(interpolated_quats).as_euler("xyz") + + return interpolated + + @staticmethod + def _interpolate_position_trajectory( + traj, original_indices, target_indices, target_length + ): + """Optimized position trajectory interpolation.""" + interpolated = np.zeros((target_length, 3)) + for i in range(3): + interpolated[:, i] = np.interp(target_indices, original_indices, traj[:, i]) + return interpolated + + @staticmethod + def _smooth_trajectory(trajectory): + """Vectorized trajectory smoothing.""" + if len(trajectory) < 5: + return trajectory + + try: + # Smooth all dimensions + smoothed = np.zeros_like(trajectory) + for dim in range(trajectory.shape[1]): + smoothed[:, dim] = savgol_filter( + trajectory[:, dim], + min( + 5, + ( + len(trajectory) + if len(trajectory) % 2 == 1 + else len(trajectory) - 1 + ), + ), + 3, + mode="nearest", + ) + return smoothed + except Exception: + return trajectory + + @staticmethod + def calculate_optimal_trajectory_length(left_traj, right_traj): + """Compute optimal trajectory length.""" + + # Vectorized distance computation + def calc_distance(traj): + if len(traj) < 2: + return 0.0 + pos_diff = traj[1:, :3] - traj[:-1, :3] + return np.sum(np.linalg.norm(pos_diff, axis=1)) + + distances = [calc_distance(left_traj), calc_distance(right_traj)] + max_distance = max(distances) + + if max_distance > 1e-6: + execution_time = np.clip( + max_distance / ARM_MAX_VELOCITY, + ARM_MIN_EXECUTION_TIME, + ARM_MAX_EXECUTION_TIME, + ) + else: + execution_time = ARM_MIN_EXECUTION_TIME + + return max(int(execution_time * ARM_EXECUTION_HZ), len(left_traj)) + + +class VehiclePoseHandler: + """Vehicle pose and velocity utilities.""" + + def __init__(self): + self.current_pose = None + self.previous_pose = None + self.pose_history = deque(maxlen=10) + self.logger = InferLogger.get_utils_logger("VehiclePoseHandler") + + def update_pose(self, new_pose): + """Update vehicle pose.""" + if new_pose is not None: + self.previous_pose = self.current_pose + self.current_pose = np.array(new_pose) + self.pose_history.append(self.current_pose.copy()) + self.logger.info("current_pose %s", self.current_pose) + return self.current_pose + + def velocity_to_pose(self, vx_body, vy_body, vyaw, dt, start_pose=None): + """Integrate body-frame velocity to global pose.""" + if start_pose is None: + if self.current_pose is not None: + start_pose = self.current_pose.copy() + else: + start_pose = np.array([0.0, 0.0, 0.0]) + + x, y, theta = start_pose + + # Body-frame velocity to global displacement + cos_theta = np.cos(theta) + sin_theta = np.sin(theta) + + # Transform: body frame -> global frame + dx_global = (vx_body * cos_theta - vy_body * sin_theta) * dt + dy_global = (vx_body * sin_theta + vy_body * cos_theta) * dt + dtheta = vyaw * dt + + # Compute new pose + x_new = x + dx_global + y_new = y + dy_global + theta_new = theta + dtheta + + # Wrap angle to [-pi, pi] + theta_new = (theta_new + np.pi) % (2 * np.pi) - np.pi + + return np.array([x_new, y_new, theta_new]) + + def compute_body_velocities_from_poses( + self, current_pose, previous_pose, dt=1 / 20 + ): + """Compute body-frame velocity from pose delta.""" + if current_pose is None or previous_pose is None: + return np.array([0.0, 0.0, 0.0]) + + # Displacement in global frame + dx_global = current_pose[0] - previous_pose[0] + dy_global = current_pose[1] - previous_pose[1] + dtheta = current_pose[2] - previous_pose[2] + + # Use previous frame angle for transform + theta = previous_pose[2] + cos_theta = np.cos(theta) + sin_theta = np.sin(theta) + + # Global displacement to body-frame velocity + vx_body = (dx_global * cos_theta + dy_global * sin_theta) / dt + vy_body = (-dx_global * sin_theta + dy_global * cos_theta) / dt + vyaw = dtheta / dt + + return np.array([vx_body, vy_body, vyaw]) diff --git a/wall_x/_vendor/harrix/serving/client.py b/wall_x/_vendor/harrix/serving/client.py new file mode 100644 index 00000000..1dfdbaf4 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/client.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python3 +""" +Example client for Wall-X model server with sync support. + +This script demonstrates how to connect to a Wall-X server and request +action predictions from observations in both sync and async contexts. +""" + +import asyncio +import logging +from typing import Dict +import threading +import time +import numpy as np +import torch +import os + + +try: + import msgpack + import msgpack_numpy as m + + m.patch() +except ImportError: + print("Please install msgpack-numpy: pip install msgpack-numpy") + exit(1) + +try: + import websockets +except ImportError: + print("Please install websockets: pip install websockets") + exit(1) + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +class RobotClient: + """Client for connecting to Wall-X model server.""" + + def __init__(self, config_path: str, uri: str = "ws://localhost:8000"): + """Initialize client. + + Args: + uri: WebSocket URI of the server (e.g., ws://localhost:8000) + """ + self.uri = uri + self.websocket = None + self.metadata = None + self._loop = None + self._thread = None + + async def connect(self): + """Connect to the server and receive metadata.""" + logger.info(f"Connecting to {self.uri}...") + self.websocket = await websockets.connect( + self.uri, + ping_interval=None, + ping_timeout=None, + max_size=None, + ) + + self.metadata = msgpack.unpackb(await self.websocket.recv()) + logger.info(f"Connected! Server metadata: {self.metadata}") + + async def predict(self, obs: Dict) -> Dict: + """Get action prediction from observation. + + Args: + obs: Observation dictionary containing: + - 'image': Image array (H, W, C) + - 'prompt': Optional text prompt + - 'state': Optional robot state + + Returns: + Dictionary with: + - 'action': Predicted action array + - 'server_timing': Timing information + """ + if self.websocket is None: + raise RuntimeError("Not connected. Call connect() first.") + + ziped = msgpack.packb(obs) + size_bytes = len(ziped) + print(f"{size_bytes / 1024:.2f} KB") + print(f"{size_bytes / 1024 / 1024:.2f} MB") + await self.websocket.send(ziped) + response = msgpack.unpackb(await self.websocket.recv()) + return response + + async def close(self): + """Close the connection.""" + if self.websocket: + await self.websocket.close() + logger.info("Connection closed") + + async def reset(self): + """Reset the policy (if supported).""" + pass + + # ============ Synchronous methods (using independent thread event loop) ============ + + def _start_background_loop(self): + """Start event loop in background thread.""" + self._loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + def _ensure_loop(self): + """Ensure background event loop is running.""" + if self._loop is None or not self._loop.is_running(): + self._thread = threading.Thread( + target=self._start_background_loop, daemon=True + ) + self._thread.start() + # Wait for loop to start + import time + + while self._loop is None: + time.sleep(0.01) + + def _run_async(self, coro): + """Run coroutine in background event loop.""" + self._ensure_loop() + future = asyncio.run_coroutine_threadsafe(coro, self._loop) + return future.result() + + def connect_sync(self): + """Synchronously connect to server.""" + return self._run_async(self.connect()) + + def predict_sync(self, obs: Dict) -> Dict: + """Synchronous prediction method. + + Args: + obs: Observation dictionary + + Returns: + Prediction result dictionary + """ + return self._run_async(self.predict(obs)) + + def close_sync(self): + """Synchronously close connection.""" + result = self._run_async(self.close()) + # Stop event loop + if self._loop: + self._loop.call_soon_threadsafe(self._loop.stop) + return result + + +# ============ Synchronous version of main function ============ + + +def main_sync(args): + """Synchronous version of main function.""" + + # Create client and connect + client = RobotClient(args.config_path, uri=args.uri) + client.connect_sync() + import pickle + + with open("./example_data", "rb") as f: + all_views = pickle.load(f) + + with open("./example_data", "rb") as f: + all_actions = pickle.load(f) + print(all_actions.keys()) + print(all_views.keys()) + + pos = { + "follow1_pos": all_actions["follow1_pos"], + "follow2_pos": all_actions["follow2_pos"], + } + use_base64 = False + if use_base64: + import cv2 + import base64 + + for k, v in all_views.items(): + v = v.squeeze(0) + encoded_img = cv2.imencode(".jpg", v)[1] + print(encoded_img.shape) + + base64_str = base64.b64encode(encoded_img).decode("utf-8") + all_views[k] = base64_str + obs = { + "state": pos, + "views": all_views, + "instruction": "Pick up the cup.", + } + else: + path = "./example_data" + follow_pos_1 = torch.load(os.path.join(path, "state_follow1_pos.pt")).numpy() + follow_pos_2 = torch.load(os.path.join(path, "state_follow2_pos.pt")).numpy() + camera_front = torch.load(os.path.join(path, "views_camera_front.pt")).numpy() + camera_left = torch.load(os.path.join(path, "views_camera_left.pt")).numpy() + camera_right = torch.load(os.path.join(path, "views_camera_right.pt")).numpy() + obs = { + "state": { + "follow1_pos": follow_pos_1, + "follow2_pos": follow_pos_2, + }, + "views": { + "camera_front": camera_front, + "camera_left": camera_left, + "camera_right": camera_right, + }, + "instruction": "Pick up the cup.", + } + for _ in range(1): + start_time = time.perf_counter() + response = client.predict_sync(obs) + print(response["action"].shape) + print(response["action"]) + end_time = time.perf_counter() + print(f"Time taken: {(end_time - start_time) * 1000} ms") + action_label = torch.load( + "./example_data" + ).numpy() + np.testing.assert_array_equal(response["action"], action_label) + print("compare success") + client.close_sync() + + +if __name__ == "__main__": + """Asynchronous version of main function.""" + import argparse + + parser = argparse.ArgumentParser(description="Wall-X client examples") + parser.add_argument( + "--example", + choices=["single", "multiple", "benchmark"], + default="single", + help="Example to run", + ) + parser.add_argument( + "--uri", + default="ws://localhost:8000", + help="Server URI", + ) + parser.add_argument( + "--pred_horizon", type=int, default=32, help="Prediction horizon" + ) + parser.add_argument("--action_dim", type=int, default=7, help="Action dimension") + parser.add_argument( + "--config_path", + default="./example_data", + help="Train config path", + ) + parser.add_argument( + "--save_dir", + default="./example_data", + help="Save directory", + ) + args = parser.parse_args() + + # Synchronous mode + main_sync(args) + + # Asynchronous mode + # asyncio.run(main(args)) diff --git a/wall_x/_vendor/harrix/serving/launch_serving.py b/wall_x/_vendor/harrix/serving/launch_serving.py new file mode 100644 index 00000000..47a3420a --- /dev/null +++ b/wall_x/_vendor/harrix/serving/launch_serving.py @@ -0,0 +1,281 @@ +#!/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. + +Based on the OpenPI serve_policy.py script structure. +""" + +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): + from wall_x._vendor.harrix.serving.policy.wall_x_policy 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) + + + + +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 + + # ── 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'." + ) + policy = get_wallx_policy( + model_config, args.image_passing_mode, args.serialize_actions + ) + 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}") + + # 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 + + # Get network info + hostname = socket.gethostname() + try: + local_ip = socket.gethostbyname(hostname) + except Exception: + local_ip = "unknown" + + logger.info(f"Server hostname: {hostname}") + logger.info(f"Server IP: {local_ip}") + logger.info(f"Server will be available at: ws://{args.host}:{args.port}") + logger.info(f"Health check endpoint: http://{args.host}:{args.port}/healthz") + + batching_str = ( + f"batch_size={args.max_batch_size}, wait={args.max_wait_time_ms}ms" + if args.max_batch_size + else "disabled" + ) + logger.info(f"Batching: {batching_str}") + + # Create and start server + server = WebsocketPolicyServer( + policy=policy, + host=args.host, + port=args.port, + metadata=policy_metadata, + max_batch_size=args.max_batch_size, + max_wait_time_ms=args.max_wait_time_ms, + max_queue_size=args.max_queue_size, + timeout_ms=args.timeout_ms, + ) + + logger.info("Starting server...") + try: + server.serve_forever() + except KeyboardInterrupt: + logger.info("Server stopped by user") + except Exception as e: + logger.error(f"Server error: {e}") + sys.exit(1) + + +if __name__ == "__main__": + main(tyro.cli(Args)) diff --git a/wall_x/_vendor/harrix/serving/policy/__init__.py b/wall_x/_vendor/harrix/serving/policy/__init__.py new file mode 100644 index 00000000..e59a287b --- /dev/null +++ b/wall_x/_vendor/harrix/serving/policy/__init__.py @@ -0,0 +1,5 @@ +"""Wall-X policy for harrix websocket serving.""" + +from .wall_x_policy import WallXPolicy + +__all__ = ["WallXPolicy"] diff --git a/wall_x/_vendor/harrix/serving/policy/_smoothing.py b/wall_x/_vendor/harrix/serving/policy/_smoothing.py new file mode 100644 index 00000000..9b6c127d --- /dev/null +++ b/wall_x/_vendor/harrix/serving/policy/_smoothing.py @@ -0,0 +1,104 @@ +"""Test-time Laplacian smoothing for serving model outputs. + +Mirrors the open-loop implementation in +`wallx_bus2604/wall-x/run_scripts/infer_openloop.py` +(_laplacian_smooth @ 203-209, smooth_action / smooth_gripper fields @ 186-187, +application @ 458-463). Invoked by WallXPolicy before downstream 6D->euler +conversion so that smoothed rotations propagate through serialization. +""" + +from __future__ import annotations + +from typing import List, Optional, Sequence + +import numpy as np +import torch + +from wall_x._vendor.harrix.serving._wallx_infer.base_dataclass import dof_dims + + +def _laplacian_smooth(a: np.ndarray, lam: float = 1.0, iters: int = 30) -> np.ndarray: + """Iterative Laplacian smoothing along axis 0; endpoints pinned.""" + a = a.copy() + orig = a.copy() + for _ in range(iters): + a[1:-1] = (orig[1:-1] + lam * (a[:-2] + a[2:])) / (1 + 2 * lam) + return a + + +def _gripper_column_indices( + predict_action_keys: Sequence[str], + action_padding_dof: Optional[int] = None, +) -> List[int]: + """Column indices in the flat (T, D) layout that correspond to gripper dofs.""" + cols: List[int] = [] + dof_start = 0 + for key in predict_action_keys: + if key == "action_padding": + dof_start += action_padding_dof or 0 + continue + short = key.replace("follow_", "").replace("master_", "") + width = dof_dims[short] + if "gripper" in short: + cols.extend(range(dof_start, dof_start + width)) + dof_start += width + return cols + + +_LAZY_ACTION_KEYS = ( + "action_left_ee_cartesian_pos", + "action_right_ee_cartesian_pos", + "action_left_ee_rotation", + "action_right_ee_rotation", + "action_left_ee_rotation_6D", + "action_right_ee_rotation_6D", +) + + +def apply_smoothing( + model_output: dict, + smooth_action: bool, + smooth_gripper: bool, + predict_action_keys: Sequence[str], + action_padding_dof: Optional[int] = None, + action_dim: Optional[int] = None, +) -> None: + """Smooth `model_output['predict_action']` in place and refresh per-arm keys.""" + if not smooth_action: + return + pa = model_output.get("predict_action") + if pa is None: + return + + was_tensor = isinstance(pa, torch.Tensor) + arr = pa.detach().cpu().numpy() if was_tensor else np.asarray(pa) + + orig_ndim = arr.ndim + if orig_ndim == 3: + if arr.shape[0] != 1: + return + arr = arr[0] + if arr.ndim != 2 or arr.shape[0] < 3: + return + if action_dim is not None and arr.shape[-1] != action_dim: + return + + orig = arr.copy() + smoothed = _laplacian_smooth(arr) + + if not smooth_gripper: + for c in _gripper_column_indices(predict_action_keys, action_padding_dof): + if c < smoothed.shape[-1]: + smoothed[:, c] = orig[:, c] + + out = smoothed[None] if orig_ndim == 3 else smoothed + if was_tensor: + out = torch.from_numpy(out).to(device=pa.device, dtype=pa.dtype) + model_output["predict_action"] = out + + rsd = model_output.get("robot_state_action_data") + if rsd is not None: + for k in _LAZY_ACTION_KEYS: + if k in rsd.data: + rsd.data[k] = None + rsd.save_action_data(out) diff --git a/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py b/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py new file mode 100644 index 00000000..73984592 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py @@ -0,0 +1,328 @@ +import base64 +import logging +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 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 + +logger = logging.getLogger(__name__) + +# Inference modes aligned with model_wrapper (extensible for VQA/batch) +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, + ): + """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.model_wrapper = WallxModelWrapper(config) + self.robot_preprocessor = self._register_robot_preprocessor() + 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}") + + 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), + } + 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 _run_action_infer( + self, observation: Dict, instruction: str, mode: str + ) -> Dict[str, Any]: + """Run model_wrapper action inference by mode; returns model_output with robot_state_action_data. + + Modes: flow | ar | flow_with_subtask. VQA can be added here later. + """ + if mode == INFER_MODE_FLOW: + return self.model_wrapper.infer_flow_action(observation, instruction) + 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: + """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(): + img_bytes = base64.b64decode(v) + img_array = np.frombuffer(img_bytes, dtype=np.uint8) + decoded_img = cv2.imdecode(img_array, cv2.IMREAD_COLOR) + # decoded_img = cv2.cvtColor(decoded_img, cv2.COLOR_BGR2RGB) # modify + views[k] = np.expand_dims(decoded_img, axis=0) + + with ScopeTimer("get_observation"): + observation = self.robot_preprocessor.get_observation(state, views) + + mode = obs.get("infer_mode", self.default_infer_mode) + with ScopeTimer(f"infer_{mode}"): + model_output = self._run_action_infer(observation, instruction, mode) + + 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"] + ) + + # 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(): + return self._pack_action_chunk_response(model_output) + + return self.robot_preprocessor.get_serialized_actions( + model_output, robot_action_interpolate_multiplier=1 + ) # Client interpolates over websocket + + # ── 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": + import cv2 + + for k, v in views.items(): + img_bytes = base64.b64decode(v) + img_array = np.frombuffer(img_bytes, dtype=np.uint8) + decoded_img = cv2.imdecode(img_array, cv2.IMREAD_COLOR) + decoded_img = cv2.cvtColor(decoded_img, cv2.COLOR_BGR2RGB) + views[k] = np.expand_dims(decoded_img, axis=0) + else: + # raw numpy: ensure (1, H, W, C) + for k, img in views.items(): + if isinstance(img, np.ndarray) and img.ndim == 3: + views[k] = np.expand_dims(img, axis=0) + + observation = self.robot_preprocessor.get_observation(state, views) + return observation, instruction + + def infer_batch( + self, obs_list: List[Dict[str, Any]], skip_serialize: bool = False + ) -> List[Dict[str, Any]]: + """Perform batch inference. + + Args: + obs_list: List of observations, each containing: + - "views": dict of camera images + - "state": robot state dict or array + - "instruction": text instruction + + Returns: + List of action dicts. + """ + batch_size = len(obs_list) + logger.info(f"WallXPolicy.infer_batch: processing {batch_size} observations") + + try: + observations = [] + instructions = [] + for obs in obs_list: + observation, instruction = self._preprocess_obs(obs) + observations.append(observation) + instructions.append(instruction) + + mode = obs_list[0].get("infer_mode", self.default_infer_mode) + + with torch.no_grad(): + if mode == INFER_MODE_FLOW: + model_outputs = self.model_wrapper.infer_flow_action_batch( + observations, instructions + ) + else: + # AR / flow_with_subtask: fall back to per-sample inference + model_outputs = [] + for obs_dict, instruction in zip(observations, instructions): + output = self._run_action_infer(obs_dict, instruction, mode) + model_outputs.append(output) + + if skip_serialize: + return [{}] * len(model_outputs) + + results = [] + for model_output in model_outputs: + self._apply_smoothing(model_output) + if self.serialize_actions and not self._is_single_arm_right_only(): + action = self.robot_preprocessor.get_serialized_actions( + model_output, robot_action_interpolate_multiplier=1 + ) + results.append(action) + else: + results.append(self._pack_action_chunk_response(model_output)) + + return results + + except Exception as e: + logger.error(f"Batch inference failed: {e}", exc_info=True) + return [{"action": {}, "error": str(e)} for _ in obs_list] + + @property + def metadata(self) -> Dict[str, Any]: + return {"batch_enabled": True, "model": "wall-x"} diff --git a/wall_x/_vendor/harrix/serving/scheduler.py b/wall_x/_vendor/harrix/serving/scheduler.py new file mode 100644 index 00000000..d114f022 --- /dev/null +++ b/wall_x/_vendor/harrix/serving/scheduler.py @@ -0,0 +1,156 @@ +import asyncio +import logging +import time +import uuid +from collections import deque +from dataclasses import dataclass +from typing import Any, Dict, List + +logger = logging.getLogger(__name__) + + +@dataclass +class Request: + request_id: str + obs: Dict[str, Any] + future: asyncio.Future + timestamp: float + + +class RequestScheduler: + def __init__( + self, + policy, + max_batch_size: int = 8, + max_wait_time_ms: float = 100, + max_queue_size: int = 128, + timeout_ms: float = 5000, + ): + self.policy = policy + self.max_batch_size = max_batch_size + self.max_wait_time = max_wait_time_ms / 1000.0 + self.max_queue_size = max_queue_size + self.timeout = timeout_ms / 1000.0 + + self.queue = deque() + self.queue_lock = asyncio.Lock() + self.queue_not_empty = asyncio.Condition(self.queue_lock) + + self.running = False + self.batch_task = None + + async def start(self): + self.running = True + self.batch_task = asyncio.create_task(self._batch_loop()) + logger.info(f"RequestScheduler started: max_batch={self.max_batch_size}, max_wait={self.max_wait_time*1000}ms") + + async def stop(self): + self.running = False + # Process all remaining requests before shutting down + async with self.queue_lock: + remaining = len(self.queue) + if remaining: + logger.info(f"Graceful shutdown: processing {remaining} remaining request(s)") + while True: + async with self.queue_lock: + if len(self.queue) == 0: + break + batch = [] + while self.queue and len(batch) < self.max_batch_size: + batch.append(self.queue.popleft()) + if batch: + await self._process_batch(batch) + # Now stop the batch loop + async with self.queue_lock: + self.queue_not_empty.notify_all() + if self.batch_task: + await self.batch_task + logger.info("RequestScheduler stopped") + + async def add_request(self, obs: Dict[str, Any]) -> Dict[str, Any]: + request_id = str(uuid.uuid4()) + future = asyncio.Future() + request = Request( + request_id=request_id, + obs=obs, + future=future, + timestamp=time.monotonic() + ) + + async with self.queue_lock: + if len(self.queue) >= self.max_queue_size: + raise RuntimeError(f"Queue full: {len(self.queue)}/{self.max_queue_size}") + self.queue.append(request) + self.queue_not_empty.notify() + + try: + result = await asyncio.wait_for(future, timeout=self.timeout) + return result + except asyncio.TimeoutError: + logger.error(f"Request {request_id} timeout after {self.timeout}s") + raise + + async def _batch_loop(self): + while self.running: + try: + batch = await self._collect_batch() + if batch: + await self._process_batch(batch) + except Exception as e: + logger.error(f"Unexpected error in batch loop: {e}", exc_info=True) + + async def _collect_batch(self) -> List[Request]: + async with self.queue_not_empty: + while self.running and len(self.queue) == 0: + await self.queue_not_empty.wait() + + if not self.running: + return [] + + batch = [] + deadline = time.monotonic() + self.max_wait_time + + while len(batch) < self.max_batch_size: + if len(self.queue) > 0: + batch.append(self.queue.popleft()) + + if len(batch) >= self.max_batch_size: + break + + if len(self.queue) == 0: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + await asyncio.wait_for( + self.queue_not_empty.wait(), + timeout=remaining + ) + except asyncio.TimeoutError: + break + + return batch + + async def _process_batch(self, batch: List[Request]): + if not batch: + return + + start_time = time.monotonic() + logger.info(f"Processing batch of {len(batch)} requests") + + try: + obs_list = [req.obs for req in batch] + results = await asyncio.to_thread(self.policy.infer_batch, obs_list) + + for req, result in zip(batch, results): + if not req.future.done(): + req.future.set_result(result) + + infer_time = time.monotonic() - start_time + logger.info(f"Batch processed in {infer_time*1000:.1f}ms, throughput: {len(batch)/infer_time:.1f} req/s") + + except Exception as e: + logger.error(f"Batch processing failed: {e}", exc_info=True) + for req in batch: + if not req.future.done(): + req.future.set_exception(e) diff --git a/wall_x/serving/websocket_policy_server.py b/wall_x/_vendor/harrix/serving/websocket_policy_server.py similarity index 57% rename from wall_x/serving/websocket_policy_server.py rename to wall_x/_vendor/harrix/serving/websocket_policy_server.py index 719a1104..89aadc93 100644 --- a/wall_x/serving/websocket_policy_server.py +++ b/wall_x/_vendor/harrix/serving/websocket_policy_server.py @@ -3,7 +3,7 @@ import logging import time import traceback -from typing import Any, Dict, Optional +from typing import Dict, Any, List try: import msgpack @@ -29,6 +29,10 @@ def infer(self, obs: Dict) -> Dict: """Infer actions from observations.""" raise NotImplementedError + def infer_batch(self, obs_list: List[Dict]) -> List[Dict]: + """Batch inference. Default implementation calls infer() per sample.""" + return [self.infer(obs) for obs in obs_list] + def reset(self) -> None: """Reset the policy to its initial state.""" pass @@ -45,8 +49,11 @@ class WebsocketPolicyServer: Implements a websocket server that: 1. Sends policy metadata on connection 2. Receives observations - 3. Returns predicted actions + 3. Returns predicted actions (single or batched) 4. Tracks timing information + + When batching parameters are provided, dynamic batching is enabled: + requests from multiple clients are queued and processed in batches. """ def __init__( @@ -54,30 +61,62 @@ def __init__( policy: BasePolicy, host: str = "0.0.0.0", port: int = 8000, - metadata: Optional[Dict] = None, + metadata: Dict | None = None, + # Dynamic batching parameters (None = disabled) + max_batch_size: int | None = None, + max_wait_time_ms: float | None = None, + max_queue_size: int = 100, + timeout_ms: float = 30000, ) -> None: self._policy = policy self._host = host self._port = port self._metadata = metadata or {} + + # Dynamic batching + self._scheduler = None + if max_batch_size is not None: + from .scheduler import RequestScheduler + + self._scheduler = RequestScheduler( + policy=policy, + max_batch_size=max_batch_size, + max_wait_time_ms=max_wait_time_ms if max_wait_time_ms is not None else 0, + max_queue_size=max_queue_size, + timeout_ms=timeout_ms, + ) + logging.getLogger("websockets.server").setLevel(logging.INFO) + @property + def batching_enabled(self) -> bool: + return self._scheduler is not None + def serve_forever(self) -> None: asyncio.run(self.run()) async def run(self): - async with _server.serve( - self._handler, - self._host, - self._port, - compression=None, - max_size=None, - ping_interval=None, # Disable automatic ping for long-running inference - ping_timeout=None, # Disable ping timeout - process_request=_health_check, - ) as server: - logger.info(f"Server started on {self._host}:{self._port}") - await server.serve_forever() + # Start the scheduler if batching is enabled + if self._scheduler is not None: + await self._scheduler.start() + + try: + async with _server.serve( + self._handler, + self._host, + self._port, + compression=None, + max_size=None, + ping_interval=None, + ping_timeout=None, + process_request=_health_check, + ) as server: + mode_str = "batched" if self.batching_enabled else "single" + logger.info(f"Server started on {self._host}:{self._port} (mode={mode_str})") + await server.serve_forever() + finally: + if self._scheduler is not None: + await self._scheduler.stop() async def _handler(self, websocket: _server.ServerConnection): logger.info(f"Connection from {websocket.remote_address} opened") @@ -99,7 +138,14 @@ async def _handler(self, websocket: _server.ServerConnection): obs = msgpack.unpackb(await websocket.recv()) infer_time = time.monotonic() - action = self._policy.infer(obs) + + if self._scheduler is not None: + # Dynamic batching path + action = await self._scheduler.add_request(obs) + else: + # Single inference path + action = await asyncio.to_thread(self._policy.infer, obs) + infer_time = time.monotonic() - infer_time action["server_timing"] = { @@ -126,7 +172,9 @@ async def _handler(self, websocket: _server.ServerConnection): def _health_check( connection: _server.ServerConnection, request: _server.Request -) -> Optional[_server.Response]: +) -> _server.Response | None: if request.path == "/healthz": return connection.respond(http.HTTPStatus.OK, "OK\n") + if request.path == "/v2/health/ready": + return connection.respond(http.HTTPStatus.OK, "OK\n") return None diff --git a/wall_x/_vendor/harrix/utils/__init__.py b/wall_x/_vendor/harrix/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/wall_x/_vendor/harrix/utils/ckpt_load.py b/wall_x/_vendor/harrix/utils/ckpt_load.py new file mode 100644 index 00000000..4e003500 --- /dev/null +++ b/wall_x/_vendor/harrix/utils/ckpt_load.py @@ -0,0 +1,171 @@ +"""Load checkpoint weights and apply fused-format conversion when needed. + +Model-instance operations such as ``load_state_dict`` and ``set_normalizer`` +are intentionally left to the adapter. +""" + +from __future__ import annotations + +import os +from typing import Callable, Optional + +import torch +from safetensors.torch import load_file + + +def _noop_log(_msg: str, **_kw) -> None: + pass + + +def _align_checkpoint_tensor( + param: torch.Tensor, + target: torch.Tensor, + name: str, + log_fn: Callable, +) -> torch.Tensor | None: + """Crop or pad a checkpoint tensor to match the current model parameter.""" + if param.shape == target.shape: + return param + if param.ndim != target.ndim: + log_fn( + f"Skipping '{name}': ndim mismatch " + f"checkpoint={param.ndim} model={target.ndim}" + ) + return None + + overlap = tuple( + slice(0, min(src, dst)) for src, dst in zip(param.shape, target.shape) + ) + + if all(src >= dst for src, dst in zip(param.shape, target.shape)): + aligned = param[overlap].contiguous() + log_fn( + f"Cropped '{name}': checkpoint {tuple(param.shape)} " + f"-> model {tuple(aligned.shape)}" + ) + return aligned + + if all(src <= dst for src, dst in zip(param.shape, target.shape)): + aligned = target.detach().clone() + aligned[overlap] = param[overlap] + log_fn( + f"Padded '{name}': checkpoint {tuple(param.shape)} " + f"-> model {tuple(aligned.shape)} (tail keeps model init)" + ) + return aligned + + aligned = target.detach().clone() + aligned[overlap] = param[overlap] + log_fn( + f"Partially aligned '{name}': checkpoint {tuple(param.shape)} " + f"-> model {tuple(aligned.shape)} (non-overlap keeps model init)" + ) + return aligned + + +def reshape_compatible_state_dict( + state_dict: dict, model_sd: dict, log_fn: Optional[Callable] = None +) -> dict: + """Align checkpoint tensors to the target model shapes via crop / pad.""" + log_fn = log_fn or _noop_log + out = {} + for name, param in state_dict.items(): + if name not in model_sd: + log_fn(f"Not used parameter: {name}") + continue + target = model_sd[name] + if param.shape == target.shape: + out[name] = param + continue + aligned = _align_checkpoint_tensor(param, target, name, log_fn) + if aligned is not None: + out[name] = aligned + return out + + +def load_state_dict(checkpoint_path: str, model_class) -> dict: + """Load a state dict from a checkpoint directory. + + Supported formats: + - pytorch_model_fsdp.bin, optionally wrapped as {"state_dict": ...} + - model.safetensors + + If the model class reports that the state dict is not fused, it is converted + through ``model_class.convert_to_fused``. + """ + 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): + state_dict = torch.load(fsdp_ckpt, map_location="cpu") + if isinstance(state_dict, dict) and "state_dict" in state_dict: + state_dict = state_dict["state_dict"] + elif os.path.exists(safetensor_ckpt): + state_dict = load_file(safetensor_ckpt, device="cpu") + else: + raise FileNotFoundError( + "checkpoint contains neither pytorch_model_fsdp.bin nor model.safetensors: " + f"{checkpoint_path}" + ) + + if not model_class.is_fused(state_dict): + state_dict = model_class.convert_to_fused(state_dict) + + return state_dict + + +def read_global_step(checkpoint_path: str) -> int | None: + """Read ``global_step.pth`` when present.""" + p = os.path.join(checkpoint_path, "global_step.pth") + if not os.path.exists(p): + return None + payload = torch.load(p) + return int(payload["global_step"]) + + +def _dir_has_weights(path: str) -> bool: + return os.path.exists( + os.path.join(path, "pytorch_model_fsdp.bin") + ) or os.path.exists(os.path.join(path, "model.safetensors")) + + +def resolve_checkpoint_dir(checkpoint_path: str) -> str: + """Return a directory that directly contains model weights. + + Training saves under a root such as ``libero6/`` with step subdirs + ``libero6/0/``, ``libero6/3/``, etc. Inference callers may pass either the + root or a concrete step directory. + """ + if os.path.isfile(checkpoint_path): + checkpoint_path = os.path.dirname(checkpoint_path) + + if _dir_has_weights(checkpoint_path): + return checkpoint_path + + if not os.path.isdir(checkpoint_path): + raise FileNotFoundError(f"checkpoint path does not exist: {checkpoint_path}") + + candidates: list[tuple[int, float, str]] = [] + for entry in os.listdir(checkpoint_path): + sub = os.path.join(checkpoint_path, entry) + if not os.path.isdir(sub) or not _dir_has_weights(sub): + continue + step = read_global_step(sub) + sort_step = step if step is not None else -1 + candidates.append((sort_step, os.path.getmtime(sub), sub)) + + if not candidates: + return checkpoint_path + + candidates.sort() + resolved = candidates[-1][2] + if resolved != checkpoint_path: + import logging + + logging.getLogger(__name__).info( + "Resolved checkpoint root %s -> %s (global_step=%s)", + checkpoint_path, + resolved, + read_global_step(resolved), + ) + return resolved diff --git a/wall_x/_vendor/harrix/utils/normalizer.py b/wall_x/_vendor/harrix/utils/normalizer.py new file mode 100644 index 00000000..2494b83e --- /dev/null +++ b/wall_x/_vendor/harrix/utils/normalizer.py @@ -0,0 +1,135 @@ +"""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. +It does not fall back to internal default action statistics. +""" + +from __future__ import annotations + +import json +import logging +import os + +import torch + +from wall_x.data.backends.lerobot.utils import NormStats +from wall_x.model.core.action.normalizer import Normalizer, pad_normalizer_to_dim +from wall_x._vendor.harrix.utils.train_config import ( + resolve_agent_pos_config, + resolve_dof_config, +) + +logger = logging.getLogger(__name__) + + +def _load_norm_stats(norm_stats_path: str, action_key: str) -> NormStats: + with open(norm_stats_path, "r") as f: + norm_stats = json.load(f) + q01 = torch.tensor(norm_stats["norm_stats"][action_key]["q01"]) + q99 = torch.tensor(norm_stats["norm_stats"][action_key]["q99"]) + return NormStats(min=q01, max=q99, delta=q99 - q01) + + +def _load_custom_action_stats(train_config: dict) -> dict | None: + custom = train_config.get("customized_action_statistic_dof", None) + if not custom: + return None + with open(custom, "r") as f: + return json.load(f) + + +def _normalizer_from_stats(action_stats: dict, train_config: dict, key: str) -> Normalizer: + return Normalizer( + action_stats, + train_config[key], + min_key=train_config.get("min_key", "min"), + delta_key=train_config.get("delta_key", "delta"), + ) + + +def _missing_normalizer_error(checkpoint_path: str, train_config: dict) -> FileNotFoundError: + custom = train_config.get("customized_action_statistic_dof", None) + return FileNotFoundError( + "Public inference requires normalization data. Expected one of: " + f"{os.path.join(checkpoint_path, 'norm_stats.json')}; checkpoint-side " + "normalizer_action.pth and normalizer_propri.pth; or an explicit " + f"customized_action_statistic_dof path. Current customized_action_statistic_dof={custom!r}." + ) + + +def build_normalizers( + checkpoint_path: str, + train_config: dict, + norm_key: str, +) -> 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): + 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") + 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) + + if os.path.exists(action_pth): + normalizer_action = Normalizer.from_ckpt(action_pth) + else: + normalizer_action = _normalizer_from_stats(custom_stats, train_config, "dof_config") + + if os.path.exists(propri_pth): + normalizer_propri = Normalizer.from_ckpt(propri_pth) + else: + normalizer_propri = _normalizer_from_stats(custom_stats, train_config, "agent_pos_config") + + action_dim = sum(resolve_dof_config(train_config).values()) + propri_dim = sum(resolve_agent_pos_config(train_config).values()) + pad_normalizer_to_dim(normalizer_action, action_dim, "action") + pad_normalizer_to_dim(normalizer_propri, propri_dim, "propri") + + resolved = _resolve_norm_key(norm_key, normalizer_action, normalizer_propri) + return normalizer_action, normalizer_propri, resolved + + +def _resolve_norm_key( + norm_key: str, + normalizer_action: Normalizer, + normalizer_propri: Normalizer, +) -> str: + """Resolve a requested norm key against normalizer keys.""" + available = sorted( + set(normalizer_action.min.keys()) & set(normalizer_propri.min.keys()) + ) + if norm_key in available: + return norm_key + if not available: + return norm_key + + prefix_matches = [k for k in available if k.startswith(f"{norm_key}_")] + if len(prefix_matches) == 1: + logger.warning( + "norm_key=%r not found; using prefix fallback %r", + norm_key, + prefix_matches[0], + ) + return prefix_matches[0] + if len(available) == 1: + logger.warning( + "norm_key=%r not found; using the only available key %r", + norm_key, + available[0], + ) + return available[0] + + logger.warning( + "norm_key=%r not found; available=%s; returning the requested key unchanged", + norm_key, + available, + ) + return norm_key diff --git a/wall_x/_vendor/harrix/utils/seed.py b/wall_x/_vendor/harrix/utils/seed.py new file mode 100644 index 00000000..9c4fd8d2 --- /dev/null +++ b/wall_x/_vendor/harrix/utils/seed.py @@ -0,0 +1,23 @@ +"""Seed helpers for driver and environment worker processes.""" + +from __future__ import annotations + +import os +import random + +import numpy as np +import torch + + +def set_seed_everywhere(seed: int) -> None: + """Set random seeds and deterministic backend options.""" + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + np.random.seed(seed) + random.seed(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + os.environ["PYTHONHASHSEED"] = str(seed) + # Deterministic cuBLAS workspace plus warning-only deterministic op checks. + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + torch.use_deterministic_algorithms(True, warn_only=True) diff --git a/wall_x/_vendor/harrix/utils/train_config.py b/wall_x/_vendor/harrix/utils/train_config.py new file mode 100644 index 00000000..435c28af --- /dev/null +++ b/wall_x/_vendor/harrix/utils/train_config.py @@ -0,0 +1,371 @@ +"""Load checkpoint-side train config and construct runtime config objects. + +Inference adapters use this module to load the train config, apply checkpoint +overrides for moved processor/tokenizer files, build model config, and build the +data config required by preprocessing. +""" + +from __future__ import annotations + +import os +import logging + +import yaml + +logger = logging.getLogger(__name__) + + +_TASK_INFERENCE_KEYS = ( + "dof_config", + "agent_pos_config", + "ar_dof_config", + "action_horizon", + "action_horizon_flow", + "noise_scheduler", + "use_state_string_representation", +) + +_VIRTUAL_TAIL_KEYS = frozenset(("action_padding",)) + + +def _move_virtual_keys_to_tail(layout: dict | None) -> dict | None: + """Keep real LeRobot dims before virtual padding dims in inference layouts.""" + if not isinstance(layout, dict): + return layout + head = {k: v for k, v in layout.items() if k not in _VIRTUAL_TAIL_KEYS} + tail = {k: v for k, v in layout.items() if k in _VIRTUAL_TAIL_KEYS} + if not tail: + return layout + return {**head, **tail} + + +def _canonicalize_task_layouts(target: dict) -> None: + if os.environ.get("WALLX_CANONICALIZE_VIRTUAL_DOF_ORDER", "1") == "0": + return + task = target.get("task") + data = target.get("data") + for key in ("dof_config", "agent_pos_config", "ar_dof_config"): + value = _move_virtual_keys_to_tail(target.get(key)) + if value is not None: + target[key] = value + if isinstance(task, dict): + task_value = _move_virtual_keys_to_tail(task.get(key)) + if task_value is not None: + task[key] = task_value + if isinstance(data, dict): + data_value = _move_virtual_keys_to_tail(data.get(key)) + if data_value is not None: + data[key] = data_value + + +def resolve_use_state_string_representation(train_config: dict) -> bool: + """Read ``use_state_string_representation`` with task YAML as authority.""" + task = train_config.get("task") or {} + if task.get("use_state_string_representation") is not None: + return bool(task["use_state_string_representation"]) + if train_config.get("use_state_string_representation") is not None: + return bool(train_config["use_state_string_representation"]) + data = train_config.get("data") or {} + if data.get("use_state_string_representation") is not None: + return bool(data["use_state_string_representation"]) + return False + + +def resolve_state_bins(train_config: dict, default: int = 256) -> int: + """Read discretization bin count from flat or nested train config.""" + if train_config.get("state_bins") is not None: + return int(train_config["state_bins"]) + data = train_config.get("data") or {} + if data.get("state_bins") is not None: + return int(data["state_bins"]) + return default + + +def resolve_agent_pos_config(train_config: dict) -> dict: + """Return ``agent_pos_config`` from flat or ``task``-nested train YAML.""" + task = train_config.get("task") or {} + agent_pos_config = train_config.get("agent_pos_config") or task.get( + "agent_pos_config" + ) + if not agent_pos_config: + raise KeyError( + "agent_pos_config missing from train config " + "(expected top-level or task.agent_pos_config)" + ) + return dict(agent_pos_config) + + +def resolve_dof_config(train_config: dict) -> dict: + """Return ``dof_config`` from flat or ``task``-nested train YAML.""" + task = train_config.get("task") or {} + dof_config = train_config.get("dof_config") or task.get("dof_config") + if not dof_config: + raise KeyError( + "dof_config missing from train config (expected top-level or task.dof_config)" + ) + return dict(dof_config) + + +def resolve_cam_names_from_train_config(train_config: dict) -> list[str] | None: + """Infer model camera keys from ``data.key_mappings.camera`` (e.g. LIBERO).""" + data = train_config.get("data") or {} + key_mappings = data.get("key_mappings") or {} + camera = key_mappings.get("camera") or {} + if not camera: + return None + names: list[str] = [] + for value in camera.values(): + name = str(value) + if name not in names: + names.append(name) + return names or None + + +def resolve_camera_label(cam_name: str, camera_name_mapping: dict | None = None) -> str: + """Match training ``get_wallx_normal_text`` camera display names.""" + mapping = camera_name_mapping or {} + return mapping.get(cam_name, cam_name) + + +def resolve_max_length(train_config: dict, default: int = 768) -> int: + """Read ``max_length`` from flat or nested train YAML (collator default 768).""" + data = train_config.get("data") or {} + if train_config.get("max_length") is not None: + return int(train_config["max_length"]) + if data.get("max_length") is not None: + return int(data["max_length"]) + raw = train_config.get("_raw_data") or {} + if raw.get("max_length") is not None: + return int(raw["max_length"]) + return default + + +# Fields ``load_wallx_processors`` and model construction read from a flat dict. +_ACTION_TOKENIZER_KEYS = ( + "action_tokenizer_type", + "action_tokenizer_path", + "action_tokenizer_checkpoint_path", + "action_tokenizer_config_dir", + "action_tokenizer", +) + +_INFERENCE_MODEL_KEYS = ( + "processor_path", + "pretrained_path", + "config_path", + *_ACTION_TOKENIZER_KEYS, + "ar_loss_weight", + "attn_deterministic", + "flow_loss_weight", + "use_ema", +) + + +def _merge_model_fields(target: dict, source: dict, *, overwrite: bool = False) -> None: + """Copy model-related keys from ``source`` into flat ``target``.""" + if not source: + return + for key in _INFERENCE_MODEL_KEYS: + value = source.get(key) + if value is None: + continue + if overwrite or target.get(key) in (None, ""): + target[key] = value + + +def _mirror_task_fields(target: dict, task: dict) -> None: + """Copy typed task fields into the flat/data legacy inference mirrors.""" + if not isinstance(task, dict): + return + data = target.setdefault("data", {}) + if not isinstance(data, dict): + data = {} + target["data"] = data + for key in _TASK_INFERENCE_KEYS: + value = task.get(key) + if value is None: + continue + target[key] = value + data[key] = value + + +def _load_ckpt_config_overlay(checkpoint_path: str) -> dict: + """Read ``config.yml`` saved beside a training checkpoint.""" + ckpt_yml = os.path.join(checkpoint_path, "config.yml") + if not os.path.exists(ckpt_yml): + return {} + with open(ckpt_yml, "r") as f: + raw = yaml.load(f, Loader=yaml.FullLoader) or {} + + overlay: dict = {} + _merge_model_fields(overlay, raw.get("model") or {}, overwrite=True) + raw_yaml = raw.get("_raw_yaml") or {} + _merge_model_fields(overlay, raw_yaml.get("model") or {}, overwrite=False) + _merge_model_fields(overlay, raw, overwrite=False) + return overlay + + +def strip_action_tokenizer_fields(train_config: dict) -> None: + """Drop action-tokenizer fields so flow inference can skip AR tokenizer setup.""" + for key in _ACTION_TOKENIZER_KEYS: + train_config.pop(key, None) + model = train_config.get("model") + if isinstance(model, dict): + for key in _ACTION_TOKENIZER_KEYS: + model.pop(key, None) + + +def load_train_config_with_ckpt_overlay( + train_config_path: str, + checkpoint_path: str, +) -> dict: + """Load train YAML and apply checkpoint-local processor/tokenizer overlays. + + This lets a checkpoint remain portable when the original training-machine + processor or action-tokenizer paths are no longer available. + """ + with open(train_config_path, "r") as f: + train_config = yaml.load(f, Loader=yaml.FullLoader) + + preprocessor_file = os.path.join(checkpoint_path, "preprocessor_config.json") + if os.path.exists(preprocessor_file): + train_config["processor_path"] = checkpoint_path + + orig_action_tok = train_config.get("action_tokenizer_path", None) + if orig_action_tok is not None and not os.path.exists(orig_action_tok): + tokenizer_file = os.path.join(checkpoint_path, "tokenizer.json") + tokenizer_config_file = os.path.join(checkpoint_path, "tokenizer_config.json") + if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file): + train_config["action_tokenizer_path"] = checkpoint_path + + ckpt_overlay = _load_ckpt_config_overlay(checkpoint_path) + _merge_model_fields(train_config, ckpt_overlay, overwrite=False) + if isinstance(train_config.get("model"), dict): + _merge_model_fields(train_config["model"], ckpt_overlay, overwrite=False) + + return train_config + + +def normalize_train_config_for_inference( + train_config: dict, + train_config_path: str, +) -> dict: + """Flatten typed TrainConfig YAML into the legacy dict inference expects. + + Serving and ``update_model_config`` still read top-level ``dof_config`` and + nested ``data.*`` fields. New training yamls store those under ``task:``. + """ + try: + from wall_x.config.loader import load_config + + typed_cfg = load_config(train_config_path) + except Exception as e: + logger.debug( + "Train config is not typed TrainConfig schema (%s): %s", + train_config_path, + e, + ) + _mirror_task_fields(train_config, train_config.get("task") or {}) + _canonicalize_task_layouts(train_config) + return train_config + + import dataclasses + + normalized = typed_cfg.build_data_loader_dict() + model_dict = dataclasses.asdict(typed_cfg.model) + _merge_model_fields(normalized, model_dict, overwrite=True) + + normalized["data"] = dict(normalized.get("data") or {}) + if isinstance(train_config.get("model"), dict): + _merge_model_fields(normalized, train_config["model"], overwrite=False) + for key, value in train_config.items(): + if key in _INFERENCE_MODEL_KEYS or key == "qwen_vl_act_config_path": + if value is not None: + normalized[key] = value + elif key == "data" and isinstance(value, dict): + normalized["data"].update(value) + + _mirror_task_fields(normalized, dataclasses.asdict(typed_cfg.task)) + _canonicalize_task_layouts(normalized) + return normalized + + +def register_data_backend(train_config: dict) -> None: + """Register the data backend before typed data config construction.""" + from wall_x.data._registry import _set_data_backend + + data_section = train_config.get("data", {}) + dataset_type = train_config.get("dataset_type") or data_section.get( + "dataset_type", "lerobot" + ) + _set_data_backend(dataset_type) + + +def build_model_config( + config_class, + checkpoint_path: str, + train_config: dict, + train_config_path: str | None = None, +): + """Build the HF model config and inject train_config-derived fields. + + ``config_class`` is supplied by the variant adapter. If the checkpoint does + not include ``config.json``, this falls back to + ``train_config["qwen_vl_act_config_path"]``. + """ + ckpt_config_path = os.path.join(checkpoint_path, "config.json") + if os.path.exists(ckpt_config_path): + resolved = ckpt_config_path + else: + resolved = train_config.get("qwen_vl_act_config_path") + if resolved is None or not os.path.exists(resolved): + raise ValueError( + f"cannot load model config: checkpoint file {ckpt_config_path} " + f"does not exist, and fallback qwen_vl_act_config_path={resolved!r} " + "does not exist either" + ) + train_config["qwen_vl_act_config_path"] = resolved + + if resolved.endswith(".json"): + model_config = config_class.from_json_file(resolved) + else: + model_config = config_class.from_pretrained(resolved) + + legacy_train_config = ( + normalize_train_config_for_inference(train_config, train_config_path) + if train_config_path + else train_config + ) + model_config.update_model_config(legacy_train_config) + model_config._attn_implementation = "sdpa" + model_config.vision_config._attn_implementation = "flash_attention_2" + + return model_config + + +def build_data_config(train_config_path: str, train_config: dict): + """Build the data config used by image resizing and preprocessing. + + New checkpoints use the typed schema. Older checkpoints may still use a flat + legacy schema, so this falls back to the active data backend. + """ + from wall_x.config.loader import load_config + from wall_x.data import data_backend + + try: + typed_cfg = load_config(train_config_path) + backend = data_backend() + if backend.supports("load_trainer_data_config"): + return backend.load_trainer_data_config(typed_cfg) + except Exception as e: + logger.warning( + "typed TrainConfig loading failed; falling back to raw data config: %s", + e, + ) + + backend = data_backend() + if backend.supports("load_trainer_data_config_from_yaml_dict"): + return backend.load_trainer_data_config_from_yaml_dict(train_config) + raise RuntimeError( + f"active data backend {backend!r} cannot build a trainer data config" + ) diff --git a/wall_x/_vendor/x2robot_utils/__init__.py b/wall_x/_vendor/x2robot_utils/__init__.py new file mode 100644 index 00000000..535937fe --- /dev/null +++ b/wall_x/_vendor/x2robot_utils/__init__.py @@ -0,0 +1,5 @@ +"""Vendored subset of x2robot_utils for Wall-X.""" + +__version__ = "0.2.0-vendored" + +__all__ = ["__version__"] diff --git a/wall_x/_vendor/x2robot_utils/geometry.py b/wall_x/_vendor/x2robot_utils/geometry.py new file mode 100644 index 00000000..0125da4c --- /dev/null +++ b/wall_x/_vendor/x2robot_utils/geometry.py @@ -0,0 +1,309 @@ +"""Rotation / pose geometry utilities (pure numpy + numba). + +Shared between wall-x and internal_dataset_backend: +- ``euler_to_matrix_zyx_6d_nb``: ZYX Euler → flattened top-2-rows of R (N, 6) +- ``so3_to_euler_zyx_batch_nb``: 6D rotation → ZYX Euler (canonicalized) +- ``compose_state_and_delta_to_abs_{rpy,6d}``: state ⊕ delta → absolute pose +""" + +import numpy as np +from numba import jit, prange + + +@jit(nopython=True) +def euler_to_matrix_zyx_6d_nb(eulers): + """Euler angles (N, 3) → flattened top two rows of rotation matrix (N, 6).""" + N = eulers.shape[0] + R6 = np.empty((N, 6), dtype=np.float64) + for i in prange(N): + roll = eulers[i, 0] + pitch = eulers[i, 1] + yaw = eulers[i, 2] + + cy, sy = np.cos(yaw), np.sin(yaw) + cp, sp = np.cos(pitch), np.sin(pitch) + cr, sr = np.cos(roll), np.sin(roll) + + r00 = cy * cp + r01 = cy * sp * sr - sy * cr + r02 = cy * sp * cr + sy * sr + + r10 = sy * cp + r11 = sy * sp * sr + cy * cr + r12 = sy * sp * cr - cy * sr + + R6[i, 0] = r00 + R6[i, 1] = r01 + R6[i, 2] = r02 + R6[i, 3] = r10 + R6[i, 4] = r11 + R6[i, 5] = r12 + return R6 + + +@jit(nopython=True) +def euler_to_matrix_zyx_batch_nb(eulers): + N = eulers.shape[0] + R = np.empty((N, 3, 3), dtype=np.float64) + for i in prange(N): + roll = eulers[i, 0] + pitch = eulers[i, 1] + yaw = eulers[i, 2] + + cy, sy = np.cos(yaw), np.sin(yaw) + cp, sp = np.cos(pitch), np.sin(pitch) + cr, sr = np.cos(roll), np.sin(roll) + + R[i, 0, 0] = cy * cp + R[i, 0, 1] = cy * sp * sr - sy * cr + R[i, 0, 2] = cy * sp * cr + sy * sr + + R[i, 1, 0] = sy * cp + R[i, 1, 1] = sy * sp * sr + cy * cr + R[i, 1, 2] = sy * sp * cr - cy * sr + + R[i, 2, 0] = -sp + R[i, 2, 1] = cp * sr + R[i, 2, 2] = cp * cr + return R + + +@jit(nopython=True) +def matrix_to_euler_zyx_batch_nb(Rs): + """R = Rz(yaw) * Ry(pitch) * Rx(roll) → (roll, pitch, yaw).""" + N = Rs.shape[0] + eulers = np.empty((N, 3), dtype=np.float64) + for i in prange(N): + r00 = Rs[i, 0, 0] + r10 = Rs[i, 1, 0] + r20 = Rs[i, 2, 0] + r21 = Rs[i, 2, 1] + r22 = Rs[i, 2, 2] + + x = -r20 + if x > 1.0: + x = 1.0 + elif x < -1.0: + x = -1.0 + + pitch = np.arcsin(x) + roll = np.arctan2(r21, r22) + yaw = np.arctan2(r10, r00) + + eulers[i, 0] = roll + eulers[i, 1] = pitch + eulers[i, 2] = yaw + return eulers + + +@jit(nopython=True) +def so3_to_matrix_batch_nb(batch_so3): + N = batch_so3.shape[0] + R_all = np.empty((N, 3, 3), dtype=np.float64) + eps = 1e-12 + for i in prange(N): + r1x, r1y, r1z = batch_so3[i, 0], batch_so3[i, 1], batch_so3[i, 2] + r2x, r2y, r2z = batch_so3[i, 3], batch_so3[i, 4], batch_so3[i, 5] + + n1 = np.sqrt(r1x * r1x + r1y * r1y + r1z * r1z) + eps + r1x /= n1 + r1y /= n1 + r1z /= n1 + + dot12 = r1x * r2x + r1y * r2y + r1z * r2z + r2x -= dot12 * r1x + r2y -= dot12 * r1y + r2z -= dot12 * r1z + n2 = np.sqrt(r2x * r2x + r2y * r2y + r2z * r2z) + eps + r2x /= n2 + r2y /= n2 + r2z /= n2 + + r3x = r1y * r2z - r1z * r2y + r3y = r1z * r2x - r1x * r2z + r3z = r1x * r2y - r1y * r2x + + R_all[i, 0, 0] = r1x + R_all[i, 0, 1] = r1y + R_all[i, 0, 2] = r1z + R_all[i, 1, 0] = r2x + R_all[i, 1, 1] = r2y + R_all[i, 1, 2] = r2z + R_all[i, 2, 0] = r3x + R_all[i, 2, 1] = r3y + R_all[i, 2, 2] = r3z + return R_all + + +@jit(nopython=True) +def canonicalize_euler_zyx_batch_nb(rpy_batch): + """Canonicalize ZYX Euler angles so each component falls in (-π, π].""" + N = rpy_batch.shape[0] + out = np.empty_like(rpy_batch) + two_pi = 2.0 * np.pi + + for i in prange(N): + r = rpy_batch[i, 0] + p = rpy_batch[i, 1] + y = rpy_batch[i, 2] + + r = (r + np.pi) % two_pi - np.pi + p = (p + np.pi) % two_pi - np.pi + y = (y + np.pi) % two_pi - np.pi + + if p > np.pi / 2.0: + p = np.pi - p + r = r + np.pi + y = y + np.pi + elif p <= -np.pi / 2.0: + p = -np.pi - p + r = r + np.pi + y = y + np.pi + + r = (r + np.pi) % two_pi - np.pi + p = (p + np.pi) % two_pi - np.pi + y = (y + np.pi) % two_pi - np.pi + + out[i, 0] = r + out[i, 1] = p + out[i, 2] = y + + return out + + +def so3_to_euler_zyx_batch_nb(batch_so3): + matrix = so3_to_matrix_batch_nb(batch_so3) + eulers = matrix_to_euler_zyx_batch_nb(matrix) + return canonicalize_euler_zyx_batch_nb(eulers) + + +@jit(nopython=True) +def compose_state_and_delta_to_abs_rpy(delta, state): + """Compose a delta (ZYX rpy or 6D) with an absolute state → absolute rpy(ZYX). + + delta: (N, 3) Δrpy or (N, 6) Δ6D. state: (3,) rpy or (6,) 6D. + Output: (N, 3) rpy canonicalized into (-π, π]. + """ + if delta.shape[-1] == 3: + R_delta = euler_to_matrix_zyx_batch_nb(delta) + elif delta.shape[-1] == 6: + R_delta = so3_to_matrix_batch_nb(delta) + else: + raise ValueError(f"delta last dim must be 3 or 6, got {delta.shape[-1]}") + + if state.shape[-1] == 3: + R_state = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0] + elif state.shape[-1] == 6: + R_state = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] + else: + raise ValueError(f"state last dim must be 3 or 6, got {state.shape[-1]}") + + N = R_delta.shape[0] + R_abs = np.empty((N, 3, 3), dtype=np.float64) + + S00 = R_state[0, 0] + S01 = R_state[0, 1] + S02 = R_state[0, 2] + S10 = R_state[1, 0] + S11 = R_state[1, 1] + S12 = R_state[1, 2] + S20 = R_state[2, 0] + S21 = R_state[2, 1] + S22 = R_state[2, 2] + + for i in prange(N): + A00 = R_delta[i, 0, 0] + A01 = R_delta[i, 0, 1] + A02 = R_delta[i, 0, 2] + A10 = R_delta[i, 1, 0] + A11 = R_delta[i, 1, 1] + A12 = R_delta[i, 1, 2] + A20 = R_delta[i, 2, 0] + A21 = R_delta[i, 2, 1] + A22 = R_delta[i, 2, 2] + + R_abs[i, 0, 0] = A00 * S00 + A01 * S10 + A02 * S20 + R_abs[i, 0, 1] = A00 * S01 + A01 * S11 + A02 * S21 + R_abs[i, 0, 2] = A00 * S02 + A01 * S12 + A02 * S22 + + R_abs[i, 1, 0] = A10 * S00 + A11 * S10 + A12 * S20 + R_abs[i, 1, 1] = A10 * S01 + A11 * S11 + A12 * S21 + R_abs[i, 1, 2] = A10 * S02 + A11 * S12 + A12 * S22 + + R_abs[i, 2, 0] = A20 * S00 + A21 * S10 + A22 * S20 + R_abs[i, 2, 1] = A20 * S01 + A21 * S11 + A22 * S21 + R_abs[i, 2, 2] = A20 * S02 + A21 * S12 + A22 * S22 + + abs_rpy = matrix_to_euler_zyx_batch_nb(R_abs) + abs_rpy = canonicalize_euler_zyx_batch_nb(abs_rpy) + + return abs_rpy + + +@jit(nopython=True) +def compose_state_and_delta_to_abs_6d(delta, state): + """Compose a 6D delta with a 6D state → absolute 6D rotation. + + delta: (N, 6). state: (6,). Output: (N, 6). + """ + R_delta = so3_to_matrix_batch_nb(delta) + R_state = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] + + N = R_delta.shape[0] + R_abs = np.empty((N, 3, 3), dtype=np.float64) + + S00 = R_state[0, 0] + S01 = R_state[0, 1] + S02 = R_state[0, 2] + S10 = R_state[1, 0] + S11 = R_state[1, 1] + S12 = R_state[1, 2] + S20 = R_state[2, 0] + S21 = R_state[2, 1] + S22 = R_state[2, 2] + + for i in prange(N): + A00 = R_delta[i, 0, 0] + A01 = R_delta[i, 0, 1] + A02 = R_delta[i, 0, 2] + A10 = R_delta[i, 1, 0] + A11 = R_delta[i, 1, 1] + A12 = R_delta[i, 1, 2] + A20 = R_delta[i, 2, 0] + A21 = R_delta[i, 2, 1] + A22 = R_delta[i, 2, 2] + + R_abs[i, 0, 0] = A00 * S00 + A01 * S10 + A02 * S20 + R_abs[i, 0, 1] = A00 * S01 + A01 * S11 + A02 * S21 + R_abs[i, 0, 2] = A00 * S02 + A01 * S12 + A02 * S22 + + R_abs[i, 1, 0] = A10 * S00 + A11 * S10 + A12 * S20 + R_abs[i, 1, 1] = A10 * S01 + A11 * S11 + A12 * S21 + R_abs[i, 1, 2] = A10 * S02 + A11 * S12 + A12 * S22 + + R_abs[i, 2, 0] = A20 * S00 + A21 * S10 + A22 * S20 + R_abs[i, 2, 1] = A20 * S01 + A21 * S11 + A22 * S21 + R_abs[i, 2, 2] = A20 * S02 + A21 * S12 + A22 * S22 + + abs_6d = np.empty((N, 6), dtype=np.float64) + for i in prange(N): + abs_6d[i, 0] = R_abs[i, 0, 0] + abs_6d[i, 1] = R_abs[i, 0, 1] + abs_6d[i, 2] = R_abs[i, 0, 2] + abs_6d[i, 3] = R_abs[i, 1, 0] + abs_6d[i, 4] = R_abs[i, 1, 1] + abs_6d[i, 5] = R_abs[i, 1, 2] + + return abs_6d + + +__all__ = [ + "euler_to_matrix_zyx_6d_nb", + "euler_to_matrix_zyx_batch_nb", + "matrix_to_euler_zyx_batch_nb", + "so3_to_matrix_batch_nb", + "canonicalize_euler_zyx_batch_nb", + "so3_to_euler_zyx_batch_nb", + "compose_state_and_delta_to_abs_rpy", + "compose_state_and_delta_to_abs_6d", +] diff --git a/wall_x/_vendor/x2robot_utils/grounding.py b/wall_x/_vendor/x2robot_utils/grounding.py new file mode 100644 index 00000000..0852cd01 --- /dev/null +++ b/wall_x/_vendor/x2robot_utils/grounding.py @@ -0,0 +1,172 @@ +"""Grounding-point helpers (point / bbox coordinate remap, pure regex).""" + +import re +from typing import List, Optional + + +def process_grounding_points( + text: str, orig_height, orig_width, resized_height, resized_width, model_type +) -> str: + """Remap // coordinates inside ``text`` from the original + image size to the resized space used by the given model type. + """ + point_pattern = re.compile(r"<(point|box|bbox)>(.*?)") + + def process_match(match): + tag_name = match.group(1) + coords_str = match.group(2) + try: + coords = list(map(int, re.findall(r"\d+", coords_str))) + + scale_w = resized_width / orig_width + scale_h = resized_height / orig_height + + if len(coords) == 2: + x, y = coords + if model_type == "qwen2_5": + new_x = max(0, min(round(x * scale_w), resized_width - 1)) + new_y = max(0, min(round(y * scale_h), resized_height - 1)) + elif model_type in ["qwen2", "qwen3", "qwen3_5"]: + new_x = max(0, min(999.999, (x / orig_width) * 1000)) + new_y = max(0, min(999.999, (y / orig_height) * 1000)) + else: + raise ValueError("Unsupported model type") + coords = [new_x, new_y] + + if len(coords) == 4: + x1, y1, x2, y2 = coords + if model_type == "qwen2_5": + new_x1 = max(0, min(round(x1 * scale_w), resized_width - 1)) + new_y1 = max(0, min(round(y1 * scale_h), resized_height - 1)) + new_x2 = max(0, min(round(x2 * scale_w), resized_width - 1)) + new_y2 = max(0, min(round(y2 * scale_h), resized_height - 1)) + elif model_type in ["qwen2", "qwen3", "qwen3_5"]: + new_x1 = max(0, min(999.999, (x1 / orig_width) * 1000)) + new_y1 = max(0, min(999.999, (y1 / orig_height) * 1000)) + new_x2 = max(0, min(999.999, (x2 / orig_width) * 1000)) + new_y2 = max(0, min(999.999, (y2 / orig_height) * 1000)) + else: + raise ValueError("Unsupported model type") + coords = [new_x1, new_y1, new_x2, new_y2] + + return f'<{tag_name}>[{", ".join(map(str, coords))}]' + + except (ValueError, TypeError): + return match.group(0) + + return point_pattern.sub(process_match, text) + + +def extract_grounding_points(text: str) -> List[List[float]]: + """Extract all // coordinates from ``text`` as list-of-list.""" + point_pattern = re.compile(r"<(point|box|bbox)>\s*\[([^\]]+)\]\s*") + + points: List[List[float]] = [] + for match in point_pattern.finditer(text): + coords_str = match.group(2) + raw_values = re.findall(r"-?\d+\.?\d*", coords_str) + converted: List[float] = [] + for value in raw_values: + number = float(value) + converted.append(int(number) if number.is_integer() else number) + if converted: + points.append(converted) + + return points + + +def reverse_grounding_points( + text: str, orig_height, orig_width, resized_height, resized_width, model_type +) -> str: + """Inverse of ``process_grounding_points`` — map resized coords back to original.""" + point_pattern = re.compile(r"<(point|box|bbox)>(.*?)") + + def reverse_match(match): + tag_name = match.group(1) + coords_str = match.group(2) + try: + coords = list(map(float, re.findall(r"-?\d+\.?\d*", coords_str))) + + scale_w = resized_width / orig_width + scale_h = resized_height / orig_height + + if len(coords) == 2: + x, y = coords + if model_type == "qwen2_5": + orig_x = max(0, min(orig_width - 1, round(x / scale_w))) + orig_y = max(0, min(orig_height - 1, round(y / scale_h))) + elif model_type in ["qwen2", "qwen3", "qwen3_5"]: + orig_x = max( + 0, min(orig_width - 1, round((x / 1000) * orig_width)) + ) + orig_y = max( + 0, min(orig_height - 1, round((y / 1000) * orig_height)) + ) + else: + raise ValueError("Unsupported model type") + coords = [orig_x, orig_y] + + if len(coords) == 4: + x1, y1, x2, y2 = coords + if model_type == "qwen2_5": + orig_x1 = max(0, min(orig_width - 1, round(x1 / scale_w))) + orig_y1 = max(0, min(orig_height - 1, round(y1 / scale_h))) + orig_x2 = max(0, min(orig_width - 1, round(x2 / scale_w))) + orig_y2 = max(0, min(orig_height - 1, round(y2 / scale_h))) + elif model_type in ["qwen2", "qwen3", "qwen3_5"]: + orig_x1 = max( + 0, min(orig_width - 1, round((x1 / 1000) * orig_width)) + ) + orig_y1 = max( + 0, min(orig_height - 1, round((y1 / 1000) * orig_height)) + ) + orig_x2 = max( + 0, min(orig_width - 1, round((x2 / 1000) * orig_width)) + ) + orig_y2 = max( + 0, min(orig_height - 1, round((y2 / 1000) * orig_height)) + ) + else: + raise ValueError("Unsupported model type") + coords = [orig_x1, orig_y1, orig_x2, orig_y2] + + return f'<{tag_name}>[{", ".join(map(str, map(int, coords)))}]' + + except (ValueError, TypeError): + return match.group(0) + + return point_pattern.sub(reverse_match, text) + + +def calculate_point_l1_distance(gt_text: str, pred_text: str) -> Optional[float]: + """Average L1 distance between 2D points extracted from ``gt_text`` / ``pred_text``. + + Returns None if either side has no points or counts differ. + """ + point_pattern = re.compile(r"\[(\d+),\s*(\d+)\]") + + gt_matches = point_pattern.findall(gt_text) + pred_matches = point_pattern.findall(pred_text) + + if not gt_matches or not pred_matches or len(gt_matches) != len(pred_matches): + return None + + total_l1_distance = 0.0 + for (gt_x, gt_y), (pred_x, pred_y) in zip(gt_matches, pred_matches): + try: + gt_x, gt_y = int(gt_x), int(gt_y) + pred_x, pred_y = int(pred_x), int(pred_y) + l1_dist = abs(gt_x - pred_x) + abs(gt_y - pred_y) + total_l1_distance += l1_dist + except ValueError: + continue + + return total_l1_distance / len(gt_matches) if gt_matches else None + + +__all__ = [ + "process_grounding_points", + "extract_grounding_points", + "reverse_grounding_points", + "calculate_point_l1_distance", +] diff --git a/wall_x/_vendor/x2robot_utils/text_templates.py b/wall_x/_vendor/x2robot_utils/text_templates.py new file mode 100644 index 00000000..ae1ca20e --- /dev/null +++ b/wall_x/_vendor/x2robot_utils/text_templates.py @@ -0,0 +1,235 @@ +"""Multimodal text preprocessing helpers. + +This file is generated by ``scripts/export_opensource.py``. It keeps only the +processor wrapper needed by harrix inference and uses a generic system prompt. +Internal robot-id, dataset, camera, and frequency maps are not bundled here. +""" + +from __future__ import annotations + +import numpy as np +import torch +from transformers import BatchFeature +from transformers.tokenization_utils_base import BatchEncoding + + +def pad_text_input_to_target_length( + text_inputs, target_length, pad_token_id=151643, padding_side="right" +): + """Pad or truncate tokenized text to ``target_length``.""" + batch_size, current_length = text_inputs.input_ids.shape + if current_length < target_length: + padding_size = target_length - current_length + padding = torch.full( + (batch_size, padding_size), + pad_token_id, + dtype=text_inputs.input_ids.dtype, + device=text_inputs.input_ids.device, + ) + attention_padding = torch.zeros( + (batch_size, padding_size), + dtype=text_inputs.attention_mask.dtype, + device=text_inputs.attention_mask.device, + ) + if padding_side == "right": + text_inputs["input_ids"] = torch.cat([text_inputs.input_ids, padding], dim=1) + text_inputs["attention_mask"] = torch.cat( + [text_inputs.attention_mask, attention_padding], dim=1 + ) + else: + text_inputs["input_ids"] = torch.cat([padding, text_inputs.input_ids], dim=1) + text_inputs["attention_mask"] = torch.cat( + [attention_padding, text_inputs.attention_mask], dim=1 + ) + elif current_length > target_length: + if padding_side == "right": + text_inputs["input_ids"] = text_inputs.input_ids[:, :target_length] + text_inputs["attention_mask"] = text_inputs.attention_mask[:, :target_length] + else: + text_inputs["input_ids"] = text_inputs.input_ids[:, -target_length:] + text_inputs["attention_mask"] = text_inputs.attention_mask[:, -target_length:] + return text_inputs + + +def _replace_media_placeholders(text, grid_thw, token, merge_length): + if grid_thw is None: + return text + index = 0 + for i in range(len(text)): + while token in text[i]: + if index >= len(grid_thw): + raise ValueError( + f"More {token} placeholders than media tensors in sample {i}" + ) + token_count = int(grid_thw[index].prod() // merge_length) + text[i] = text[i].replace(token, "<|placeholder|>" * token_count, 1) + index += 1 + text[i] = text[i].replace("<|placeholder|>", token) + return text + + +_PUBLIC_CAMERA_LABELS = { + "face_view": "front view", + "right_wrist_view": "right wrist view", + "left_wrist_view": "left wrist view", +} + + +def _camera_label(cam_name): + return _PUBLIC_CAMERA_LABELS.get(str(cam_name), str(cam_name).replace("_", " ")) + + +def preprocesser_call( + processor, + norm_state=None, + agent_pos_mask=None, + images=None, + prefix_text=None, + postfix_text=None, + videos=None, + padding=False, + padding_side="left", + truncation=None, + max_length=None, + return_tensors="pt", + pad_prefix_to_same_length=False, + pad_to_128_multiple=True, + state_augmentation_prob=0.0, + state_augmentation_ratio=0.0, + state_bins=256, + inference_mode=False, + **_, +): + """Build a ``BatchFeature`` for Wall-X VLA inference. + + This is the inference subset of the internal preprocessing helper: text, + image/video placeholder expansion, optional discretized proprioception + strings, padding, and labels=None for inference. + """ + if prefix_text is None: + raise ValueError("prefix_text is required") + if postfix_text is None: + postfix_text = [""] * len(prefix_text) + if not isinstance(prefix_text, list): + prefix_text = [prefix_text] + if not isinstance(postfix_text, list): + postfix_text = [postfix_text] + batch_size = len(prefix_text) + + if images is not None and len(images) > 0: + image_inputs = processor.image_processor(images=images, return_tensors=return_tensors) + image_grid_thw = image_inputs["image_grid_thw"] + else: + image_inputs = {} + image_grid_thw = None + + if videos is not None: + if hasattr(processor, "video_processor"): + videos_inputs = processor.video_processor(videos=videos, return_tensors=return_tensors) + else: + videos_inputs = processor.image_processor( + images=None, videos=videos, return_tensors=return_tensors + ) + video_grid_thw = videos_inputs["video_grid_thw"] + else: + videos_inputs = {} + video_grid_thw = None + + merge_length = processor.image_processor.merge_size**2 + prefix_text = _replace_media_placeholders( + list(prefix_text), image_grid_thw, "<|image_pad|>", merge_length + ) + prefix_text = _replace_media_placeholders( + prefix_text, video_grid_thw, "<|video_pad|>", merge_length + ) + + if norm_state is not None: + norm_state = norm_state.cpu().numpy() if isinstance(norm_state, torch.Tensor) else norm_state + agent_pos_mask = ( + agent_pos_mask[:, 0, :].cpu().numpy().astype(bool) + if isinstance(agent_pos_mask, torch.Tensor) + else agent_pos_mask[:, 0, :].astype(bool) + ) + discretized = np.digitize(norm_state, bins=np.linspace(-1, 1, state_bins + 1)[:-1]) - 1 + discretized = discretized[:, 0, :] + for i in range(batch_size): + if "<|propri|>" not in prefix_text[i]: + continue + state_str = " ".join(map(str, discretized[i, agent_pos_mask[i]])) + prefix_text[i] = prefix_text[i].replace("<|propri|>", state_str) + + if not pad_prefix_to_same_length: + text = [pre + post for pre, post in zip(prefix_text, postfix_text)] + text_inputs = processor.tokenizer( + text, + return_tensors=return_tensors, + padding=padding, + padding_side=padding_side, + truncation=truncation, + max_length=max_length, + ) + text_inputs["prefix_length"] = None + else: + prefix_inputs = processor.tokenizer( + prefix_text, + return_tensors=return_tensors, + padding=padding, + padding_side="left", + truncation=truncation, + max_length=max_length, + ) + postfix_inputs = processor.tokenizer( + postfix_text, + return_tensors=return_tensors, + padding=padding, + padding_side="right", + truncation=truncation, + max_length=max_length, + ) + text_inputs = BatchEncoding( + data={ + "input_ids": torch.cat([prefix_inputs.input_ids, postfix_inputs.input_ids], dim=1), + "attention_mask": torch.cat( + [prefix_inputs.attention_mask, postfix_inputs.attention_mask], dim=1 + ), + "prefix_length": prefix_inputs.input_ids.shape[1], + } + ) + + pad_token_id = processor.tokenizer.pad_token_id + if pad_token_id is None: + pad_token_id = processor.tokenizer.eos_token_id + if pad_to_128_multiple: + target_length = 128 * ((max(len(t) for t in text_inputs.input_ids) + 127) // 128) + text_inputs = pad_text_input_to_target_length( + text_inputs, target_length, pad_token_id=pad_token_id, padding_side=padding_side + ) + + text_inputs["labels"] = None if inference_mode else None + return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}) + + +def get_prologue_with_embodied_information(dataset_name, cam_mapping, robot_id, uid, config): + """Return a generic VLA system prompt without private robot maps.""" + role_start = "<|im_start|>" + role_end = "<|im_end|>" + prologue = ( + f"{role_start}system\n" + "You are an embodied vision-language-action model controlling a robot " + "with language instructions." + ) + if cam_mapping: + cameras = ", ".join(_camera_label(name) for name in cam_mapping.values()) + prologue += f"\nCamera Setup: {cameras}" + if not getattr(config, "use_relative_action", False): + prologue += "\nAction Space: Abs EEF" + else: + prologue += "\nAction Space: Rel EEF" + return f"{prologue}\n{role_end}\n" + + +__all__ = [ + "preprocesser_call", + "get_prologue_with_embodied_information", + "pad_text_input_to_target_length", +] diff --git a/wall_x/config/__init__.py b/wall_x/config/__init__.py new file mode 100644 index 00000000..329780a9 --- /dev/null +++ b/wall_x/config/__init__.py @@ -0,0 +1,62 @@ +"""Typed training configuration system.""" + +from .data_config import DataConfig, LeRobotDataConfig +from .hyperparams_config import ( + AdamWConfig, + ConstantSchedulerConfig, + CosineSchedulerConfig, + DMuonConfig, + OptimizerConfig, + SchedulerConfig, + StepSchedulerConfig, + TrainHyperParams, +) +from .infra_config import ( + CheckpointConfig, + DebugConfig, + DistributedConfig, + LoggingConfig, +) +from .loader import load_config, save_config +from .model_config import ModelConfig, QActModelConfig +from .registry import ( + register_data_config, + register_model_config, + register_optimizer_config, + register_scheduler_config, +) +from .task_config import TaskConfig +from .train_config import TrainConfig + +__all__ = [ + # Core + "TrainConfig", + "TaskConfig", + # Model configs + "ModelConfig", + "QActModelConfig", + # Data configs + "DataConfig", + "LeRobotDataConfig", + # Hyperparams + "TrainHyperParams", + "OptimizerConfig", + "AdamWConfig", + "DMuonConfig", + "SchedulerConfig", + "CosineSchedulerConfig", + "ConstantSchedulerConfig", + "StepSchedulerConfig", + # Infra configs + "DistributedConfig", + "LoggingConfig", + "CheckpointConfig", + "DebugConfig", + # Functions + "load_config", + "save_config", + "register_data_config", + "register_model_config", + "register_optimizer_config", + "register_scheduler_config", +] diff --git a/wall_x/config/data_config.py b/wall_x/config/data_config.py new file mode 100644 index 00000000..9ed7f94f --- /dev/null +++ b/wall_x/config/data_config.py @@ -0,0 +1,63 @@ +"""Public data config dataclasses. + +Only data backends shipped in the public package should define config classes +here. Internal backends register their dataclasses from their own packages via +``wall_x.config.registry.register_data_config``. +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, Optional + +from .registry import register_data_config + + +@dataclass +class DataConfig: + """Base fields shared by data backends. + + ``normalizer_config`` may contain: + - ``min_key``: stats key for the minimum value. + - ``delta_key``: stats key for the value range. + - ``customized_action_statistic_dof``: explicit action-stats JSON path. + """ + + dataset_type: str = "lerobot" + resolution: Dict[str, int] = field( + default_factory=lambda: { + "face_view": 256, + "left_wrist_view": 256, + "right_wrist_view": 256, + } + ) + train_test_split: float = 0.95 + normalizer_config: Optional[Dict[str, Any]] = None + + +@register_data_config("lerobot") +@dataclass +class LeRobotDataConfig(DataConfig): + """LeRobot data config. + + ``lerobot_config`` is expected to contain fields such as ``repo_id`` and + ``root`` for a HuggingFace LeRobot dataset. ``norm_stats_path`` points to + explicit action normalizer stats; the core package does not bundle private + defaults. + """ + + dataset_type: str = "lerobot" + lerobot_config: Optional[Dict[str, Any]] = None + key_mappings: Optional[Dict[str, Any]] = None + norm_stats_path: Optional[str] = None + priority_order: Optional[Dict[str, float]] = None + camera_name_mapping: Optional[Dict[str, str]] = None + num_workers: int = 4 + action_tokenizer_path: Optional[str] = None + use_fast_tokenizer: bool = False + padding_side: str = "left" + noise_scheduler: Optional[Dict[str, Any]] = None + + +__all__ = [ + "DataConfig", + "LeRobotDataConfig", +] diff --git a/wall_x/config/hyperparams_config.py b/wall_x/config/hyperparams_config.py new file mode 100644 index 00000000..10a57b68 --- /dev/null +++ b/wall_x/config/hyperparams_config.py @@ -0,0 +1,129 @@ +"""Training hyperparameter config dataclasses.""" + +from dataclasses import dataclass, field +from typing import List, Optional, Tuple + +from .registry import register_optimizer_config, register_scheduler_config + + +@dataclass +class LRGroupConfig: + """Named optimizer LR group matched by parameter-name substrings.""" + + name: str + lr: float + include: List[str] = field(default_factory=list) + fail_on_empty: bool = True + + +@dataclass +class OptimizerConfig: + """Base optimizer config. ``optimizer_type`` selects a registered subclass.""" + + optimizer_type: str = "adamw" + learning_rate: float = 1e-4 + max_grad_norm: float = 1.0 + enable_grad_clip: bool = True + # Named parameter groups with independent learning rates. Unmatched + # trainable parameters remain in the base group using ``learning_rate``. + lr_groups: Optional[List[LRGroupConfig]] = None + # Optional action-expert LR split. + train_action_expert_only: bool = False + action_expert_learning_rate: Optional[float] = None + action_lr_keywords: Optional[List[str]] = None + + +@register_optimizer_config("adamw") +@dataclass +class AdamWConfig(OptimizerConfig): + """AdamW optimizer config.""" + + optimizer_type: str = "adamw" + betas: Tuple[float, float] = (0.9, 0.98) + weight_decay: float = 1e-8 + eps: float = 1e-8 + fused: bool = True + foreach: Optional[bool] = None + + +@register_optimizer_config("dmuon") +@dataclass +class DMuonConfig(OptimizerConfig): + """DMuon optimizer config.""" + + optimizer_type: str = "dmuon" + muon_lr: float = 0.02 + momentum: float = 0.95 + ns_steps: int = 5 + muon_weight_decay: float = 0.0 + adamw_lr: float = 1e-3 + adamw_betas: Tuple[float, float] = (0.9, 0.999) + adamw_weight_decay: float = 0.01 + adamw_eps: float = 1e-8 + ns_backend: str = "gram" + ns_coefficients: str = "default" + nesterov: bool = True + + +@dataclass +class SchedulerConfig: + """Base scheduler config. ``scheduler_type`` selects a registered subclass.""" + + scheduler_type: str = "cosine" + # Optional training-step cap. When > 0, the trainer sets + # loss_guard_should_stop=True once global_step >= num_training_steps, + # regardless of scheduler type. Cosine reads it for its own decay + # horizon; constant / step schedulers use it only for the stop signal. + num_training_steps: int = 0 + + +@register_scheduler_config("cosine") +@dataclass +class CosineSchedulerConfig(SchedulerConfig): + """Cosine annealing with warmup.""" + + scheduler_type: str = "cosine" + num_warmup_steps: int = 0 + num_training_steps: int = 0 + min_lr: Optional[float] = None # None means 0.1 * learning_rate at runtime. + + +@register_scheduler_config("constant") +@dataclass +class ConstantSchedulerConfig(SchedulerConfig): + """Constant learning rate with no decay.""" + + scheduler_type: str = "constant" + + +@register_scheduler_config("step") +@dataclass +class StepSchedulerConfig(SchedulerConfig): + """Step decay scheduler.""" + + scheduler_type: str = "step" + step_size: int = 10000 + gamma: float = 0.1 + + +@dataclass +class TrainHyperParams: + num_epoch: int = 1 + batch_size_per_gpu: int = 1 + gradient_accumulation_steps: int = 1 + seed: int = 42 + optimizer: OptimizerConfig = field(default_factory=AdamWConfig) + scheduler: SchedulerConfig = field(default_factory=CosineSchedulerConfig) + + +__all__ = [ + "AdamWConfig", + "ConstantSchedulerConfig", + "CosineSchedulerConfig", + "DMuonConfig", + "LRGroupConfig", + "OptimizerConfig", + "SchedulerConfig", + "StepSchedulerConfig", + "TrainHyperParams", +] diff --git a/wall_x/config/infra_config.py b/wall_x/config/infra_config.py new file mode 100644 index 00000000..19853de5 --- /dev/null +++ b/wall_x/config/infra_config.py @@ -0,0 +1,98 @@ +"""Infrastructure config: distributed runtime, logging, checkpoints, and debug.""" + +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class DistributedConfig: + # FSDP + use_fsdp: bool = False + fsdp_sharding_strategy: str = "full_shard" + fsdp_backward_prefetch: str = "backward_pre" + fsdp_cpu_offload: bool = False + fsdp_use_orig_params: bool = True + fsdp_limit_all_gathers: bool = True + fsdp_forward_prefetch: bool = False + fsdp_sync_module_states: bool = True + fsdp_save_policy: str = "full" + fsdp_hsdp_replicate_size: Optional[int] = None + # Mixed precision + use_mixed_precision: bool = True + bf16: bool = True + fsdp_reduce_dtype: str = "bf16" + use_amp: bool = False + use_gradient_checkpointing: bool = False + use_gradient_checkpointing_offload: bool = False + use_selective_recompute: bool = False + # DDP fallback + find_unused_parameters: bool = False + broadcast_buffers: bool = True + bucket_cap_mb: int = 25 + + +@dataclass +class LoggingConfig: + log_name: str = "exp" + log_project: str = "wallx" + log_entity: Optional[str] = None + use_wandb: bool = True + wandb_offline: bool = False + log_interval: int = 1 + save_interval: int = 1000 + val_interval: int = 4000 + epoch_save_interval: int = 1 + gc_interval_steps: int = 1000 + ignore_until_interval: int = 0 + # Rolling-window for smoothing per-step training metrics displayed on the + # console (and reused by tqdm). 1 = raw per-step (historical behavior). + # 10 = DZ-style 10-step rolling average — diffusion losses are dominated + # by timestep-sampling noise per step; the smoothing only changes display, + # not training. Independent of log_interval (which buffers for wandb). + loss_log_smooth_window: int = 1 + + +@dataclass +class CheckpointConfig: + """Checkpoint save and resume options. + + ``resume_from`` is shorthand for setting every component-specific resume + path to the same checkpoint. Component-specific fields take precedence. + """ + + save_path: str = "./ckpt" + validate_first: bool = False + # Shorthand path for all components. + resume_from: Optional[str] = None + # Component-specific resume paths. + resume_model: Optional[str] = None + resume_optimizer: Optional[str] = None + resume_scheduler: Optional[str] = None + resume_ema: Optional[str] = None + resume_rng: Optional[str] = None + resume_data: Optional[str] = None + resume_step: Optional[str] = None + + def get_resume_path(self, component: str) -> Optional[str]: + """Return the resume path for one checkpoint component.""" + specific = getattr(self, f"resume_{component}", None) + if specific is not None: + return specific + return self.resume_from + + +@dataclass +class DebugConfig: + profile: bool = False + profile_save_path: str = "./profile" + profile_wait_iters: int = 1 + profile_warmup_iters: int = 1 + profile_active_iters: int = 3 + show_time_details: bool = False + visualize_sample: bool = False + save_debug_batch_path: Optional[str] = None + nvtx: bool = False + # Formula MFU uses the local FLOPs estimate and measured step time. + enable_mfu: bool = False + # Optional FLOPs profiling runs an extra forward at step 0. + enable_mfu_profile: bool = False diff --git a/wall_x/config/loader.py b/wall_x/config/loader.py new file mode 100644 index 00000000..c22bcd5b --- /dev/null +++ b/wall_x/config/loader.py @@ -0,0 +1,285 @@ +"""Wall-X config loader.""" + +import dataclasses +import os +import shutil +from typing import Any, Type, TypeVar + +import yaml + +from .data_config import LeRobotDataConfig +from .hyperparams_config import ( + AdamWConfig, + LRGroupConfig, + OptimizerConfig, + SchedulerConfig, + TrainHyperParams, +) +from .infra_config import ( + CheckpointConfig, + DebugConfig, + DistributedConfig, + LoggingConfig, +) +from .model_config import ModelConfig, QActModelConfig +from .registry import ( + get_data_config, + get_model_config, + get_optimizer_config, + get_scheduler_config, + registered_data_configs, + registered_model_configs, + registered_optimizer_configs, + registered_scheduler_configs, +) +from .task_config import TaskConfig +from .train_config import TrainConfig + +T = TypeVar("T") + +_CONFIG_PLUGINS_LOADED = False + + +class _TrainConfigSafeLoader(yaml.SafeLoader): + pass + + +def _construct_python_tuple(loader: yaml.SafeLoader, node: yaml.Node) -> tuple: + return tuple(loader.construct_sequence(node)) + + +_TrainConfigSafeLoader.add_constructor( + "tag:yaml.org,2002:python/tuple", _construct_python_tuple +) + + +def _ensure_config_plugins_loaded() -> None: + """Load optional internal config plugins when they are present.""" + global _CONFIG_PLUGINS_LOADED + if _CONFIG_PLUGINS_LOADED: + return + _CONFIG_PLUGINS_LOADED = True + try: + from .internal_plugins import register_internal_config_plugins + except ImportError: + return + register_internal_config_plugins() + + +def load_config(config_path: str, cli_args: Any = None) -> TrainConfig: + """Load a ``TrainConfig`` from a YAML file.""" + _ensure_config_plugins_loaded() + + with open(config_path, "r") as f: + raw = yaml.load(f, Loader=_TrainConfigSafeLoader) + + if raw is None: + raise ValueError(f"Config file is empty: {config_path}") + + model_type = raw.get("model_type") + if model_type is None: + raise ValueError(f"Missing required field 'model_type' in {config_path}") + + config = TrainConfig( + model_type=model_type, + task=_build_dataclass(TaskConfig, raw.get("task", {})), + model=_build_model_config(model_type, raw.get("model", {})), + data=_build_data_config(raw.get("data", {})), + hyperparams=_build_hyperparams(raw.get("hyperparams", {})), + distributed=_build_dataclass(DistributedConfig, raw.get("distributed", {})), + logging=_build_dataclass(LoggingConfig, raw.get("logging", {})), + checkpoint=_build_dataclass(CheckpointConfig, raw.get("checkpoint", {})), + debug=_build_dataclass(DebugConfig, raw.get("debug", {})), + _raw_data=raw.get("data", {}), + _raw_yaml=raw, + dataset_path=raw.get("dataset_path"), + ) + + if cli_args is not None: + _apply_cli_overrides(config, cli_args) + + _validate(config) + + # Register the active data backend now that cfg is fully resolved + # (post-CLI-override, post-validate). This is the single source of + # truth for "which dataset backend is this run using". + from wall_x.data._registry import _set_data_backend + + _set_data_backend(config.data.dataset_type) + + return config + + +def save_config(config: TrainConfig, save_dir: str) -> str: + """Save ``TrainConfig`` to ``config.yml`` under ``save_dir``.""" + os.makedirs(save_dir, exist_ok=True) + config_path = os.path.join(save_dir, "config.yml") + + data = _sanitize_for_yaml(dataclasses.asdict(config)) + with open(config_path, "w") as f: + yaml.dump( + data, f, default_flow_style=False, allow_unicode=True, sort_keys=False + ) + + dataset_config_path = getattr(config.data, "dataset_config_path", None) + if dataset_config_path and os.path.exists(dataset_config_path): + dst = os.path.join(save_dir, "dataset_config.yml") + shutil.copy(dataset_config_path, dst) + + return config_path + + +def _sanitize_for_yaml(obj: Any) -> Any: + """Convert dataclass output to YAML-safe containers.""" + if isinstance(obj, dict): + return {k: _sanitize_for_yaml(v) for k, v in obj.items()} + elif isinstance(obj, (list, tuple)): + return [_sanitize_for_yaml(v) for v in obj] + return obj + + +def _build_dataclass(cls: Type[T], raw: dict) -> T: + """Build a dataclass from a dict, ignoring unknown keys.""" + if not raw: + return cls() + + field_names = {f.name for f in dataclasses.fields(cls)} + field_types = {f.name: f.type for f in dataclasses.fields(cls)} + filtered = {} + + for k, v in raw.items(): + if k not in field_names: + continue + ft = field_types[k] + if ( + isinstance(ft, type) + and dataclasses.is_dataclass(ft) + and isinstance(v, dict) + ): + filtered[k] = _build_dataclass(ft, v) + else: + filtered[k] = v + + return cls(**filtered) + + +def _build_model_config(model_type: str, raw: dict) -> ModelConfig: + """Build the registered model config for ``model_type``.""" + cls = get_model_config(model_type) + if cls is None: + raise ValueError( + f"Unknown model_type: '{model_type}'. " + f"Supported: {registered_model_configs()}" + ) + return _build_dataclass(cls, raw) + + +def _build_data_config(raw: dict): + """Build the registered data config for ``dataset_type``.""" + if not raw: + return LeRobotDataConfig() + + dataset_type = raw.get("dataset_type", "lerobot") + cls = get_data_config(dataset_type) + if cls is None: + raise ValueError( + f"Unknown dataset_type: '{dataset_type}'. " + f"Supported: {registered_data_configs()}" + ) + + return _build_dataclass(cls, raw) + + +def _build_optimizer_config(raw: dict) -> OptimizerConfig: + """Build the registered optimizer config for ``optimizer_type``.""" + if not raw: + return AdamWConfig() + + raw = dict(raw) + optimizer_type = raw.get("optimizer_type", "adamw") + cls = get_optimizer_config(optimizer_type) + if cls is None: + raise ValueError( + f"Unknown optimizer_type: '{optimizer_type}'. " + f"Supported: {registered_optimizer_configs()}" + ) + + if "betas" in raw and isinstance(raw["betas"], list): + raw["betas"] = tuple(raw["betas"]) + if "adamw_betas" in raw and isinstance(raw["adamw_betas"], list): + raw["adamw_betas"] = tuple(raw["adamw_betas"]) + if raw.get("lr_groups") is not None: + raw["lr_groups"] = [ + _build_dataclass(LRGroupConfig, group) for group in raw["lr_groups"] + ] + + return _build_dataclass(cls, raw) + + +def _build_scheduler_config(raw: dict) -> SchedulerConfig: + """Build the registered scheduler config for ``scheduler_type``.""" + if not raw: + cls = get_scheduler_config("cosine") + if cls is None: + raise ValueError("Scheduler config 'cosine' is not registered") + return cls() + + scheduler_type = raw.get("scheduler_type", "cosine") + cls = get_scheduler_config(scheduler_type) + if cls is None: + raise ValueError( + f"Unknown scheduler_type: '{scheduler_type}'. " + f"Supported: {registered_scheduler_configs()}" + ) + return _build_dataclass(cls, raw) + + +def _build_hyperparams(raw: dict) -> TrainHyperParams: + """Build ``TrainHyperParams`` with polymorphic optimizer/scheduler config.""" + if not raw: + return TrainHyperParams() + + raw = dict(raw) # shallow copy to avoid mutating caller's dict + optimizer_raw = raw.pop("optimizer", {}) + scheduler_raw = raw.pop("scheduler", {}) + + optimizer = _build_optimizer_config(optimizer_raw) + scheduler = _build_scheduler_config(scheduler_raw) + + hp = _build_dataclass(TrainHyperParams, raw) + hp.optimizer = optimizer + hp.scheduler = scheduler + return hp + + +def _apply_cli_overrides(config: TrainConfig, args: Any) -> None: + """Apply CLI overrides to the loaded config.""" + if getattr(args, "fsdp_sharding_strategy", None) is not None: + config.distributed.fsdp_sharding_strategy = args.fsdp_sharding_strategy + + if getattr(args, "debug", False): + config.logging.log_name = "debug" + config.logging.log_project = "debug" + config.checkpoint.save_path = "./ckpt/debug" + + if getattr(args, "visualize", False): + config.debug.visualize_sample = True + + if getattr(args, "wandb_offline", None) is not None: + config.logging.wandb_offline = args.wandb_offline in ("true", "True", "1") + + +def _validate(config: TrainConfig) -> None: + """Validate required fields and model-specific config.""" + if not config.model_type: + raise ValueError("model_type is required") + + if not config.task.dof_config: + raise ValueError("task.dof_config is required (cannot be empty)") + + if isinstance(config.model, QActModelConfig): + model = config.model + if not model.config_path: + raise ValueError("model.config_path is required for QAct models") + if not model.processor_path: + raise ValueError("model.processor_path is required for QAct models") diff --git a/wall_x/config/model_config.py b/wall_x/config/model_config.py new file mode 100644 index 00000000..4edaf445 --- /dev/null +++ b/wall_x/config/model_config.py @@ -0,0 +1,42 @@ +"""Base model config dataclasses. + +Optional model variants register their config dataclasses from their own +packages via ``wall_x.config.registry.register_model_config``. +""" + +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +from .registry import register_model_config + + +@dataclass +class ModelConfig: + """Base model architecture fields.""" + + use_ema: bool = False + attn_implementation: Optional[str] = None + attn_deterministic: Optional[bool] = None + ar_loss_weight: float = 1.0 + + +@register_model_config("qwen2_5") +@dataclass +class QActModelConfig(ModelConfig): + """QAct model config for the Qwen2.5 VLA path.""" + + config_path: str = "" + processor_path: str = "" + pretrained_path: Optional[str] = None + backbone: str = "qwen2_5" + action_tokenizer_type: Optional[str] = None + action_tokenizer_path: Optional[str] = None + action_tokenizer_checkpoint_path: Optional[str] = None + action_tokenizer_config_dir: Optional[str] = None + new_special_tokens: Optional[List[str]] = None + flow_loss_weight: float = 1.0 + enable_customized_robot_config: bool = False + customized_robot_config: Optional[Dict[str, Any]] = None + + +__all__ = ["ModelConfig", "QActModelConfig"] diff --git a/wall_x/config/registry.py b/wall_x/config/registry.py new file mode 100644 index 00000000..3ecfdb25 --- /dev/null +++ b/wall_x/config/registry.py @@ -0,0 +1,100 @@ +"""Runtime registries for typed config variants. + +Core Wall-X keeps only base config types in ``wall_x.config``. Optional +datasets, model families, and optimizer variants register their dataclasses +from their own packages, so trimmed distributions do not leave dead config +entries behind. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, TypeVar + + +T = TypeVar("T") + +_DATA_CONFIG_REGISTRY: dict[str, type[Any]] = {} +_MODEL_CONFIG_REGISTRY: dict[str, type[Any]] = {} +_OPTIMIZER_CONFIG_REGISTRY: dict[str, type[Any]] = {} +_SCHEDULER_CONFIG_REGISTRY: dict[str, type[Any]] = {} + + +def _register( + registry: dict[str, type[Any]], names: tuple[str, ...], cls: type[T] +) -> type[T]: + for name in names: + if not name: + raise ValueError("Config registry name must be non-empty") + existing = registry.get(name) + if existing is not None and existing is not cls: + raise ValueError( + f"Config name {name!r} is already registered by " + f"{existing.__module__}.{existing.__name__}" + ) + registry[name] = cls + return cls + + +def register_data_config(*names: str) -> Callable[[type[T]], type[T]]: + return lambda cls: _register(_DATA_CONFIG_REGISTRY, names, cls) + + +def register_model_config(*names: str) -> Callable[[type[T]], type[T]]: + return lambda cls: _register(_MODEL_CONFIG_REGISTRY, names, cls) + + +def register_optimizer_config(*names: str) -> Callable[[type[T]], type[T]]: + return lambda cls: _register(_OPTIMIZER_CONFIG_REGISTRY, names, cls) + + +def register_scheduler_config(*names: str) -> Callable[[type[T]], type[T]]: + return lambda cls: _register(_SCHEDULER_CONFIG_REGISTRY, names, cls) + + +def get_data_config(name: str) -> type[Any] | None: + return _DATA_CONFIG_REGISTRY.get(name) + + +def get_model_config(name: str) -> type[Any] | None: + return _MODEL_CONFIG_REGISTRY.get(name) + + +def get_optimizer_config(name: str) -> type[Any] | None: + return _OPTIMIZER_CONFIG_REGISTRY.get(name) + + +def get_scheduler_config(name: str) -> type[Any] | None: + return _SCHEDULER_CONFIG_REGISTRY.get(name) + + +def registered_data_configs() -> list[str]: + return sorted(_DATA_CONFIG_REGISTRY) + + +def registered_model_configs() -> list[str]: + return sorted(_MODEL_CONFIG_REGISTRY) + + +def registered_optimizer_configs() -> list[str]: + return sorted(_OPTIMIZER_CONFIG_REGISTRY) + + +def registered_scheduler_configs() -> list[str]: + return sorted(_SCHEDULER_CONFIG_REGISTRY) + + +__all__ = [ + "get_data_config", + "get_model_config", + "get_optimizer_config", + "get_scheduler_config", + "register_data_config", + "register_model_config", + "register_optimizer_config", + "register_scheduler_config", + "registered_data_configs", + "registered_model_configs", + "registered_optimizer_configs", + "registered_scheduler_configs", +] diff --git a/wall_x/config/task_config.py b/wall_x/config/task_config.py new file mode 100644 index 00000000..1cd02b82 --- /dev/null +++ b/wall_x/config/task_config.py @@ -0,0 +1,17 @@ +"""Task config shared by model and data code.""" + +from dataclasses import dataclass, field +from typing import Any, Dict, Optional + + +@dataclass +class TaskConfig: + """Robot task definition shared by model construction and data slicing.""" + + dof_config: Dict[str, int] = field(default_factory=dict) + agent_pos_config: Dict[str, int] = field(default_factory=dict) + ar_dof_config: Optional[Dict[str, int]] = None + action_horizon: int = 32 + action_horizon_flow: int = 32 + noise_scheduler: Optional[Dict[str, Any]] = None + use_state_string_representation: bool = False diff --git a/wall_x/config/train_config.py b/wall_x/config/train_config.py new file mode 100644 index 00000000..979622e3 --- /dev/null +++ b/wall_x/config/train_config.py @@ -0,0 +1,81 @@ +"""Top-level training config.""" + +import dataclasses +from dataclasses import dataclass, field +from typing import Any, Dict + +from .data_config import DataConfig +from .hyperparams_config import TrainHyperParams +from .infra_config import ( + CheckpointConfig, + DebugConfig, + DistributedConfig, + LoggingConfig, +) +from .model_config import ModelConfig +from .task_config import TaskConfig + + +@dataclass +class TrainConfig: + """Top-level Wall-X training config.""" + + model_type: str = "qwen2_5" + task: TaskConfig = field(default_factory=TaskConfig) + model: ModelConfig = field(default_factory=ModelConfig) + data: DataConfig = field(default_factory=DataConfig) + hyperparams: TrainHyperParams = field(default_factory=TrainHyperParams) + distributed: DistributedConfig = field(default_factory=DistributedConfig) + logging: LoggingConfig = field(default_factory=LoggingConfig) + checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig) + debug: DebugConfig = field(default_factory=DebugConfig) + # Raw YAML sections preserved verbatim for backend APIs that read fields + # not captured by the typed DataConfig dataclass. + _raw_data: Dict[str, Any] = field(default_factory=dict) + # Full raw YAML dict for backend-specific compatibility paths. + _raw_yaml: Dict[str, Any] = field(default_factory=dict) + # dataset_path lives at top level in YAML, consumed by data loaders directly + dataset_path: Any = None + + @property + def action_dim(self) -> int: + return sum(self.task.dof_config.values()) + + @property + def propri_dim(self) -> int: + return sum(self.task.agent_pos_config.values()) + + def build_data_loader_dict(self) -> Dict[str, Any]: + """Build the raw dict consumed by backend-specific compatibility paths. + + Merges ``_raw_data`` (verbatim YAML ``data:`` section) with task + fields (dof_config, action_horizon, etc.) and hyperparams + (batch_size). Backend compatibility APIs may read fields that the + typed DataConfig dataclass does not carry. + + This keeps legacy flat configs working while the main config surface + stays typed. + """ + # Start with the raw YAML data section, then add typed task defaults. + data_dict = dict(self._raw_data) + task_dict = dataclasses.asdict(self.task) + for key in ( + "dof_config", + "agent_pos_config", + "action_horizon", + "action_horizon_flow", + "ar_dof_config", + "use_state_string_representation", + ): + if key in task_dict and task_dict[key] is not None: + data_dict.setdefault(key, task_dict[key]) + data_dict.setdefault("batch_size_per_gpu", self.hyperparams.batch_size_per_gpu) + data_dict.setdefault("batch_size", self.hyperparams.batch_size_per_gpu) + result: Dict[str, Any] = { + "model_type": self.model_type, + "data": data_dict, + **task_dict, + } + if self.dataset_path is not None: + result["dataset_path"] = self.dataset_path + return result diff --git a/wall_x/data/__init__.py b/wall_x/data/__init__.py index e69de29b..bd53c4f0 100644 --- a/wall_x/data/__init__.py +++ b/wall_x/data/__init__.py @@ -0,0 +1,31 @@ +"""Public data facade for backend selection and dataset construction.""" + +from wall_x.data._bundle import DataBundle +from wall_x.data._protocol import BuildContext, DatasetBackend +from wall_x.data._registry import ( + MissingOperationError, + data_backend, + available_backends, + backend_for, + build_data, + has_data_backend, + register, + register_module, + resolve_dataset_type, +) +from wall_x.data import backends # noqa: F401 (side-effect registration) + +__all__ = [ + "DataBundle", + "BuildContext", + "DatasetBackend", + "MissingOperationError", + "data_backend", + "available_backends", + "backend_for", + "build_data", + "has_data_backend", + "register", + "register_module", + "resolve_dataset_type", +] diff --git a/wall_x/data/_bundle.py b/wall_x/data/_bundle.py new file mode 100644 index 00000000..9a250aa1 --- /dev/null +++ b/wall_x/data/_bundle.py @@ -0,0 +1,39 @@ +"""Dataset bundle returned by data backends.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Iterable, Optional + + +def _noop_set_epoch(_epoch: int) -> None: + """Default ``set_epoch`` for backends with epoch-agnostic shuffling.""" + return None + + +@dataclass +class DataBundle: + """Container returned by every backend ``build()``. + + Attributes: + dataset: backend-private; trainer code should treat it as opaque. + train_loader: anything iterable that yields training batches. + val_loader: val iterable or None if the backend does not split val. + train_iters: one-epoch step count for the train loader. Backends + with dynamic resizing should set this to a stable snapshot + and expose the live value separately on ``dataset``. + val_iters: one-epoch step count for the val loader; 0 if no val. + set_epoch: per-epoch seed hook. Called before each epoch by + the trainer. Backends that don't need per-epoch reshuffle + should use ``_noop_set_epoch``. + """ + + dataset: Any + train_loader: Iterable + val_loader: Optional[Iterable] = None + train_iters: int = 0 + val_iters: int = 0 + set_epoch: Callable[[int], None] = field(default=_noop_set_epoch) + + +__all__ = ["DataBundle"] diff --git a/wall_x/data/_protocol.py b/wall_x/data/_protocol.py new file mode 100644 index 00000000..50694eb6 --- /dev/null +++ b/wall_x/data/_protocol.py @@ -0,0 +1,37 @@ +"""Shared data backend protocol definitions.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional, Protocol, runtime_checkable + +from wall_x.data._bundle import DataBundle + + +@dataclass +class BuildContext: + """Shared runtime state passed to every backend ``build()``. + + Fields are Optional so backends not using a particular piece can + simply leave it ``None``. The trainer populates whatever it has. + """ + + rank: int = 0 + world_size: int = 1 + tokenizer: Optional[Any] = None + processor: Optional[Any] = None + tokenizer_mixin: Optional[Any] = None + normalizer_action: Optional[Any] = None + normalizer_propri: Optional[Any] = None + model_config: Optional[Any] = None + resume_state: Optional[dict] = None + + +@runtime_checkable +class DatasetBackend(Protocol): + """Callable every backend registers under its ``dataset_type`` name.""" + + def __call__(self, cfg: Any, ctx: BuildContext) -> DataBundle: ... + + +__all__ = ["BuildContext", "DatasetBackend"] diff --git a/wall_x/data/_registry.py b/wall_x/data/_registry.py new file mode 100644 index 00000000..5e036e46 --- /dev/null +++ b/wall_x/data/_registry.py @@ -0,0 +1,263 @@ +"""Data backend registry + multi-verb dispatch. + +Each backend registers an entire **module** (not just a single function). +Consumers obtain a backend via :func:`backend_for` and then call any +operation the backend supports (``backend.build`` or optional helper +operations published by that backend). Operations a backend does *not* +implement raise +:class:`MissingOperationError` with a list of backends that do. + +Three failure shapes, each pointing at the right next step: + +- **Unknown name** (typo / never registered): :class:`KeyError` + ``Unknown backend 'foo_bar'. Known: [...]`` +- **Known but failed to import** (optional dependency missing): + :class:`RuntimeError` + ``Backend 'example' failed to import: . Install its + dependency or switch to one of [...]`` +- **Backend exists but operation missing**: :class:`MissingOperationError` + ``Backend 'example' does not implement 'make_processor'. + Supported by: ['other_backend']`` +""" + +from __future__ import annotations + +import logging +from types import ModuleType +from typing import Any, Dict + +from wall_x.data._bundle import DataBundle +from wall_x.data._protocol import BuildContext + +logger = logging.getLogger(__name__) + +_BACKENDS: Dict[str, ModuleType] = {} +_import_errors: Dict[str, BaseException] = {} + +# The single, process-wide active backend. Set exactly once at config load +# time by :func:`_set_data_backend`; all consumer code reads it via +# :func:`data_backend`. Keeping this as module-level state (rather than +# threading cfg through every callsite) is the whole point of the design. +_DATA_BACKEND: str | None = None + + +# --- Registration -------------------------------------------------------- + + +def register_module(dataset_type: str, module: ModuleType) -> None: + """Register a backend module under ``dataset_type``. + + The module must expose a ``build(cfg, ctx) -> DataBundle`` callable. + Other optional operations are + discovered lazily via ``getattr`` when ``backend_for(name).`` is + accessed; backends only declare what they support. + """ + if not hasattr(module, "build"): + raise TypeError( + f"Backend module for {dataset_type!r} must expose a " + f"``build(cfg, ctx) -> DataBundle`` callable; got {module!r}." + ) + if dataset_type in _BACKENDS: + logger.warning("overwriting existing backend registration for %r", dataset_type) + _BACKENDS[dataset_type] = module + + +def register(dataset_type: str, build_callable) -> None: + """Legacy single-callable registration. Wraps ``build_callable`` in a + minimal module so the new dispatch path still works. + + New backends should prefer :func:`register_module` so they can publish + multiple operations. + """ + shim = ModuleType(f"_legacy_backend_shim_{dataset_type}") + shim.build = build_callable # type: ignore[attr-defined] + register_module(dataset_type, shim) + + +def record_import_error(dataset_type: str, error: BaseException) -> None: + """Stash the exception that prevented a backend from registering.""" + _import_errors[dataset_type] = error + + +def available_backends() -> list[str]: + """Return the names of currently-registered backends.""" + return sorted(_BACKENDS) + + +# --- Lookup -------------------------------------------------------------- + + +class MissingOperationError(NotImplementedError): + """Raised when a backend module does not implement a requested op.""" + + def __init__(self, backend_name: str, op_name: str) -> None: + impls = [n for n, mod in _BACKENDS.items() if hasattr(mod, op_name)] + if impls: + hint = f"Supported by: {impls}." + else: + hint = ( + f"No registered backend implements {op_name!r} — check the " + f"spelling or add it to the backend module." + ) + super().__init__( + f"Backend {backend_name!r} does not implement {op_name!r}. {hint}" + ) + self.backend_name = backend_name + self.op_name = op_name + + +class _BackendProxy: + """Thin wrapper that forwards ``proxy.(...)`` to the backend module. + + Wrapping (instead of returning the module directly) lets us emit + ``MissingOperationError`` with a useful "supported by" hint instead + of plain ``AttributeError``. + """ + + __slots__ = ("_name", "_module") + + def __init__(self, name: str, module: ModuleType) -> None: + self._name = name + self._module = module + + def __repr__(self) -> str: + return f"" + + def __getattr__(self, op: str): + attr = getattr(self._module, op, None) + if attr is None: + raise MissingOperationError(self._name, op) + return attr + + def supports(self, op: str) -> bool: + """Cheap predicate: does this backend implement ``op``?""" + return hasattr(self._module, op) + + +def backend_for(name: str) -> _BackendProxy: + """Look up a backend by ``dataset_type`` name. + + Raises :class:`RuntimeError` (with chained ImportError) if the + backend is known but failed to import; :class:`KeyError` if the + name was never registered. + """ + if name in _BACKENDS: + return _BackendProxy(name, _BACKENDS[name]) + if name in _import_errors: + err = _import_errors[name] + raise RuntimeError( + f"Backend {name!r} failed to import: {err}. " + f"Install its dependency or switch to one of " + f"{available_backends()}." + ) from err + raise KeyError(f"Unknown backend {name!r}. Known: {available_backends()}.") + + +# --- Convenience helpers ------------------------------------------------- + + +def resolve_dataset_type(cfg_or_yaml: Any, default: str = "lerobot") -> str: + """Pull ``dataset_type`` out of either a typed TrainConfig or a raw yaml dict. + + Lookup order: + 1. ``cfg.data.dataset_type`` (typed TrainConfig) + 2. ``yaml_dict["dataset_type"]`` (legacy flat yaml) + 3. ``default`` + """ + data = getattr(cfg_or_yaml, "data", None) + if data is not None: + v = getattr(data, "dataset_type", None) + if v: + return v + if isinstance(cfg_or_yaml, dict): + v = cfg_or_yaml.get("dataset_type") + if v: + return v + return default + + +def build_data(cfg: Any, ctx: BuildContext) -> DataBundle: + """Dispatch ``backend.build(cfg, ctx)`` to the backend named by cfg. + + ``cfg`` is expected to be a typed :class:`TrainConfig` — that's what + ``wall_x.config.load_config`` returns and what every trainer entry + point passes in. Raw yaml dicts are not supported here on purpose: + the typed schema is the contract that gives backends their + ``cfg.data.dataset_type`` access. If you have a raw dict, use + :func:`resolve_dataset_type` + :func:`backend_for` directly. + """ + if not hasattr(cfg, "data") or not hasattr(cfg.data, "dataset_type"): + raise TypeError( + f"build_data() expects a typed TrainConfig, got {type(cfg).__name__}. " + f"Load via wall_x.config.load_config() instead of passing a raw dict." + ) + return backend_for(cfg.data.dataset_type).build(cfg, ctx) + + +# --- Active backend (process-global) ------------------------------------ + + +def _set_data_backend(name: str) -> None: + """Internal — only config loaders should call this. + + Strict semantics: first call wins; same-value re-set is a no-op; + different-value re-set raises. The error is loud on purpose — it means + two config loaders disagreed about which backend to use, which is + almost always a bug (e.g. business code calling this directly, or two + yamls being loaded in the same process). + """ + global _DATA_BACKEND + if _DATA_BACKEND is None: + _DATA_BACKEND = name + return + if _DATA_BACKEND == name: + return # idempotent on same value + raise RuntimeError( + f"Active backend already set to {_DATA_BACKEND!r}, refusing to " + f"overwrite with {name!r}. This usually means a config loader was " + f"called twice with different dataset_type, or business code called " + f"_set_data_backend directly. Use _reset_data_backend() in tests " + f"if you need to switch." + ) + + +def _reset_data_backend() -> None: + """Internal — clear the active backend. Tests only.""" + global _DATA_BACKEND + _DATA_BACKEND = None + + +def has_data_backend() -> bool: + """Whether a backend has been registered for this process.""" + return _DATA_BACKEND is not None + + +def data_backend() -> _BackendProxy: + """Return the active backend proxy. Raises if no config has been loaded. + + Direct construction of ``TrainConfig(...)`` does **not** register a + backend — that's intentional. Business code must go through + :func:`wall_x.config.load_config`, which registers the backend before + returning. + """ + if _DATA_BACKEND is None: + raise RuntimeError( + "No active data backend. Load a TrainConfig via " + "wall_x.config.load_config() before accessing backend verbs. " + "Direct TrainConfig() construction does not register a backend." + ) + return backend_for(_DATA_BACKEND) + + +__all__ = [ + "register", + "register_module", + "record_import_error", + "available_backends", + "backend_for", + "build_data", + "resolve_dataset_type", + "MissingOperationError", + "data_backend", + "has_data_backend", +] diff --git a/wall_x/data/backends/__init__.py b/wall_x/data/backends/__init__.py new file mode 100644 index 00000000..5b9a1287 --- /dev/null +++ b/wall_x/data/backends/__init__.py @@ -0,0 +1,16 @@ +"""Data backend registration. + +Importing this package registers the default shipped backends. Additional +backends are loaded by plugin modules so their names and dependencies do not +have to appear in the default import path. +""" + +import importlib + +from wall_x.data._registry import record_import_error + +for _name in ("lerobot",): + try: + importlib.import_module(f"wall_x.data.backends.{_name}") + except ImportError as _e: + record_import_error(_name, _e) diff --git a/wall_x/data/backends/lerobot/__init__.py b/wall_x/data/backends/lerobot/__init__.py new file mode 100644 index 00000000..d85849bb --- /dev/null +++ b/wall_x/data/backends/lerobot/__init__.py @@ -0,0 +1,26 @@ +"""LeRobot backend registration.""" + +from __future__ import annotations + +import sys + +try: + from wall_x.data._registry import register_module + from wall_x.data.backends.lerobot.build import ( + build, + load_trainer_data_config, + load_trainer_data_config_from_yaml_dict, + ) + + register_module("lerobot", sys.modules[__name__]) +except ImportError as _e: + from wall_x.data._registry import record_import_error + + record_import_error("lerobot", _e) + + +__all__ = [ + "build", + "load_trainer_data_config", + "load_trainer_data_config_from_yaml_dict", +] diff --git a/wall_x/data/backends/lerobot/build.py b/wall_x/data/backends/lerobot/build.py new file mode 100644 index 00000000..2badf51e --- /dev/null +++ b/wall_x/data/backends/lerobot/build.py @@ -0,0 +1,337 @@ +"""LeRobot data loading bridge for typed training configs.""" + +from __future__ import annotations + +import logging +import multiprocessing as mp +from typing import Any, Dict, Tuple + +import torch +import torch.distributed as dist + +from wall_x.data.backends.lerobot.config import LerobotConfig +from wall_x.data.backends.lerobot.utils import load_norm_stats +from wall_x.model.core.action.normalizer import ( + create_normalizers_from_lerobot_norm_stats, +) + +logger = logging.getLogger(__name__) + + +def load_lerobot_normalizers(cfg: Any): + """Create model normalizers from the LeRobot norm stats configured for a run.""" + data = cfg.data + + norm_stats_path = getattr(data, "norm_stats_path", None) + if not norm_stats_path: + return None + + key_mappings = getattr(data, "key_mappings", None) + if not key_mappings: + raise ValueError( + "LeRobot normalizer from norm_stats_path requires data.key_mappings" + ) + + lerobot_config = getattr(data, "lerobot_config", None) + if not isinstance(lerobot_config, dict) or not lerobot_config.get("repo_id"): + raise ValueError( + "LeRobot normalizer from norm_stats_path requires " + "data.lerobot_config.repo_id" + ) + + dataset_name = str(lerobot_config["repo_id"]) + norm_stats = load_norm_stats( + norm_stats_path, + key_mappings, + dof_config=dict(cfg.task.dof_config or {}), + agent_pos_config=dict(cfg.task.agent_pos_config or {}), + ) + normalizer_action, normalizer_propri = create_normalizers_from_lerobot_norm_stats( + norm_stats, + dataset_name, + cfg.action_dim, + cfg.propri_dim, + ) + return normalizer_action, normalizer_propri, norm_stats_path, dataset_name + + +class _LerobotDatasetWrapper: + """Trainer-facing wrapper aligning PreprocessedDataset with v1 API. + + PreprocessedDataset internally switches ``self._dataset`` between + its train/val splits via ``_train()`` / ``_eval()``. Its + ``get_train_dataloader`` / ``get_val_dataloader`` return + ``(dataloader, sampler)`` tuples and no-argument calls are supported + (they read rank/world_size/seed from the inner object itself). + + This wrapper: + - Caches the rebuilt train dataloader / sampler so + ``set_epoch(epoch)`` can reset shuffling per-epoch. + - Owns the val dataloader so the trainer's ``val_loop`` can do + ``self.dataset.get_val_dataloader()`` and iterate directly (matching + what the v1/v2 wrappers return). + """ + + def __init__( + self, + inner, + train_dataloader: torch.utils.data.DataLoader, + train_sampler, + train_num: int, + val_dataloader: torch.utils.data.DataLoader = None, + val_num: int = 0, + ): + self._inner = inner + self._train_dataloader = train_dataloader + self._train_sampler = train_sampler + self._train_num = train_num + self._val_dataloader = val_dataloader + self.global_train_iters = mp.Value("i", train_num) + self.global_val_iters = mp.Value("i", val_num) + + def __len__(self) -> int: + return self._train_num + + def _activate_train_split(self) -> None: + if hasattr(self._inner, "_train"): + self._inner._train() + + def get_train_dataloader(self): + self._activate_train_split() + return self._train_dataloader + + def get_val_dataloader(self): + # PreprocessedDataset shares one ``_dataset`` pointer between its + # train and val splits (flipped by ``_train()`` / ``_eval()``). + # The val DataLoader's DistributedSampler caches total_size sized + # to the val split but ``__iter__`` reads ``len(self.dataset)`` + # live — if a preceding train_loop left the pointer at train, that + # live len is ~20× total_size and DistributedSampler asserts. + # Rebuild each time so ``_eval()`` runs and a fresh sampler is + # snapped to the current (val) split length. Mirrors the train-side + # rebuild-on-every-epoch pattern. + if self._val_dataloader is None: + return None + self._val_dataloader, _ = self._inner.get_val_dataloader() + return self._val_dataloader + + def set_epoch(self, epoch: int) -> None: + """Seed the per-epoch shuffle in the train DistributedSampler.""" + self._activate_train_split() + if self._train_sampler is not None and hasattr( + self._train_sampler, "set_epoch" + ): + self._train_sampler.set_epoch(epoch) + + +def load_trainer_data_config(cfg: Any) -> LerobotConfig: + """Build the inference/trainer data config from a typed TrainConfig.""" + raw_yaml = dict(getattr(cfg, "_raw_yaml", {}) or {}) + raw_data = dict(getattr(cfg, "_raw_data", {}) or {}) + data = getattr(cfg, "data", None) + + data_section = dict(raw_yaml.get("data", {}) or {}) + data_section.update(raw_data) + + if data is not None: + for key in ( + "resolution", + "train_test_split", + "priority_order", + "camera_name_mapping", + ): + value = getattr(data, key, None) + if value is not None: + data_section.setdefault(key, value) + + raw_yaml["data"] = data_section + raw_yaml.setdefault("model_type", getattr(cfg, "model_type", "qwen2_5")) + return load_trainer_data_config_from_yaml_dict(raw_yaml) + + +def load_trainer_data_config_from_yaml_dict(yaml_dict: Dict[str, Any]) -> LerobotConfig: + """Build the LeRobot runtime config from a raw training YAML dict.""" + return LerobotConfig.from_yaml_dict(yaml_dict) + + +def _build_flat_config(cfg: Any) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """Map typed TrainConfig → (flat_config, lerobot_config) for legacy entry. + + ``load_lerobot_data`` expects a 2509-style flat dict plus a separate + ``lerobot_config`` carrying ``repo_id`` / ``root``. This function is + the one place that translation lives; keep it surgical so future + field additions on ``LeRobotDataConfig`` do not require touching the + legacy loader. + """ + model = cfg.model + data = cfg.data + hp = cfg.hyperparams + raw = getattr(cfg, "_raw_yaml", {}) or {} + raw_data = dict(getattr(cfg, "_raw_data", {}) or {}) + + lerobot_cfg = dict(data.lerobot_config or {}) + if "repo_id" not in lerobot_cfg: + raise ValueError( + "lerobot requires data.lerobot_config.repo_id to be set " + "(HuggingFace LeRobot dataset id)." + ) + + data_section: Dict[str, Any] = { + "key_mappings": data.key_mappings, + "action_horizon": cfg.task.action_horizon, + "train_test_split": data.train_test_split, + "seed": hp.seed, + "resolution": data.resolution, + } + if raw_data.get("max_length") is not None: + data_section["max_length"] = raw_data["max_length"] + if data.priority_order is not None: + data_section["priority_order"] = data.priority_order + if data.camera_name_mapping is not None: + data_section["camera_name_mapping"] = data.camera_name_mapping + data_section.setdefault( + "use_state_string_representation", + cfg.task.use_state_string_representation, + ) + data_section.setdefault( + "state_bins", + raw_data.get("state_bins", raw.get("state_bins", 256)), + ) + + # Dof/agent_pos totals for the collator's zero-pad step. When resuming + # from a checkpoint trained on a larger action space (e.g. bus2602 26-dim + # dof while libero data is 7-dim), task.dof_config should include an + # ``action_padding`` key that absorbs the diff; the collator right-pads + # action/agent_pos tensors to these totals with dof_mask/agent_pos_mask + # zeroed on padded dims so loss doesn't flow through them. + dof_total = int(sum((cfg.task.dof_config or {}).values())) + agent_pos_total = int(sum((cfg.task.agent_pos_config or {}).values())) + + flat: Dict[str, Any] = { + "model_type": cfg.model_type, + "processor_path": getattr(model, "processor_path", "") or "", + "norm_stats_path": data.norm_stats_path or raw.get("norm_stats_path"), + "batch_size_per_gpu": hp.batch_size_per_gpu, + "eval_batch_size_per_gpu": raw.get( + "eval_batch_size_per_gpu", hp.batch_size_per_gpu + ), + "num_workers": data.num_workers, + "padding_side": data.padding_side, + "use_fast_tokenizer": data.use_fast_tokenizer, + "action_tokenizer_path": data.action_tokenizer_path, + "noise_scheduler": data.noise_scheduler or {}, + "dof_total_dim": dof_total, + "agent_pos_total_dim": agent_pos_total, + "dof_config": dict(cfg.task.dof_config or {}), + "agent_pos_config": dict(cfg.task.agent_pos_config or {}), + "use_state_string_representation": cfg.task.use_state_string_representation, + "state_bins": int( + raw_data.get("state_bins") + or data_section.get("state_bins") + or raw.get("state_bins") + or 256 + ), + "data": data_section, + } + return flat, lerobot_cfg + + +def load_lerobot_v2( + cfg: Any, +) -> Tuple[_LerobotDatasetWrapper, torch.utils.data.DataLoader, int]: + """Build lerobot (wrapper, dataloader, train_num) from TrainConfig. + + The third return value ``train_num`` is a snapshot of + ``len(train_dataloader)`` at construction time. It matches + ``wrapper.global_train_iters.value`` initially but does not track + subsequent rebuilds inside ``set_epoch`` — callers doing dynamic + resampling should read from the mp.Value, not from this snapshot. + """ + from wall_x.data.backends.lerobot.loader import load_lerobot_data + + flat_cfg, lerobot_cfg = _build_flat_config(cfg) + + if dist.is_initialized(): + rank = dist.get_rank() + world_size = dist.get_world_size() + else: + rank = 0 + world_size = 1 + + seed = cfg.hyperparams.seed + inner, _ = load_lerobot_data( + flat_cfg, + lerobot_cfg, + rank=rank, + world_size=world_size, + seed=seed, + ) + + # PreprocessedDataset.get_*_dataloader returns (dataloader, sampler). + # Build val first, train second, so the inner ``_dataset`` pointer is + # left at the train split when we finish — workers fork from that + # state on first iteration. + val_dataloader, _ = inner.get_val_dataloader() + val_num = len(val_dataloader) if val_dataloader is not None else 0 + + train_dataloader, train_sampler = inner.get_train_dataloader() + train_num = len(train_dataloader) + + if rank == 0: + logger.info( + "\n%s\nLeRobot Data Loading Configuration:\n" + " RANK: %d\n WORLD SIZE: %d\n" + " BATCH SIZE PER DEVICE: %d\n GLOBAL BATCH SIZE: %d\n" + " TRAIN BATCHES: %d\n VAL BATCHES: %d\n" + " NUM WORKERS: %d\n REPO ID: %s\n%s", + "=" * 50, + rank, + world_size, + flat_cfg["batch_size_per_gpu"], + flat_cfg["batch_size_per_gpu"] * world_size, + train_num, + val_num, + flat_cfg["num_workers"], + lerobot_cfg.get("repo_id"), + "=" * 50, + ) + + wrapper = _LerobotDatasetWrapper( + inner, + train_dataloader, + train_sampler, + train_num, + val_dataloader=val_dataloader, + val_num=val_num, + ) + return wrapper, train_dataloader, train_num + + +def build(cfg, ctx): + """Backend Protocol entry — returns a ``DataBundle``. + + Wraps ``load_lerobot_v2`` (which returns the trainer-facing triple) + into the unified ``DataBundle`` shape every backend exposes. + """ + from wall_x.data._bundle import DataBundle + + wrapper, train_dataloader, train_num = load_lerobot_v2(cfg) + + # PreprocessedDataset shares one ``self._dataset`` pointer between + # train and val splits (flipped by ``_train()`` / ``_eval()``). + # ``wrapper.get_val_dataloader()`` flips the pointer to val. Flip back once + # here so the initial train loop starts from the right split even if callers + # inspect the raw ``train_dataloader`` before invoking ``set_epoch``. + val_loader = wrapper.get_val_dataloader() + inner = wrapper._inner + if hasattr(inner, "_train"): + inner._train() + + return DataBundle( + dataset=wrapper, + train_loader=train_dataloader, + val_loader=val_loader, + train_iters=train_num, + val_iters=wrapper.global_val_iters.value, + set_epoch=wrapper.set_epoch, + ) diff --git a/wall_x/data/backends/lerobot/config.py b/wall_x/data/backends/lerobot/config.py new file mode 100644 index 00000000..5a6c8075 --- /dev/null +++ b/wall_x/data/backends/lerobot/config.py @@ -0,0 +1,133 @@ +from typing import Any, Dict, Optional +from dataclasses import dataclass, field +from qwen_vl_utils.vision_process import MIN_PIXELS, MAX_PIXELS, IMAGE_FACTOR + + +@dataclass +class LerobotConfig: + """Configuration for the LeRobot preprocessing pipeline. + + Dataset-specific camera display names are optional config inputs. Other + dataset behavior is derived from the current LeRobot sample. + """ + + # Image resolution settings for different views + resolution: Dict[str, int] = field( + default_factory=lambda: { + "face_view": -1, + "left_wrist_view": 128, + "right_wrist_view": 128, + } + ) + + # Dataset splitting + train_test_split: float = 0.9 + seed: int = 42 + + # Instruction handling + priority_order: Optional[Dict[str, float]] = None + camera_name_mapping: Optional[Dict[str, str]] = None + + # Vision model parameters + model_type: str = "qwen2_5" + max_pixels: int = MAX_PIXELS + min_pixels: int = MIN_PIXELS + image_factor: int = IMAGE_FACTOR + + generate_subtask_ratio: float = 0.0 + + def __post_init__(self): + """Post-initialization validation and setup.""" + # Validate train/test split + if not 0 < self.train_test_split < 1: + raise ValueError( + f"train_test_split must be between 0 and 1, got {self.train_test_split}" + ) + + def as_dict(self) -> Dict: + """Convert configuration to dictionary format. + + Returns: + Dict: Configuration as dictionary + """ + return self.__dict__ + + def update(self, **kwargs) -> "LerobotConfig": + """Update configuration parameters. + + Args: + **kwargs: Key-value pairs to update + + Returns: + LerobotConfig: Updated configuration instance + """ + for key, value in kwargs.items(): + if hasattr(self, key): + setattr(self, key, value) + else: + raise ValueError(f"Unknown configuration parameter: {key}") + return self + + def __getitem__(self, key: str): + return getattr(self, key) + + @classmethod + def from_yaml_dict(cls, yaml_dict: Dict[str, Any]) -> "LerobotConfig": + """ + Build a LerobotConfig instance from a YAML dictionary. + + Supports two styles: + + 1) Top-level fields: + train_test_split: 0.8 + model_type: qwen2_5 + + 2) Nested under `data:` (higher priority): + data: + train_test_split: 0.8 + model_type: qwen2_5 + + Keys inside `data:` override top-level keys. + """ + + data_config = yaml_dict.get("data", {}) + + def get(key: str, default: Any = None): + """ + Helper function: + Read from `data` first, then fallback to the top-level YAML. + """ + return data_config.get(key, yaml_dict.get(key, default)) + + # Construct only fields that actually exist in LerobotConfig + params: Dict[str, Any] = { + # Action prediction settings + # Image resolution per camera view + "resolution": get( + "resolution", + { + "face_view": -1, + "left_wrist_view": 128, + "right_wrist_view": 128, + }, + ), + # Dataset train/test split configuration + "train_test_split": get("train_test_split", 0.9), + "seed": get("seed", 42), + # Instruction priority ordering (optional) + "priority_order": get("priority_order", None), + "camera_name_mapping": get("camera_name_mapping", None), + # Vision model parameters + "model_type": get("model_type", "qwen2_5"), + "max_pixels": get("max_pixels", MAX_PIXELS), + "min_pixels": get("min_pixels", MIN_PIXELS), + "image_factor": get("image_factor", IMAGE_FACTOR), + # Subtask generation ratio + "generate_subtask_ratio": get("generate_subtask_ratio", 0.0), + } + + # Keep only valid dataclass fields (ignore unknown YAML keys) + valid_fields = {f.name for f in cls.__dataclass_fields__.values()} + filtered_params = {k: v for k, v in params.items() if k in valid_fields} + + return cls(**filtered_params) diff --git a/wall_x/data/backends/lerobot/loader.py b/wall_x/data/backends/lerobot/loader.py new file mode 100644 index 00000000..d7698534 --- /dev/null +++ b/wall_x/data/backends/lerobot/loader.py @@ -0,0 +1,886 @@ +""" +LeRobot Dataset Loader - Distributed Version +""" + +import logging +import os +import numpy as np +import torch +from torch.utils.data import DistributedSampler, random_split +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from typing import Protocol, SupportsIndex, TypeVar +from qwen_vl_utils.vision_process import smart_resize +from wall_x.data.backends.lerobot.config import LerobotConfig +from wall_x.data.backends.lerobot.utils import ( + process_grounding_points, + get_wallx_normal_text, + replace_action_token, + preprocesser_call, + load_norm_stats, +) +from wall_x.data.backends.lerobot.rotation_layout import ( + LAYOUT_SKIP_KEYS, + euler_layout_dim as _euler_layout_dim, + layout_uses_6d_rotation as _layout_uses_6d_rotation, + maybe_convert_euler_to_6d, +) +from wall_x._vendor.x2robot_utils.geometry import ( + canonicalize_euler_zyx_batch_nb, + euler_to_matrix_zyx_batch_nb, + matrix_to_euler_zyx_batch_nb, + so3_to_matrix_batch_nb, +) + +from transformers import AutoProcessor + +T_co = TypeVar("T_co", covariant=True) +logger = logging.getLogger(__name__) + +RELATIVE_KEYWORD = "relative" +ROTATION_KEYWORD = "rotation" +RELATIVE_SKIP_KEYS = LAYOUT_SKIP_KEYS + + +def _compute_delta_from_state_and_abs_rot( + rotations: np.ndarray, state: np.ndarray +) -> np.ndarray: + """Relative rotation: R_rel = R_abs @ R_state^T.""" + if rotations.shape[-1] == 3: + rotations_matrix = euler_to_matrix_zyx_batch_nb(rotations) + out_is_euler = True + elif rotations.shape[-1] == 6: + rotations_matrix = so3_to_matrix_batch_nb(rotations) + out_is_euler = False + else: + raise ValueError( + f"Only 3D euler or 6D rotation supported, got {rotations.shape[-1]}D" + ) + + if state.shape[-1] == 3: + state_matrix = euler_to_matrix_zyx_batch_nb(state[np.newaxis, :])[0] + elif state.shape[-1] == 6: + state_matrix = so3_to_matrix_batch_nb(state[np.newaxis, :])[0] + else: + raise ValueError( + f"Only 3D euler or 6D rotation supported, got {state.shape[-1]}D" + ) + + r_rel = np.matmul(rotations_matrix, state_matrix.T) + if out_is_euler: + d_euler = matrix_to_euler_zyx_batch_nb(r_rel) + return canonicalize_euler_zyx_batch_nb(d_euler) + return r_rel[:, :2, :].reshape(r_rel.shape[0], 6) + + +# Abstract class for dataset +class Dataset(Protocol[T_co]): + """Interface for a dataset with random access.""" + + def __getitem__(self, index: SupportsIndex) -> T_co: + raise NotImplementedError("Subclasses of Dataset should implement __getitem__.") + + def __len__(self) -> int: + raise NotImplementedError("Subclasses of Dataset should implement __len__.") + + +class PreprocessedDataset(Dataset[T_co]): + def __init__( + self, + dataset, + config, + norm_stats, + dataload_config, + lerobot_config, + seed=42, + rank=0, + world_size=1, + test_only=False, + ): + self.hf_dataset = dataset + + if test_only: + self._dataset = dataset + else: + self._dataset = None + self.train_dataset, self.val_dataset = random_split( + dataset, + [0.95, 0.05], + torch.Generator().manual_seed(seed) if seed is not None else None, + ) + self._train() + + self.seed = seed + self.rank = rank + self.world_size = world_size + + # init configs + self.config = config + self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) + self.dataload_config = dataload_config + self.norm_stats = norm_stats + self.lerobot_config = lerobot_config + + self.data_config = LerobotConfig().update( + train_test_split=self.dataload_config["train_test_split"], + seed=self.dataload_config["seed"], + resolution=self.dataload_config.get("resolution", None), + priority_order=self.dataload_config.get("priority_order", None), + camera_name_mapping=self.dataload_config.get("camera_name_mapping", None), + ) + + self.key_mappings = self.dataload_config["key_mappings"] + + self._cam_key_mapping = self.key_mappings["camera"] + self._state_key_mapping = self.key_mappings + self._action_key_mapping = self.key_mappings + + task_cfg = self.config.get("task") or {} + self._dof_config = self.config.get("dof_config") or task_cfg.get( + "dof_config", {} + ) + self._agent_pos_config = self.config.get("agent_pos_config") or task_cfg.get( + "agent_pos_config", {} + ) + self._use_relative_action = any( + RELATIVE_KEYWORD in key for key in self._dof_config + ) + self._convert_action_euler_to_6d = _layout_uses_6d_rotation(self._dof_config) + self._convert_state_euler_to_6d = _layout_uses_6d_rotation( + self._agent_pos_config + ) + if self._convert_action_euler_to_6d or self._convert_state_euler_to_6d: + logger.info( + "LeRobot loader: Euler->6D rotation enabled " + "(action=%s, state=%s; raw action dim=%s -> %s)", + self._convert_action_euler_to_6d, + self._convert_state_euler_to_6d, + _euler_layout_dim(self._dof_config) + if self._convert_action_euler_to_6d + else "-", + sum( + d + for k, d in self._dof_config.items() + if k not in RELATIVE_SKIP_KEYS + ) + if self._convert_action_euler_to_6d + else "-", + ) + + def _maybe_convert_euler_to_6d(self, vec, layout_config: dict, enabled: bool): + converted = maybe_convert_euler_to_6d(vec, layout_config, enabled) + if ( + enabled + and layout_config + and isinstance(vec, torch.Tensor) + and converted is not vec + ): + return torch.as_tensor(converted, dtype=vec.dtype, device=vec.device) + return converted + + def _to_relative_action(self, action, agent_pos): + """Convert absolute action horizon to deltas w.r.t. current agent_pos.""" + action = np.asarray(action, dtype=np.float64) + agent_pos = np.asarray(agent_pos, dtype=np.float64) + if action.ndim == 1: + action = action[np.newaxis, :] + if agent_pos.ndim > 1: + agent_pos = agent_pos.reshape(-1) + + parts = [] + cur = 0 + for key, dim in self._dof_config.items(): + if key in RELATIVE_SKIP_KEYS: + continue + action_clip = action[:, cur : cur + dim] + agent_pos_clip = agent_pos[cur : cur + dim] + if RELATIVE_KEYWORD not in key: + parts.append(action_clip) + elif ROTATION_KEYWORD in key: + parts.append( + _compute_delta_from_state_and_abs_rot( + action_clip.astype(np.float64), + agent_pos_clip.astype(np.float64), + ) + ) + else: + parts.append(action_clip - agent_pos_clip[np.newaxis, :]) + cur += dim + + if not parts: + return action + return np.concatenate(parts, axis=1).astype(np.float32) + + def _vision_preprocess(self, frames): + processed_frames = [] + for key in self.hf_dataset.meta.camera_keys: + from PIL import Image + + current_obs = frames[key].clone().permute(1, 2, 0) + + img_pil = Image.fromarray((current_obs * 255).to(torch.uint8).cpu().numpy()) + orig_width, orig_height = img_pil.size + # 2. Apply resolution constraints (if config is not -1) + target_size = self.data_config.resolution.get( + self._cam_key_mapping[key], -1 + ) + if target_size != -1: + # Maintain aspect ratio logic + if orig_width > orig_height: # Landscape image + new_width = target_size + new_height = int(target_size * orig_height / orig_width) + else: # Portrait image + new_height = target_size + new_width = int(target_size * orig_width / orig_height) + img_pil = img_pil.resize((new_width, new_height)) + + # 3. Apply smart scaling (qwen logic) + current_width, current_height = img_pil.size + resized_height, resized_width = smart_resize( + current_height, + current_width, + factor=self.data_config.image_factor, + min_pixels=self.data_config.min_pixels, + max_pixels=self.data_config.max_pixels, + ) + resized_img = img_pil.resize((resized_width, resized_height)) + processed_frames.append(resized_img) + + return processed_frames, orig_height, orig_width, resized_height, resized_width + + def __getitem__(self, index): + data = self._dataset[index] + image_inputs, h, w, resize_h, resize_w = self._vision_preprocess(data) + agent_pos = data[self._state_key_mapping["state"]] + action = data[self._action_key_mapping["action"]] + agent_pos = self._maybe_convert_euler_to_6d( + agent_pos, self._agent_pos_config, self._convert_state_euler_to_6d + ) + action = self._maybe_convert_euler_to_6d( + action, self._dof_config, self._convert_action_euler_to_6d + ) + if self._use_relative_action: + device = action.device if isinstance(action, torch.Tensor) else None + action = torch.as_tensor( + self._to_relative_action(action, agent_pos), + dtype=torch.float32, + device=device, + ) + frame_index = data["frame_index"] + instruction_info = {"instruction": data["task"]} + generate_subtask_ratio = self.data_config.generate_subtask_ratio + complete_text, generate_subtask = get_wallx_normal_text( + instruction_info, + self.dataload_config.get("action_horizon", 33) - 1, + frame_index, + self.data_config.priority_order, + self._cam_key_mapping, + generate_subtask_ratio=generate_subtask_ratio, + camera_name_mapping=self.data_config.camera_name_mapping, + ) + text = process_grounding_points( + complete_text, h, w, resize_h, resize_w, self.data_config.model_type + ) + result = { + "image_inputs": image_inputs, + "text": text, + "action": action, + "agent_pos": agent_pos, + "frame_index": frame_index, + } + + return result + + def __len__(self) -> int: + return len(self._dataset) + + def _eval(self): + self._dataset = self.val_dataset + + def _train(self): + self._dataset = self.train_dataset + + def get_train_dataloader(self): + """ + Get distributed training dataloader + + Args: + rank: Current process rank + world_size: Total number of processes + seed: Random seed for reproducibility + """ + self._train() + + batch_size = self.config.get("batch_size_per_gpu", 8) + num_workers = self.config.get("num_workers", 4) + + # Create distributed sampler + sampler = DistributedSampler( + self, + num_replicas=self.world_size, + rank=self.rank, + shuffle=True, + seed=self.seed, + drop_last=True, # Ensure all processes have same number of batches + ) + + dataloader = torch.utils.data.DataLoader( + self, + batch_size=batch_size, + sampler=sampler, # Use distributed sampler instead of shuffle=True + num_workers=num_workers, + collate_fn=DataCollator( + self.config, self.dataload_config, self.norm_stats, self.lerobot_config + ), + pin_memory=True, # Enable for GPU training + persistent_workers=num_workers > 0, # Only if num_workers > 0 + prefetch_factor=2, # Reduce memory usage + drop_last=True, # Avoid incomplete batches + ) + + return dataloader, sampler + + def get_val_dataloader(self): + """ + Get distributed evaluation dataloader (no shuffling for consistent evaluation) + """ + self._eval() + + batch_size = self.config.get( + "eval_batch_size_per_gpu", self.config.get("batch_size_per_gpu", 8) + ) + num_workers = self.config.get("num_workers", 4) + + # Create distributed sampler for evaluation (no shuffle) + sampler = DistributedSampler( + self, + num_replicas=self.world_size, + rank=self.rank, + shuffle=False, # No shuffling for evaluation + drop_last=False, # Keep all samples for evaluation + ) + + dataloader = torch.utils.data.DataLoader( + self, + batch_size=batch_size, + sampler=sampler, + num_workers=num_workers, + collate_fn=DataCollator( + self.config, self.dataload_config, self.norm_stats, self.lerobot_config + ), + pin_memory=True, + persistent_workers=num_workers > 0, + prefetch_factor=2, + drop_last=False, + ) + + return dataloader, sampler + + +class DataCollator: + # Class-level cache for processors to avoid reloading + _processor_cache = {} + _action_tokenizer_cache = {} + _norm_stat_alignment_warnings = set() + + def __init__(self, config, dataload_config, stats, lerobot_config): + self.config = config + self.dataload_config = dataload_config + self.stats = stats + self.action_min_stat = stats["action"].min + self.action_delta = stats["action"].delta + self.state_min_stat = stats["state"].min + self.state_delta = stats["state"].delta + self.lerobot_config = lerobot_config + self.np_rng = np.random.default_rng() + + noise_scheduler_config = config.get("noise_scheduler", {}) + self.beta_alpha = noise_scheduler_config.get( + "beta_alpha", 1.5 + ) # alpha parameter of the Beta distribution + self.beta_beta = noise_scheduler_config.get( + "beta_beta", 1.0 + ) # beta parameter of the Beta distribution + self.s = noise_scheduler_config.get("s", 0.999) # scaling factor + self.time_shift = noise_scheduler_config.get( + "time_shift", 1.0 + ) # time shift factor + + self.beta_alpha = float(self.beta_alpha) + self.beta_beta = float(self.beta_beta) + self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) + self.use_state_string_representation = bool( + self.config.get("use_state_string_representation", False) + ) + self.state_bins = int(self.config.get("state_bins", 256)) + self.load_processor() + + def load_processor(self): + processor_path = self.config["processor_path"] + action_tokenizer_path = self.config.get("action_tokenizer_path", None) + + if ( + self.use_fast_tokenizer + and action_tokenizer_path not in self._action_tokenizer_cache + ): + self._action_tokenizer_cache[action_tokenizer_path] = ( + AutoProcessor.from_pretrained( + action_tokenizer_path, trust_remote_code=True + ) + ) + + # Use cached processors if available + if processor_path not in self._processor_cache: + processor = AutoProcessor.from_pretrained(processor_path, use_fast=True) + if self.config.get("padding_side", "left") == "left": + processor.tokenizer.padding_side = "left" + + new_tokens = ["<|propri|>", "<|action|>"] + processor.tokenizer.add_tokens(new_tokens) + if self.use_fast_tokenizer and self.config.get("model_type") == "qwen2_5": + action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path] + new_tokens = [ + f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size) + ] + processor.tokenizer.add_tokens(new_tokens) + begin_idx_token = "<|action_token_0|>" + token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token) + processor.tokenizer.init_kwargs["action_token_start_index"] = token_id + processor.tokenizer.init_kwargs["action_token_vocab_size"] = ( + action_tokenizer.vocab_size + ) + + self._processor_cache[processor_path] = processor + + self.processor = self._processor_cache[processor_path] + + if not self.use_fast_tokenizer: + self.train_action_tokenizer = None + else: + self.train_action_tokenizer = self._action_tokenizer_cache[ + action_tokenizer_path + ] + + @classmethod + def _normalize(cls, action, min_stat, delta): + """ + Normalize action data using min-max normalization. + """ + delta = torch.where(delta == 0, torch.ones_like(delta), delta) + x = (action - min_stat) / delta + x = x * 2 - 1 + x = torch.clamp(x, -1, 1) + return x + + @staticmethod + def _align_norm_stat(stat, value, *, pad_value: float, name: str): + """Align a 1-D norm stat with the current LeRobot tensor width.""" + stat = stat.to(device=value.device, dtype=value.dtype) + target_dim = value.shape[-1] + stat_dim = stat.shape[-1] + if stat_dim == target_dim: + return stat + if stat_dim > target_dim: + warning_key = ("truncate", name, stat_dim, target_dim) + if warning_key not in DataCollator._norm_stat_alignment_warnings: + logger.warning( + "Truncating LeRobot %s norm stat from %s to %s dims", + name, + stat_dim, + target_dim, + ) + DataCollator._norm_stat_alignment_warnings.add(warning_key) + return stat[..., :target_dim] + pad_shape = (*stat.shape[:-1], target_dim - stat_dim) + pad = torch.full(pad_shape, pad_value, device=value.device, dtype=value.dtype) + warning_key = ("pad", name, stat_dim, target_dim) + if warning_key not in DataCollator._norm_stat_alignment_warnings: + logger.warning( + "Padding LeRobot %s norm stat from %s to %s dims", + name, + stat_dim, + target_dim, + ) + DataCollator._norm_stat_alignment_warnings.add(warning_key) + return torch.cat([stat, pad], dim=-1) + + def __call__(self, batch): + additional_inputs = {} + + # Tail-pad widths when dof_config / agent_pos_config (sum) is larger + # than the lerobot action/state — typical when resuming a ckpt that + # was pretrained on a bigger action space. Extra columns are filled + # with zeros and their mask set to 0 so loss is not propagated. + dof_total = int(self.config.get("dof_total_dim", 0) or 0) + agent_pos_total = int(self.config.get("agent_pos_total_dim", 0) or 0) + + # Explicit init so the ``if action is not None`` guard and later + # ``replace_action_token`` call stay well-defined even if a batch + # unexpectedly omits the action / agent_pos keys. Without this the + # loop-local variables would leak NameError on the first miss. + action = None + dof_mask = None + agent_pos = None + agent_pos_mask = None + + for key in batch[0].keys(): + if key == "agent_pos": + agent_pos = torch.stack([item["agent_pos"] for item in batch]) + if agent_pos.dim() == 2: + agent_pos = agent_pos.unsqueeze(1) + agent_pos_mask = (~torch.isnan(agent_pos)).float() + agent_pos.nan_to_num_(nan=0.0) + state_min_stat = self._align_norm_stat( + self.state_min_stat, + agent_pos, + pad_value=0.0, + name="state.min", + ) + state_delta = self._align_norm_stat( + self.state_delta, + agent_pos, + pad_value=1.0, + name="state.delta", + ) + agent_pos = self._normalize(agent_pos, state_min_stat, state_delta) + if agent_pos_total and agent_pos.shape[-1] < agent_pos_total: + pad_w = agent_pos_total - agent_pos.shape[-1] + agent_pos = torch.nn.functional.pad(agent_pos, (0, pad_w)) + agent_pos_mask = torch.nn.functional.pad(agent_pos_mask, (0, pad_w)) + additional_inputs["proprioception"] = agent_pos + additional_inputs["agent_pos_mask"] = agent_pos_mask + elif key == "action": + action = torch.stack([item["action"] for item in batch]) + if action.dim() == 2: + action = action.unsqueeze(1) + dof_mask = (~torch.isnan(action)).float() + action.nan_to_num_(nan=0.0) + action_min_stat = self._align_norm_stat( + self.action_min_stat, + action, + pad_value=0.0, + name="action.min", + ) + action_delta = self._align_norm_stat( + self.action_delta, + action, + pad_value=1.0, + name="action.delta", + ) + action = self._normalize(action, action_min_stat, action_delta) + if dof_total and action.shape[-1] < dof_total: + pad_w = dof_total - action.shape[-1] + action = torch.nn.functional.pad(action, (0, pad_w)) + dof_mask = torch.nn.functional.pad(dof_mask, (0, pad_w)) + additional_inputs["action_chunk"] = action + additional_inputs["dof_mask"] = dof_mask + elif key == "image_inputs": + additional_inputs["image_inputs"] = [ + item["image_inputs"] for item in batch + ] + elif key == "text": + additional_inputs["text"] = [item["text"] for item in batch] + elif key == "frame_index": + additional_inputs["frame_index"] = torch.stack( + [item["frame_index"] for item in batch] + ) + else: + raise NotImplementedError( + f"{key} input not implemented in preprocesser" + ) + + # sample noise time + if action is not None: + sample_time = self.sample_time( + action.shape[0], + device=action.device, + dtype=torch.float32, + ) + additional_inputs["sample_time"] = sample_time + + additional_inputs["text"] = replace_action_token( + additional_inputs["text"], + additional_inputs["action_chunk"], + self.train_action_tokenizer if self.use_fast_tokenizer else None, + additional_inputs["dof_mask"], + ) + + inputs = preprocesser_call( + processor=self.processor, + text=additional_inputs.pop("text"), + images=additional_inputs.pop("image_inputs"), + videos=None, + padding=True, + truncation=True, + return_tensors="pt", + max_length=self.dataload_config.get("max_length", 768), + norm_state=( + additional_inputs["proprioception"] + if self.use_state_string_representation + and "proprioception" in additional_inputs + else None + ), + agent_pos_mask=additional_inputs.get("agent_pos_mask"), + state_bins=self.state_bins, + ) + + action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") + + # Gating token types + additional_inputs["moe_token_types"] = inputs.input_ids == action_token_id + + inputs.update(additional_inputs) + + inputs["dataset_names"] = [self.lerobot_config["repo_id"]] * inputs[ + "action_chunk" + ].shape[0] + + return inputs + + def sample_time(self, batch_size, device, dtype): + """ + Sample timesteps + + Use a Beta distribution to sample values in [0, 1], then scale them. + + Args: + batch_size (int): batch size + device: Device type + dtype: dtype + + Returns: + torch.Tensor: sampled timesteps with shape [batch_size] + """ + + sample_np = self.np_rng.beta( + self.beta_alpha, self.beta_beta, size=(batch_size,) + ).astype(np.float32) + sample = torch.from_numpy(sample_np).to( + device=device, dtype=dtype, non_blocking=True + ) + + # sample = self.beta_dist.sample([batch_size]).to(dtype=dtype) + time = 1 - sample + + # Apply diffusion time shift + if self.time_shift != 1.0: + time = (self.time_shift * time) / (1 + (self.time_shift - 1) * time) + + time = time * self.s # noise should denoise from 0 to 1 here + return time + + +def load_lerobot_data( + config, + lerobot_config, + rank=0, + world_size=1, + seed=42, +): + """ + Load LeRobot dataset with distributed support + + Args: + config: Model configuration + rank: Current process rank (default: 0) + world_size: Total number of processes (default: 1) + seed: Random seed for reproducibility (default: 42) + + Returns: + dataset: Training dataset + train_num: Number of training samples per process + sampler: Distributed sampler (None if world_size=1) + """ + + # Set seed for reproducibility + torch.manual_seed(seed) + + dataload_config = get_data_configs(config["data"]) + key_mappings = dataload_config["key_mappings"] + + repo_id = lerobot_config.get("repo_id", None) + assert repo_id is not None, "repo id is required" + root = lerobot_config.get("root", None) + meta_info = LeRobotDatasetMetadata(repo_id, root=root) + dataset_fps = meta_info.fps + episodes_num = meta_info.total_episodes + + norm_stats_path = config.get("norm_stats_path", None) + assert ( + norm_stats_path is not None + ), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats" + task_cfg = config.get("task") or {} + dof_config = config.get("dof_config") or task_cfg.get("dof_config", {}) + agent_pos_config = config.get("agent_pos_config") or task_cfg.get( + "agent_pos_config", {} + ) + norm_stats = load_norm_stats( + norm_stats_path, + key_mappings, + dof_config=dof_config, + agent_pos_config=agent_pos_config, + ) + + delta_timestamps = { + # action chunk + key_mappings["action"]: [ + t / dataset_fps + for t in range(dataload_config.get("action_horizon", 33) - 1) + ], + } + batch_size = config.get("batch_size_per_gpu", 8) + + # Optional episode subset. YAML ``lerobot_config.episodes`` has always + # been present in examples but previously ignored; honour it so smoke + # tests / small-dataset runs don't pay the O(N) LeRobotDataset indexing + # cost on a multi-thousand-episode repo (~10s / episode on some formats). + episodes_override = lerobot_config.get("episodes") + if episodes_override is not None: + episodes = list(episodes_override) + episodes_num_effective = len(episodes) + else: + episodes = np.arange(episodes_num).tolist() + episodes_num_effective = episodes_num + + train_test_split = dataload_config.get("train_test_split", 0.95) + split_idx = int(episodes_num_effective * train_test_split) + # Guard: tiny episode subsets + high train_test_split can floor split_idx + # to 0 (e.g. 1 ep * 0.95 = 0), which would silently hand LeRobotDataset an + # empty episode list and end training after 0 iterations. Fail loud. + if split_idx < 1: + raise ValueError( + f"train_test_split={train_test_split} applied to " + f"{episodes_num_effective} episode(s) yields 0 train episodes. " + f"Use more episodes or a lower train_test_split." + ) + train_episodes = episodes[:split_idx] + test_episodes = episodes[split_idx:] + + global_rank = torch.distributed.get_rank() + local_rank = int(os.environ["LOCAL_RANK"]) + local_world_size = int(os.environ["LOCAL_WORLD_SIZE"]) + # TODO: Some LeRobot formats need to load all metadata before splitting + # episodes; loading from all ranks at once can exhaust memory. + train_dataset = None + + # Sequential loading inside each node + for r in range(local_world_size): + if local_rank == r: + logger.info( + "[Global rank %s] Loading dataset on local_rank=%s", + global_rank, + local_rank, + ) + + train_dataset = LeRobotDataset( + repo_id=repo_id, + root=root, + episodes=train_episodes, + delta_timestamps=delta_timestamps, + video_backend="pyav", + ) + + logger.info( + "[Global rank %s] Finished loading on local_rank=%s", + global_rank, + local_rank, + ) + + # Barrier only within the node + torch.distributed.barrier(device_ids=[local_rank]) + + if rank == 0: + logger.info("Selected train episodes: %s", train_dataset.episodes) + logger.info("Number of train episodes selected: %s", train_dataset.num_episodes) + logger.info("Number of train frames selected: %s", train_dataset.num_frames) + logger.info("Selected test episodes: %s", test_episodes) + + dataset = PreprocessedDataset( + train_dataset, + config, + norm_stats, + dataload_config, + lerobot_config, + seed=seed, + rank=rank, + world_size=world_size, + ) + + # Calculate samples per process + if world_size > 1: + # With DistributedSampler, each process gets approximately len(dataset) // world_size samples + samples_per_process = len(dataset) // world_size + train_num = samples_per_process // batch_size + else: + train_num = len(dataset) // batch_size + + if rank == 0: + lines = [ + "LeRobot Data Loading Configuration:", + f" rank: {rank}", + f" world_size: {world_size}", + f" batch_size_per_gpu: {batch_size}", + f" repo_id: {repo_id}", + f" total_dataset_size: {len(dataset)}", + ] + if world_size > 1: + lines.extend( + [ + f" samples_per_process: {samples_per_process}", + f" batches_per_process: {train_num}", + f" total_batches_all_processes: {train_num * world_size}", + ] + ) + else: + lines.append(f" total_batches: {train_num}") + lines.append(f" seed: {seed}") + logger.info("\n%s", "\n".join(lines)) + + return dataset, train_num + + +def get_distributed_dataloader( + dataset, config, rank=0, world_size=1, seed=42, is_train=True +): + """ + Helper function to get distributed dataloader + + Args: + dataset: PreprocessedDataset instance + config: Configuration dict + rank: Current process rank + world_size: Total number of processes + seed: Random seed + is_train: Whether this is for training (affects shuffling) + + Returns: + dataloader: Distributed DataLoader + sampler: DistributedSampler + """ + if is_train: + return dataset.get_train_dataloader(rank=rank, world_size=world_size, seed=seed) + else: + return dataset.get_val_dataloader(rank=rank, world_size=world_size) + + +def get_data_configs(config): + default_data_config = { + "train_test_split": 0.95, + "seed": 42, + "batch_size": 8, + "action_horizon": 21, + "action_history_length": 0, + "image_horizon": 1, + "image_history_length": 0, + "left_padding": False, + "right_padding": False, + "return_first_obs": False, + "return_last_obs": False, + "randomize_obs_after": None, + "datasets": [], + "labeled_pathes": [], + "camera_name_mapping": None, + } + data_config = default_data_config | config + data_config["action_horizon"] += 1 + + return data_config diff --git a/wall_x/data/backends/lerobot/rotation_layout.py b/wall_x/data/backends/lerobot/rotation_layout.py new file mode 100644 index 00000000..ba8a780f --- /dev/null +++ b/wall_x/data/backends/lerobot/rotation_layout.py @@ -0,0 +1,95 @@ +"""Layout helpers when config expects 6D rotation but LeRobot stores 3D Euler.""" + +from __future__ import annotations + +import numpy as np + +from wall_x._vendor.x2robot_utils.geometry import euler_to_matrix_zyx_6d_nb + +LAYOUT_SKIP_KEYS = frozenset( + {"velocity_decomposed", "height", "head_actions", "action_padding"} +) +ROTATION_KEYWORD = "rotation" +ROTATION_6D_KEYWORD = "6D" + + +def layout_uses_6d_rotation(layout_config: dict) -> bool: + for key, dim in layout_config.items(): + if key in LAYOUT_SKIP_KEYS: + continue + if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6: + return True + return False + + +def euler_layout_dim(layout_config: dict) -> int: + """Vector width in LeRobot when rotation slices are still 3D Euler.""" + total = 0 + for key, dim in layout_config.items(): + if key in LAYOUT_SKIP_KEYS: + continue + if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6: + total += 3 + else: + total += int(dim) + return total + + +def convert_euler_to_6d(vec: np.ndarray, layout_config: dict) -> np.ndarray: + """Rewrite [pos, euler(3), tail...] to [pos, rot6d(6), tail...] per layout.""" + vec = np.asarray(vec, dtype=np.float64) + single = vec.ndim == 1 + if single: + vec = vec[np.newaxis, :] + + out_rows = [] + for row in vec: + parts: list[np.ndarray] = [] + raw_cur = 0 + for key, dim in layout_config.items(): + if key in LAYOUT_SKIP_KEYS: + continue + dim = int(dim) + if ROTATION_KEYWORD in key and ROTATION_6D_KEYWORD in key and dim == 6: + euler = row[raw_cur : raw_cur + 3] + rot6d = euler_to_matrix_zyx_6d_nb(euler.reshape(1, 3)).reshape(6) + parts.append(rot6d) + raw_cur += 3 + else: + parts.append(row[raw_cur : raw_cur + dim]) + raw_cur += dim + out_rows.append(np.concatenate(parts, axis=0)) + + out = np.stack(out_rows, axis=0) + return out[0] if single else out + + +def maybe_convert_norm_stats_vector( + values, + layout_config: dict, + enabled: bool | None = None, +): + """Convert a 1D norm-stat vector (q01/q99/mean/std) from Euler layout to 6D.""" + if enabled is None: + enabled = layout_uses_6d_rotation(layout_config) + if not enabled or not layout_config: + return values + arr = np.asarray(values, dtype=np.float64) + if arr.ndim != 1: + return values + raw_dim = euler_layout_dim(layout_config) + if arr.shape[0] != raw_dim: + return values + return convert_euler_to_6d(arr, layout_config).astype(np.float32) + + +def maybe_convert_euler_to_6d( + vec: np.ndarray, layout_config: dict, enabled: bool +) -> np.ndarray: + if not enabled or not layout_config: + return vec + raw_dim = euler_layout_dim(layout_config) + arr = np.asarray(vec) + if arr.shape[-1] != raw_dim: + return vec + return convert_euler_to_6d(arr, layout_config).astype(np.float32) diff --git a/wall_x/data/utils.py b/wall_x/data/backends/lerobot/utils.py similarity index 72% rename from wall_x/data/utils.py rename to wall_x/data/backends/lerobot/utils.py index 9cb0efcc..ce45a6cf 100644 --- a/wall_x/data/utils.py +++ b/wall_x/data/backends/lerobot/utils.py @@ -1,369 +1,90 @@ -""" -Data processing utilities for Wall-X multimodal robotic learning. - -This module provides utilities for preprocessing text, images, and actions -for multimodal transformer models in robotic learning tasks. -""" - import json +import logging import random import re from collections import OrderedDict from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple, Union -import numpy as np +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", - }, - "x2_normal": { - "camera": { - "observation.images.faceImg": "face_view", - "observation.images.leftImg": "left_wrist_view", - "observation.images.rightImg": "right_wrist_view", - }, - "state": "observation.state", - "action": "action", - }, - "libero": { - "camera": { - "observation.images.faceImg": "face_view", - "observation.images.rightImg": "right_wrist_view", - }, - "state": "observation.state", - "action": "action", - }, - "robochallenge_aloha": { - "camera": { - "observation.images.cam_high_rgb": "face_view", - "observation.images.cam_wrist_left_rgb": "left_wrist_view", - "observation.images.cam_wrist_right_rgb": "right_wrist_view", - }, - "state": "observation.state", - "action": "action", - }, -} - -CAMERA_NAME_MAPPING = { - "face_view": "front view", - "left_wrist_view": "left wrist view", - "right_wrist_view": "right wrist view", - "move1_view": "move view", - "move2_view": "move view", - "wall_view": "wall view", - "top_view": "top view", -} - - -MULTIMODAL_DATASET_NAMES = [ - "x2_multimodal_from_action", - "x2_multimodal", - "x2_subtask_generation", - "multimodal_CapsFusion", - "multimodal_Robo2VLM", - "multimodal_RoboPoint", - "multimodal_EQA", - "multimodal_Cambrian", - "multimodal_pixmo", - "multimodal_VQAv2", - "multimodal_COCO", -] - -FREQUENCY_MAPPING = { - "x2_normal": 32, - "fractal": 5, - "bridge_data_v2": 5, - "droid": 15, - "agibotworld_alpha": 32, - "DobbE": 30, - "RH20T": 10, - "UMI-biarm": 10, - "austin_buds": 20, - "austin_sailor": 20, - "austin_sirius": 20, - "bc_z": 10, - "berkeley_autolab_ur5": 5, - "berkeley_cable_routing": 10, - "berkeley_fanuc_manipulation": 10, - "dlr_edan_shared_control": 5, - "fmb": 10, - "furniture_bench": 10, - "jaco_play": 10, - "nyu_rot": 10, - "stanford_hydra": 10, - "stanford_kuka_multimodal": 20, - "taco_play": 30, - "utaustin_mutex": 20, - "viola": 20, -} +logger = logging.getLogger(__name__) -def preprocesser_call( - processor, - images: Optional[Union[List, Any]] = None, - text: Optional[Union[str, List[str]]] = None, - videos: Optional[Union[List, Any]] = None, - padding: Union[bool, str] = False, - truncation: Optional[bool] = None, - max_length: Optional[int] = None, - return_tensors: str = "pt", -) -> BatchFeature: - """Unified preprocessing function for Wall-X model handling text, image and video inputs. +@dataclass +class NormStats: + min: torch.Tensor + max: torch.Tensor + delta: torch.Tensor - Processes inputs into format suitable for multimodal transformer models, including: - - Text tokenization and special token handling - - Image/video processing through image processor - - Attention mask and label generation - - Padding and truncation handling - Args: - processor: Multimodal processor containing tokenizer and image processor - images: Input images (PIL, numpy arrays, or torch tensors) - text: Text or list of texts to tokenize - videos: Input videos (numpy arrays or torch tensors) - padding: Whether to pad sequences to same length - truncation: Whether to truncate sequences longer than max_length - max_length: Maximum length for truncation/padding - return_tensors: Format for returned tensors ('pt', 'np', etc.) +def load_norm_stats( + norm_stats_path, + key_mappings, + dof_config: dict | None = None, + agent_pos_config: dict | None = None, +): + from wall_x.data.backends.lerobot.rotation_layout import ( + layout_uses_6d_rotation, + maybe_convert_norm_stats_vector, + ) - Returns: - BatchFeature containing processed inputs with keys: - - input_ids: Tokenized text - - attention_mask: Attention mask for text - - pixel_values: Processed image pixels - - pixel_values_videos: Processed video frames - - image_grid_thw: Image grid dimensions for LLM - - video_grid_thw: Video grid dimensions for LLM - - labels: Training labels with masking - """ - # Process image inputs - if images is not None and len(images) > 0: - image_inputs = processor.image_processor( - images=images, videos=None, return_tensors=return_tensors - ) - image_grid_thw = image_inputs["image_grid_thw"] - else: - image_inputs = {} - image_grid_thw = None + with open(norm_stats_path, "r") as f: + norm_stats = json.load(f) + action_key = key_mappings["action"] + state_key = key_mappings["state"] - # Process video inputs - if videos is not None: - videos_inputs = processor.image_processor( - images=None, videos=videos, return_tensors=return_tensors + action_q01 = maybe_convert_norm_stats_vector( + norm_stats["norm_stats"][action_key]["q01"], dof_config or {} + ) + action_q99 = maybe_convert_norm_stats_vector( + norm_stats["norm_stats"][action_key]["q99"], dof_config or {} + ) + if layout_uses_6d_rotation(dof_config or {}) and len(action_q01) != len( + norm_stats["norm_stats"][action_key]["q01"] + ): + logger.info( + "Converted action norm stats Euler->6D (%d -> %d dims)", + len(norm_stats["norm_stats"][action_key]["q01"]), + len(action_q01), ) - video_grid_thw = videos_inputs["video_grid_thw"] - else: - videos_inputs = {} - video_grid_thw = None - - # Ensure text input is in list format - if not isinstance(text, list): - text = [text] - - # Process image placeholder tokens in text - if image_grid_thw is not None: - merge_length = processor.image_processor.merge_size**2 - index = 0 - for i in range(len(text)): - while "<|image_pad|>" in text[i]: - # Add bounds checking to avoid index overflow - if index >= len(image_grid_thw): - print( - f"Warning: Number of image placeholders ({index + 1}) " - f"exceeds actual images ({len(image_grid_thw)}), " - f"skipping remaining placeholder processing" - ) - break - # Replace image placeholder with actual token count - token_count = image_grid_thw[index].prod() // merge_length - text[i] = text[i].replace( - "<|image_pad|>", "<|placeholder|>" * token_count, 1 - ) - index += 1 - text[i] = text[i].replace("<|placeholder|>", "<|image_pad|>") - - # Process video placeholder tokens in text - if video_grid_thw is not None: - merge_length = processor.image_processor.merge_size**2 - index = 0 - for i in range(len(text)): - while "<|video_pad|>" in text[i]: - # Replace video placeholder with actual token count - token_count = video_grid_thw[index].prod() // merge_length - text[i] = text[i].replace( - "<|video_pad|>", "<|placeholder|>" * token_count, 1 - ) - index += 1 - text[i] = text[i].replace("<|placeholder|>", "<|video_pad|>") - # Tokenize complete input text - text_inputs = processor.tokenizer( - text, - return_tensors=return_tensors, - padding=padding, - truncation=truncation, - max_length=max_length, + q01 = torch.tensor(action_q01) + q99 = torch.tensor(action_q99) + delta = q99 - q01 + action_norm_stats = NormStats( + min=q01, + max=q99, + delta=delta, ) - # Get pad token ID for label generation - pad_token_id = processor.tokenizer.pad_token_id - if pad_token_id is None: - pad_token_id = processor.tokenizer.eos_token_id - - # Generate labels for multi-turn dialogue, keeping only assistant response loss - labels = torch.full_like(text_inputs.input_ids, -100) - assistant_marker = "<|im_start|>assistant\n" - im_end_token_id = processor.tokenizer.convert_tokens_to_ids("<|im_end|>") - assistant_tokens = processor.tokenizer( - "<|im_start|>assistant\n", add_special_tokens=False - ).input_ids - - for i in range(len(text)): - assistant_regions = [] - parts = text[i].split(assistant_marker) - - # Process each part to determine which tokens belong to assistant responses - # Count left padding tokens - num_left_pads = 0 - for token_id in text_inputs.input_ids[i]: - if token_id == pad_token_id: - num_left_pads += 1 - else: - break - current_pos = num_left_pads - - for j, part in enumerate(parts): - part_tokens = processor.tokenizer(part, add_special_tokens=False).input_ids - if j == 0: - # First part is system prompt or user question, all labels are -100 - current_pos += len(part_tokens) - continue - - # From second part onwards, each part starts with assistant response - for k in range(current_pos + 1, len(text_inputs.input_ids[i])): - if text_inputs.input_ids[i][k] == im_end_token_id: - assistant_regions.append( - (current_pos + len(assistant_tokens), k + 2) - ) - break - current_pos += len(part_tokens) + 3 - - # Set labels for assistant response regions - for start, end in assistant_regions: - labels[i][start:end] = text_inputs.input_ids[i][start:end] - - # Mask special action tokens in labels - action_token_id = processor.tokenizer.encode("<|action|>")[0] - propri_token_id = processor.tokenizer.encode("<|propri|>")[0] - labels[labels == action_token_id] = -100 - labels[labels == propri_token_id] = -100 - labels[labels == processor.tokenizer.pad_token_id] = -100 - - # Set labels to None if all are invalid to skip cross entropy loss - if (labels != -100).any().item(): - text_inputs["labels"] = labels - else: - text_inputs["labels"] = None - - return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}) - - -def process_grounding_points( - text: str, - orig_height: int, - orig_width: int, - resized_height: int, - resized_width: int, - model_type: str, -) -> str: - """Process grounding point coordinates in text based on image resizing. - - Adjusts coordinate values in tags to match resized image dimensions - for different model types (qwen2, qwen2_5). - - Args: - text: Input text containing tags with coordinates - orig_height: Original image height - orig_width: Original image width - resized_height: Resized image height - resized_width: Resized image width - model_type: Model type for coordinate processing ('qwen2' or 'qwen2_5') - - Returns: - Text with adjusted coordinate values - """ - # Regex pattern to match tags and their contents - point_pattern = re.compile(r"(.*?)") - - def process_match(match): - """Process a single point match and adjust coordinates.""" - coords_str = match.group(1) - try: - # Extract coordinates from string - coords = list(map(int, re.findall(r"\d+", coords_str))) - - # Calculate resize scale factors - scale_w = resized_width / orig_width - scale_h = resized_height / orig_height - - if len(coords) == 2: - x, y = coords - if model_type == "qwen2_5": - # Qwen2.5 uses pixel coordinates - new_x = max(0, min(round(x * scale_w), resized_width - 1)) - new_y = max(0, min(round(y * scale_h), resized_height - 1)) - elif model_type == "qwen2": - # Qwen2 normalizes to [0, 1000) range - new_x = max(0, min(999.999, (x / orig_width) * 1000)) - new_y = max(0, min(999.999, (y / orig_height) * 1000)) - else: - raise ValueError(f"Unsupported model type: {model_type}") - coords = [new_x, new_y] - - elif len(coords) == 4: - x1, y1, x2, y2 = coords - if model_type == "qwen2_5": - new_x1 = max(0, min(round(x1 * scale_w), resized_width - 1)) - new_y1 = max(0, min(round(y1 * scale_h), resized_height - 1)) - new_x2 = max(0, min(round(x2 * scale_w), resized_width - 1)) - new_y2 = max(0, min(round(y2 * scale_h), resized_height - 1)) - elif model_type == "qwen2": - new_x1 = max(0, min(999.999, (x1 / orig_width) * 1000)) - new_y1 = max(0, min(999.999, (y1 / orig_height) * 1000)) - new_x2 = max(0, min(999.999, (x2 / orig_width) * 1000)) - new_y2 = max(0, min(999.999, (y2 / orig_height) * 1000)) - else: - raise ValueError(f"Unsupported model type: {model_type}") - coords = [new_x1, new_y1, new_x2, new_y2] - - # Return processed point tag - return f'[{", ".join(map(str, coords))}]' + state_q01 = maybe_convert_norm_stats_vector( + norm_stats["norm_stats"][state_key]["q01"], agent_pos_config or {} + ) + state_q99 = maybe_convert_norm_stats_vector( + norm_stats["norm_stats"][state_key]["q99"], agent_pos_config or {} + ) + if layout_uses_6d_rotation(agent_pos_config or {}) and len(state_q01) != len( + norm_stats["norm_stats"][state_key]["q01"] + ): + logger.info( + "Converted state norm stats Euler->6D (%d -> %d dims)", + len(norm_stats["norm_stats"][state_key]["q01"]), + len(state_q01), + ) - except (ValueError, TypeError): - # Return original content if processing fails - return match.group(0) + q01 = torch.tensor(state_q01) + q99 = torch.tensor(state_q99) + delta = q99 - q01 + state_norm_stats = NormStats( + min=q01, + max=q99, + delta=delta, + ) - # Replace all matching point tags - processed_text = point_pattern.sub(process_match, text) - return processed_text + return {"action": action_norm_stats, "state": state_norm_stats} def get_frame_instruction( @@ -412,6 +133,57 @@ def get_frame_instruction( return instruction_for_frame, split_end +def get_action_tokens( + normalized_actions: Union[torch.Tensor, List], action_tokenizer +) -> List[List[str]]: + """Convert normalized actions to action token strings. + + Args: + normalized_actions: Normalized action arrays/tensors + action_tokenizer: Tokenizer for converting actions to tokens + + Returns: + List of action token string lists for each sample + """ + if isinstance(normalized_actions, torch.Tensor): + normalized_actions = normalized_actions.cpu().numpy() + + all_action_tokens = [] + for i in range(len(normalized_actions)): + if isinstance(normalized_actions[i], torch.Tensor): + normalized_actions[i] = normalized_actions[i].cpu().numpy() + + token_id = action_tokenizer(normalized_actions[i]) + action_tokens = [f"<|action_token_{j}|>" for j in token_id[0]] + all_action_tokens.append(action_tokens) + + return all_action_tokens + + +def pad_action_token_strs( + actions_token_lists: List[List[str]], pad_token: str = "<|endoftext|>" +) -> List[str]: + """Pad action token lists to same length and join as strings. + + Args: + actions_token_lists: List of action token lists for each sample + pad_token: Token used for padding + + Returns: + List of padded action token strings + """ + max_len = max(len(tokens) for tokens in actions_token_lists) + padded_action_strs = [] + + for tokens in actions_token_lists: + padded_tokens = ( + tokens + ["<|im_end|>\n"] + [pad_token] * (max_len - len(tokens)) + ) + padded_action_strs.append("".join(padded_tokens)) + + return padded_action_strs + + def get_task_instruction( frame_instruction_info: Dict[str, Any], priority_order: Optional[OrderedDict] = None ) -> str: @@ -449,15 +221,95 @@ def get_task_instruction( if random.random() >= prob: continue - task_instruction += f"\n{frame_instruction_info[key]}" - got_instruction = True - break + task_instruction += f"\n{frame_instruction_info[key]}" + got_instruction = True + break + + # Fall back to base instruction if no priority components found + if not got_instruction: + task_instruction = frame_instruction_info.get("instruction", "") + + return task_instruction + + +def process_grounding_points( + text: str, + orig_height: int, + orig_width: int, + resized_height: int, + resized_width: int, + model_type: str, +) -> str: + """Process grounding point coordinates in text based on image resizing. + + Adjusts coordinate values in tags to match resized image dimensions + for different model types (qwen2, qwen2_5). + + Args: + text: Input text containing tags with coordinates + orig_height: Original image height + orig_width: Original image width + resized_height: Resized image height + resized_width: Resized image width + model_type: Model type for coordinate processing ('qwen2' or 'qwen2_5') + + Returns: + Text with adjusted coordinate values + """ + # Regex pattern to match tags and their contents + point_pattern = re.compile(r"(.*?)") + + def process_match(match): + """Process a single point match and adjust coordinates.""" + coords_str = match.group(1) + try: + # Extract coordinates from string + coords = list(map(int, re.findall(r"\d+", coords_str))) + + # Calculate resize scale factors + scale_w = resized_width / orig_width + scale_h = resized_height / orig_height + + if len(coords) == 2: + x, y = coords + if model_type == "qwen2_5": + # Qwen2.5 uses pixel coordinates + new_x = max(0, min(round(x * scale_w), resized_width - 1)) + new_y = max(0, min(round(y * scale_h), resized_height - 1)) + elif model_type == "qwen2": + # Qwen2 normalizes to [0, 1000) range + new_x = max(0, min(999.999, (x / orig_width) * 1000)) + new_y = max(0, min(999.999, (y / orig_height) * 1000)) + else: + raise ValueError(f"Unsupported model type: {model_type}") + coords = [new_x, new_y] + + elif len(coords) == 4: + x1, y1, x2, y2 = coords + if model_type == "qwen2_5": + new_x1 = max(0, min(round(x1 * scale_w), resized_width - 1)) + new_y1 = max(0, min(round(y1 * scale_h), resized_height - 1)) + new_x2 = max(0, min(round(x2 * scale_w), resized_width - 1)) + new_y2 = max(0, min(round(y2 * scale_h), resized_height - 1)) + elif model_type == "qwen2": + new_x1 = max(0, min(999.999, (x1 / orig_width) * 1000)) + new_y1 = max(0, min(999.999, (y1 / orig_height) * 1000)) + new_x2 = max(0, min(999.999, (x2 / orig_width) * 1000)) + new_y2 = max(0, min(999.999, (y2 / orig_height) * 1000)) + else: + raise ValueError(f"Unsupported model type: {model_type}") + coords = [new_x1, new_y1, new_x2, new_y2] + + # Return processed point tag + return f'[{", ".join(map(str, coords))}]' - # Fall back to base instruction if no priority components found - if not got_instruction: - task_instruction = frame_instruction_info.get("instruction", "") + except (ValueError, TypeError): + # Return original content if processing fails + return match.group(0) - return task_instruction + # Replace all matching point tags + processed_text = point_pattern.sub(process_match, text) + return processed_text def get_wallx_normal_text( @@ -467,6 +319,7 @@ def get_wallx_normal_text( priority_order: Optional[OrderedDict] = None, cam_mapping: Optional[Dict[str, str]] = None, generate_subtask_ratio: float = 0.0, + camera_name_mapping: Optional[Dict[str, str]] = None, ) -> Tuple[str, bool]: """Construct complete multimodal prompt text for Wall-X model. @@ -484,6 +337,7 @@ def get_wallx_normal_text( priority_order: Priority order for instruction sampling cam_mapping: Camera name mapping dictionary generate_subtask_ratio: Probability of generating subtask instead of actions + camera_name_mapping: Optional display-name mapping for prompt text Returns: Tuple of (formatted_prompt_text, is_subtask_generation) @@ -506,8 +360,9 @@ def get_wallx_normal_text( # User request with observation user_request = f"{role_start_symbol}user\nObservation:" if cam_mapping: + camera_name_mapping = camera_name_mapping or {} for _, cam_name in cam_mapping.items(): - view_name = CAMERA_NAME_MAPPING.get(cam_name, cam_name) + view_name = camera_name_mapping.get(cam_name, cam_name) user_request += f" {view_name}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}" user_request += "\nInstruction:" @@ -552,62 +407,10 @@ def get_wallx_normal_text( return complete_text, generate_subtask -def get_action_tokens( - normalized_actions: Union[torch.Tensor, List], action_tokenizer -) -> List[List[str]]: - """Convert normalized actions to action token strings. - - Args: - normalized_actions: Normalized action arrays/tensors - action_tokenizer: Tokenizer for converting actions to tokens - - Returns: - List of action token string lists for each sample - """ - if isinstance(normalized_actions, torch.Tensor): - normalized_actions = normalized_actions.cpu().numpy() - - all_action_tokens = [] - for i in range(len(normalized_actions)): - if isinstance(normalized_actions[i], torch.Tensor): - normalized_actions[i] = normalized_actions[i].cpu().numpy() - - token_id = action_tokenizer(normalized_actions[i]) - action_tokens = [f"<|action_token_{j}|>" for j in token_id[0]] - all_action_tokens.append(action_tokens) - - return all_action_tokens - - -def pad_action_token_strs( - actions_token_lists: List[List[str]], pad_token: str = "<|endoftext|>" -) -> List[str]: - """Pad action token lists to same length and join as strings. - - Args: - actions_token_lists: List of action token lists for each sample - pad_token: Token used for padding - - Returns: - List of padded action token strings - """ - max_len = max(len(tokens) for tokens in actions_token_lists) - padded_action_strs = [] - - for tokens in actions_token_lists: - padded_tokens = ( - tokens + ["<|im_end|>\n"] + [pad_token] * (max_len - len(tokens)) - ) - padded_action_strs.append("".join(padded_tokens)) - - return padded_action_strs - - def replace_action_token( text: List[str], norm_action: Optional[torch.Tensor], action_tokenizer, - dataset_names: List[str], dof_masks: Optional[torch.Tensor] = None, ) -> List[str]: """Replace action placeholders in text with actual action tokens. @@ -616,26 +419,17 @@ def replace_action_token( text: List of text strings with action placeholders norm_action: Normalized action tensors action_tokenizer: Tokenizer for converting actions to tokens - dataset_names: Names of datasets for each sample dof_masks: Masks for degrees of freedom Returns: List of text strings with action tokens replaced """ - # Filter out multimodal dataset names - dataset_names = [ - name for name in dataset_names if name not in MULTIMODAL_DATASET_NAMES - ] - - # Get required action chunk sizes - required_chunk_sizes = [FREQUENCY_MAPPING.get(name, 32) for name in dataset_names] - if action_tokenizer is not None and norm_action is not None: - # Extract actions based on chunk sizes and DOF masks - norm_action = [ - action[: required_chunk_sizes[i], dof_masks[i, 0].bool()] - for i, action in enumerate(norm_action) - ] + if dof_masks is not None: + norm_action = [ + action[:, dof_masks[i, 0].bool()] + for i, action in enumerate(norm_action) + ] # Convert to action tokens and pad actions_fast_tokens = get_action_tokens(norm_action, action_tokenizer) @@ -660,115 +454,220 @@ def replace_action_token( return text -@dataclass -class NormStats: - min: torch.Tensor - max: torch.Tensor - delta: torch.Tensor +def preprocesser_call( + processor, + images: Optional[Union[List, Any]] = None, + text: Optional[Union[str, List[str]]] = None, + videos: Optional[Union[List, Any]] = None, + padding: Union[bool, str] = False, + truncation: Optional[bool] = None, + max_length: Optional[int] = None, + return_tensors: str = "pt", + norm_state=None, + agent_pos_mask=None, + state_bins: int = 256, + state_drop_prob: float = 0.0, +) -> BatchFeature: + """Unified preprocessing function for Wall-X model handling text, image and video inputs. + Processes inputs into format suitable for multimodal transformer models, including: + - Text tokenization and special token handling + - Image/video processing through image processor + - Attention mask and label generation + - Padding and truncation handling -def load_norm_stats(norm_stats_path, dataset_name): - with open(norm_stats_path, "r") as f: - norm_stats = json.load(f) - action_key = KEY_MAPPINGS[dataset_name]["action"] - state_key = KEY_MAPPINGS[dataset_name]["state"] - q01 = torch.tensor(norm_stats["norm_stats"][action_key]["q01"]) - q99 = torch.tensor(norm_stats["norm_stats"][action_key]["q99"]) - delta = q99 - q01 - action_norm_stats = NormStats( - min=q01, - max=q99, - delta=delta, - ) - q01 = torch.tensor(norm_stats["norm_stats"][state_key]["q01"]) - q99 = torch.tensor(norm_stats["norm_stats"][state_key]["q99"]) - delta = q99 - q01 - state_norm_stats = NormStats( - min=q01, - max=q99, - delta=delta, - ) + Args: + processor: Multimodal processor containing tokenizer and image processor + images: Input images (PIL, numpy arrays, or torch tensors) + text: Text or list of texts to tokenize + videos: Input videos (numpy arrays or torch tensors) + padding: Whether to pad sequences to same length + truncation: Whether to truncate sequences longer than max_length + max_length: Maximum length for truncation/padding + return_tensors: Format for returned tensors ('pt', 'np', etc.) - return {"action": action_norm_stats, "state": state_norm_stats} + Returns: + BatchFeature containing processed inputs with keys: + - input_ids: Tokenized text + - attention_mask: Attention mask for text + - pixel_values: Processed image pixels + - pixel_values_videos: Processed video frames + - image_grid_thw: Image grid dimensions for LLM + - video_grid_thw: Video grid dimensions for LLM + - labels: Training labels with masking + """ + # Process image inputs. transformers>=5.2 split image/video processing + # onto distinct callables and no longer accepts ``videos=`` on + # image_processor (or vice versa); older versions accepted ``videos=None``. + # Only pass the kwargs that correspond to present inputs. + if images is not None and len(images) > 0: + image_inputs = processor.image_processor( + images=images, return_tensors=return_tensors + ) + image_grid_thw = image_inputs["image_grid_thw"] + else: + image_inputs = {} + image_grid_thw = None + if videos is not None: + # transformers>=5.2 split video processing onto a dedicated + # ``video_processor``; older versions accepted ``videos=`` on the + # image_processor. Prefer the new API when present and fail loud + # otherwise, since passing ``videos=`` to an old-style + # image_processor silently works but feeds through a different + # preprocessing pipeline than the one the checkpoint was trained + # with. + if hasattr(processor, "video_processor"): + videos_inputs = processor.video_processor( + videos=videos, return_tensors=return_tensors + ) + else: + raise RuntimeError( + "processor has no video_processor attribute — this code path " + "requires transformers>=5.2. Upgrade transformers or pass " + "videos=None." + ) + video_grid_thw = videos_inputs["video_grid_thw"] + else: + videos_inputs = {} + video_grid_thw = None -def update_action_statistics( - action_statistic_dof: Dict[str, Any], - norm_stats_path: str, - repo_id: str, - dof_config: Dict[str, int] = None, - agent_pos_config: Dict[str, int] = None, - robot_name: str = None, - customized_dof_config: Dict[str, int] = None, - customized_agent_pos_config: Dict[str, int] = None, -) -> None: - """ - Update the action statistics dictionary with new robot configuration. + # Ensure text input is in list format + if not isinstance(text, list): + text = [text] - Args: - action_statistic_dof (Dict[str, Any]): The dictionary to be updated with statistics - norm_stats_path (str): Path to the normalization statistics file - repo_id (str): Repository ID for the LeRobot configuration - dof_config (Dict[str, int]): Configuration mapping DOF names to their dimensions - agent_pos_config (Dict[str, int]): Configuration mapping agent position names to their dimensions - robot_name (str, optional): Name of the robot. If None, uses repo_id as the key - customized_dof_config (Dict[str, int], optional): Customized DOF configuration for specific robot - customized_agent_pos_config (Dict[str, int], optional): Customized agent position configuration for specific robot - """ - # Load normalization statistics - norm_stats = load_norm_stats(norm_stats_path, repo_id) - - # Extract min and delta values for action and state - action_min = norm_stats["action"].min.numpy().tolist() - action_delta = norm_stats["action"].delta.numpy().tolist() - state_min = norm_stats["state"].min.numpy().tolist() - state_delta = norm_stats["state"].delta.numpy().tolist() - - # Use customized configurations if provided, otherwise use default ones - current_dof_config = ( - customized_dof_config if customized_dof_config is not None else dof_config - ) - current_agent_pos_config = ( - customized_agent_pos_config - if customized_agent_pos_config is not None - else agent_pos_config + # Discretize normalized proprioception into the <|propri|> prompt slot. + if norm_state is not None: + norm_state = ( + norm_state.cpu().numpy() + if isinstance(norm_state, torch.Tensor) + else norm_state + ) + discretized = ( + np.digitize(norm_state, bins=np.linspace(-1, 1, state_bins + 1)[:-1]) - 1 + ) + discretized = discretized[:, 0, :] + if agent_pos_mask is not None: + mask = ( + agent_pos_mask[:, 0, :].cpu().numpy().astype(bool) + if isinstance(agent_pos_mask, torch.Tensor) + else agent_pos_mask[:, 0, :].astype(bool) + ) + else: + mask = np.ones(discretized.shape, dtype=bool) + for i in range(len(text)): + if "<|propri|>" not in text[i]: + continue + if state_drop_prob > 0 and random.random() < state_drop_prob: + text[i] = text[i].replace("<|propri|>", "") + else: + state_str = " ".join(map(str, discretized[i, mask[i]])) + text[i] = text[i].replace("<|propri|>", state_str) + + # Process image placeholder tokens in text + if image_grid_thw is not None: + merge_length = processor.image_processor.merge_size**2 + index = 0 + for i in range(len(text)): + while "<|image_pad|>" in text[i]: + # Add bounds checking to avoid index overflow + if index >= len(image_grid_thw): + logger.warning( + "Number of image placeholders (%s) exceeds actual images " + "(%s); skipping remaining placeholder processing", + index + 1, + len(image_grid_thw), + ) + break + # Replace image placeholder with actual token count + token_count = image_grid_thw[index].prod() // merge_length + text[i] = text[i].replace( + "<|image_pad|>", "<|placeholder|>" * token_count, 1 + ) + index += 1 + text[i] = text[i].replace("<|placeholder|>", "<|image_pad|>") + + # Process video placeholder tokens in text + if video_grid_thw is not None: + merge_length = processor.image_processor.merge_size**2 + index = 0 + for i in range(len(text)): + while "<|video_pad|>" in text[i]: + # Replace video placeholder with actual token count + token_count = video_grid_thw[index].prod() // merge_length + text[i] = text[i].replace( + "<|video_pad|>", "<|placeholder|>" * token_count, 1 + ) + index += 1 + text[i] = text[i].replace("<|placeholder|>", "<|video_pad|>") + + # Tokenize complete input text + text_inputs = processor.tokenizer( + text, + return_tensors=return_tensors, + padding=padding, + truncation=truncation, + max_length=max_length, ) - # Prepare keys and values for DOF and agent position configurations - dof_key = [] - agent_pos_key = [] - dof_value = [] - agent_pos_value = [] - stats_dict = {} - - # Extract DOF configuration - for k, v in current_dof_config.items(): - dof_key.append(k) - dof_value.append(v) - - # Extract agent position configuration - for k, v in current_agent_pos_config.items(): - agent_pos_key.append(k) - agent_pos_value.append(v) - - # Calculate DOF indices and extract corresponding min/delta values - dof_idx = np.array([0] + dof_value).cumsum() - for i in range(len(dof_idx) - 1): - stats_dict[dof_key[i]] = { - "min": action_min[dof_idx[i] : dof_idx[i + 1]], - "delta": action_delta[dof_idx[i] : dof_idx[i + 1]], - } + # Get pad token ID for label generation + pad_token_id = processor.tokenizer.pad_token_id + if pad_token_id is None: + pad_token_id = processor.tokenizer.eos_token_id - # Calculate agent position indices and extract corresponding min/delta values - agent_pos_idx = np.array([0] + agent_pos_value).cumsum() - for i in range(len(agent_pos_idx) - 1): - stats_dict[agent_pos_key[i]] = { - "min": state_min[agent_pos_idx[i] : agent_pos_idx[i + 1]], - "delta": state_delta[agent_pos_idx[i] : agent_pos_idx[i + 1]], - } + # Generate labels for multi-turn dialogue, keeping only assistant response loss + labels = torch.full_like(text_inputs.input_ids, -100) + assistant_marker = "<|im_start|>assistant\n" + im_end_token_id = processor.tokenizer.convert_tokens_to_ids("<|im_end|>") + assistant_tokens = processor.tokenizer( + "<|im_start|>assistant\n", add_special_tokens=False + ).input_ids + + for i in range(len(text)): + assistant_regions = [] + parts = text[i].split(assistant_marker) + + # Process each part to determine which tokens belong to assistant responses + # Count left padding tokens + num_left_pads = 0 + for token_id in text_inputs.input_ids[i]: + if token_id == pad_token_id: + num_left_pads += 1 + else: + break + current_pos = num_left_pads + + for j, part in enumerate(parts): + part_tokens = processor.tokenizer(part, add_special_tokens=False).input_ids + if j == 0: + # First part is system prompt or user question, all labels are -100 + current_pos += len(part_tokens) + continue + + # From second part onwards, each part starts with assistant response + for k in range(current_pos + 1, len(text_inputs.input_ids[i])): + if text_inputs.input_ids[i][k] == im_end_token_id: + assistant_regions.append( + (current_pos + len(assistant_tokens), k + 2) + ) + break + current_pos += len(part_tokens) + 3 + + # Set labels for assistant response regions + for start, end in assistant_regions: + labels[i][start:end] = text_inputs.input_ids[i][start:end] + + # Mask special action tokens in labels + action_token_id = processor.tokenizer.encode("<|action|>")[0] + propri_token_id = processor.tokenizer.encode("<|propri|>")[0] + labels[labels == action_token_id] = -100 + labels[labels == propri_token_id] = -100 + labels[labels == processor.tokenizer.pad_token_id] = -100 - # Use provided robot name or repo_id as the key - robot_key = robot_name if robot_name is not None else repo_id + # Set labels to None if all are invalid to skip cross entropy loss + if (labels != -100).any().item(): + text_inputs["labels"] = labels + else: + text_inputs["labels"] = None - # Update the action_statistic_dof dictionary - action_statistic_dof.update({robot_key: stats_dict}) + return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}) diff --git a/wall_x/data/backends/mm b/wall_x/data/backends/mm new file mode 100644 index 00000000..4af8b90f --- /dev/null +++ b/wall_x/data/backends/mm @@ -0,0 +1,188 @@ +# WALL-OSS + +
+ +

+ +

+ +

+ +[![Paper](https://img.shields.io/badge/📄%20Paper-PDF-EA1B22?style=for-the-badge&logo=adobeacrobatreader&logoColor=fff)](https://x2robot.com/api/files/file/wall_oss_05.pdf) +   +[![Hugging Face](https://img.shields.io/badge/Hugging%20Face-x--square--robot-FFB000?style=for-the-badge&logo=huggingface&logoColor=000)](https://huggingface.co/x-square-robot) +   +[![GitHub](https://img.shields.io/badge/GitHub-181717?style=for-the-badge&logo=github&logoColor=fff)](https://github.com/X-Square-Robot/wall-x) +   +[![Project Page](https://img.shields.io/badge/Project-1E90FF?style=for-the-badge&logo=google-chrome&logoColor=fff)](https://x2robot.com/en/oss#resources) + +
+ +
+ +## Wall-OSS-0.5: A Deployment-Ready VLA with GradientBridged Pretraining + +We introduce Wall-OSS-0.5, an open-source 4B Vision-Language-Action (VLA) foundation model built upon a 3B VLM backbone augmented with dedicated action-generation components. While traditional VLAs are treated merely as optimization initializations, Wall-OSS-0.5 is designed so that pretrained robotic capability is directly executable and measurable on physical hardware without any downstream fine-tuning. + +The model is pretrained across more than 20 distinct robot embodiments, processing over one million trajectories per epoch alongside a grounded multimodal corpus. We adopt a novel gradient-bridged co-training recipe optimizing three complementary objectives: +- **Discrete Action Prediction**: Routes strong VLM-native gradients into the backbone. +- **Multimodal Prediction**: Preserves and strengthens grounded vision-language understanding. +- **Continuous Flow Matching**: Serves as the deployment-time continuous action interface. + +## 🌟 Key Highlights +1. **Zero-Shot Real-Robot Behavior**: Achieves non-trivial zero-shot completion on a 17-task suite (including held-out deformable manipulation tasks) directly from the pretrained checkpoint. +2. **Markedly Stronger Adaptation Prior**: After task-specific fine-tuning, Wall-OSS-0.5 reaches 60.5% average task progress on 15 real-robot tasks, outperforming $\pi_{0.5}$ by 17.5%. +3. **No Capabilities Erosion**: Multimodal evaluations confirm action-pretraining preserves broad vision-language competence while significantly sharpening embodied grounding. + + + +## 🚀 Quick Start + +### Installation + +```bash +# Create conda environment +conda create --name wallx python=3.10 +conda activate wallx + +# Install base requirements +pip install torch torchvision transformers +pip install huggingface_hub + +# Install Wall-X from GitHub +git clone https://github.com/X-Square-Robot/wall-x.git +cd wall-x +pip install -e . +``` + +### Basic Usage + +```python +import torch +from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction + +# Load the model +model_path = "X-Square-Robot/wall-oss-0.5" # or your local path +model = Qwen2_5_VLMoEForAction.from_pretrained(model_path) +model.eval() + +# Configuration +device = "cuda" if torch.cuda.is_available() else "cpu" +model = model.to(device).bfloat16() + +# Your inference code here... +``` + +## 🎯 Supervised Fine-Tuning (SFT) + +For training Wall-X on your robotics datasets, please refer to our comprehensive training guide: + +**📖 [Training Documentation](https://github.com/X-Square-Robot/wall-x/blob/main/workspace/README.md)** + +The training process includes: +- **Dataset Preparation**: How to prepare your robotics datasets in LeRobot format +- **Configuration Setup**: Detailed configuration for GPU setup, model paths, and robot DOF settings +- **Training Scripts**: Ready-to-use training scripts with proper hyperparameters + +### Quick Training Start + +```bash +# Run training (see workspace/README.md for detailed configuration) +bash ./workspace/lerobot_example/run.sh +``` + +## 🔮 Inference + +For detailed inference examples and model evaluation: + +**📖 [Inference Documentation](https://github.com/X-Square-Robot/wall-x/blob/main/scripts/)** + +### Basic Inference Example + +```python +import torch +from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction + +# Load model +model_path = "X-Square-Robot/wall-oss-0.5" +model = Qwen2_5_VLMoEForAction.from_pretrained(model_path) +model.eval() + +# Setup +batch_size = 1 +seq_length = 50 +device = "cuda" if torch.cuda.is_available() else "cpu" +model = model.to(device).bfloat16() + +# Prepare inputs (example with synthetic data) +torch.manual_seed(0) +input_ids = torch.randint(0, len(model.processor.tokenizer), (batch_size, seq_length), dtype=torch.long) +attention_mask = torch.ones((batch_size, seq_length), dtype=torch.long) +moe_token_types = torch.zeros((batch_size, seq_length), dtype=torch.long) +position_ids = torch.arange(seq_length, dtype=torch.long).unsqueeze(0).expand(batch_size, -1) + +# Robotics-specific inputs +proprioception = torch.randn((batch_size, 1, 20), dtype=torch.float32) # Joint states +agent_pos_mask = torch.ones((batch_size, 1, 20), dtype=torch.float32) +dof_mask = torch.ones((batch_size, 32, 20), dtype=torch.float32) # DOF mask +dataset_names = ["x2_normal"] + +# Move to device +inputs = { + "input_ids": input_ids.to(device), + "attention_mask": attention_mask.to(device), + "moe_token_types": moe_token_types.to(device), + "position_ids": position_ids.to(device), + "proprioception": proprioception.to(device).bfloat16(), + "agent_pos_mask": agent_pos_mask.to(device).bfloat16(), + "dof_mask": dof_mask.to(device).bfloat16(), + "dataset_names": dataset_names, + "mode": "validate" +} + +# Run inference +with torch.no_grad(): + outputs = model(**inputs) + print(f"Output logits shape: {outputs.logits.shape}") +``` + +### Advanced Inference Scripts + +For production-ready inference and evaluation scripts: + +```bash +# Basic inference test +python ./scripts/fake_inference.py + +# Generate open-loop comparison plots +python ./scripts/draw_openloop_plot.py +``` + +**📁 [View all inference scripts](https://github.com/X-Square-Robot/wall-x/tree/main/scripts)** + +## 📚 Complete Documentation + +For comprehensive setup, training, and inference instructions: + +### 🚀 **[Visit our GitHub Repository](https://github.com/X-Square-Robot/wall-x)** + +The repository contains: +- **Detailed Installation Guide**: Complete environment setup with all dependencies +- **Training Tutorials**: Step-by-step SFT process with LeRobot datasets +- **Inference Examples**: Multiple inference scripts and evaluation tools +- **Configuration Templates**: Ready-to-use configs for different robot setups +- **Troubleshooting Guide**: Common issues and solutions + +## 📄 Cite Us + +If you find WALL-OSS models useful, please cite: + +```bibtex +@misc{walloss_paper_2025, + title = {WALL-OSS: Igniting VLMs toward the Embodied Space}, + author = {X Square Robot}, + year = {2025}, + howpublished = {\url{https://x2robot.cn-wlcb.ufileos.com/wall_oss.pdf}}, + note = {White paper} +} +``` diff --git a/wall_x/data/config.py b/wall_x/data/config.py deleted file mode 100644 index 098e58e2..00000000 --- a/wall_x/data/config.py +++ /dev/null @@ -1,126 +0,0 @@ -from typing import List, Dict, Optional -from dataclasses import dataclass, field -from qwen_vl_utils.vision_process import MIN_PIXELS, MAX_PIXELS, IMAGE_FACTOR - - -# Tactile sensor file mapping for data processing -TACTILE_FILE_MAPPING = { - "tactile_data_left": "left_tactile", - "tactile_data_right": "right_tactile", -} - -# Supported action datasets -ACTION_DATASET_NAMES = [ - "x2_normal", - "agibotworld_alpha", - "droid", - "fractal", - "bridge_data_v2", - "DobbE", - "RH20T", - "UMI-biarm", - "austin_buds", - "austin_sailor", - "austin_sirius", - "bc_z", - "berkeley_autolab_ur5", - "berkeley_cable_routing", - "berkeley_fanuc_manipulation", - "dlr_edan_shared_control", - "fmb", - "furniture_bench", - "jaco_play", - "nyu_rot", - "stanford_hydra", - "stanford_kuka_multimodal", - "taco_play", - "utaustin_mutex", - "viola", - "physical-intelligence/libero", - "lerobot/aloha_mobile_cabinet", -] - -# Supported multimodal datasets -MULTIMODAL_DATASET_NAMES = [ - "x2_multimodal_from_action", - "x2_multimodal", - "x2_subtask_generation", - "multimodal_CapsFusion", - "multimodal_Robo2VLM", - "multimodal_RoboPoint", - "multimodal_EQA", - "multimodal_Cambrian", - "multimodal_pixmo", - "multimodal_VQAv2", - "multimodal_COCO", -] - - -@dataclass -class X2RDataProcessingConfig: - """Configuration class for X2R data processing pipeline. - - This class contains all the necessary parameters for processing robotic data - including camera mappings, tactile sensor configurations, action predictions, - and various processing options. - """ - - # Action prediction configuration - predict_action_keys: List[str] = field(default_factory=list) - obs_action_keys: List[str] = field(default_factory=list) - - # Image resolution settings for different views - resolution: Dict[str, int] = field( - default_factory=lambda: { - "face_view": -1, - "left_wrist_view": 128, - "right_wrist_view": 128, - } - ) - - # Dataset splitting - train_test_split: float = 0.9 - split_seed: int = 42 - - # Instruction handling - priority_order: Optional[Dict[str, float]] = None - - # Vision model parameters - model_type: str = "qwen2_5" - max_pixels: int = MAX_PIXELS - min_pixels: int = MIN_PIXELS - image_factor: int = IMAGE_FACTOR - - generate_subtask_ratio: float = 0.0 - - def __post_init__(self): - """Post-initialization validation and setup.""" - # Validate train/test split - if not 0 < self.train_test_split < 1: - raise ValueError( - f"train_test_split must be between 0 and 1, got {self.train_test_split}" - ) - - def as_dict(self) -> Dict: - """Convert configuration to dictionary format. - - Returns: - Dict: Configuration as dictionary - """ - return self.__dict__ - - def update(self, **kwargs) -> "X2RDataProcessingConfig": - """Update configuration parameters. - - Args: - **kwargs: Key-value pairs to update - - Returns: - X2RDataProcessingConfig: Updated configuration instance - """ - for key, value in kwargs.items(): - if hasattr(self, key): - setattr(self, key, value) - else: - raise ValueError(f"Unknown configuration parameter: {key}") - return self diff --git a/wall_x/data/load_lerobot_dataset.py b/wall_x/data/load_lerobot_dataset.py deleted file mode 100644 index 6bef6863..00000000 --- a/wall_x/data/load_lerobot_dataset.py +++ /dev/null @@ -1,715 +0,0 @@ -""" -LeRobot Dataset Loader - Distributed Version -""" - -import numpy as np -import torch -from torch.utils.data import DistributedSampler, random_split -from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata -from typing import Protocol, SupportsIndex, TypeVar -from qwen_vl_utils.vision_process import smart_resize -from wall_x.data.config import X2RDataProcessingConfig -from wall_x.data.utils import ( - process_grounding_points, - get_wallx_normal_text, - replace_action_token, - preprocesser_call, -) - -from transformers import AutoProcessor -from .utils import KEY_MAPPINGS - -T_co = TypeVar("T_co", covariant=True) - - -# Abstract class for dataset -class Dataset(Protocol[T_co]): - """Interface for a dataset with random access.""" - - def __getitem__(self, index: SupportsIndex) -> T_co: - raise NotImplementedError("Subclasses of Dataset should implement __getitem__.") - - def __len__(self) -> int: - raise NotImplementedError("Subclasses of Dataset should implement __len__.") - - -class PreprocessedDataset(Dataset[T_co]): - def __init__( - self, - dataset, - config, - dataload_config, - normalizer_action, - normalizer_propri, - lerobot_config, - seed=42, - rank=0, - world_size=1, - test_only=False, - ): - self.hf_dataset = dataset - - if test_only: - self._dataset = dataset - else: - self._dataset = None - self.train_dataset, self.val_dataset = random_split( - dataset, - [0.95, 0.05], - torch.Generator().manual_seed(seed) if seed is not None else None, - ) - self._train() - - self.seed = seed - self.rank = rank - self.world_size = world_size - - # init configs - self.config = config - self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) - self.dataload_config = dataload_config - self.normalizer_action = (normalizer_action,) - self.normalizer_propri = normalizer_propri - # self.norm_stats = norm_stats - self.lerobot_config = lerobot_config - - self.data_config = X2RDataProcessingConfig().update( - train_test_split=self.dataload_config["train_test_split"], - split_seed=self.dataload_config["split_seed"], - predict_action_keys=self.dataload_config["predict_action_keys"], - obs_action_keys=self.dataload_config["obs_action_keys"], - resolution=self.dataload_config.get("resolution", None), - priority_order=self.dataload_config.get("priority_order", None), - ) - - self._cam_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id]["camera"] - self._state_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id] - self._action_key_mapping = KEY_MAPPINGS[self.hf_dataset.meta.repo_id] - - def _vision_preprocess(self, frames): - processed_frames = [] - for key in self.hf_dataset.meta.camera_keys: - from PIL import Image - - current_obs = frames[key].clone().permute(1, 2, 0) - - img_pil = Image.fromarray((current_obs * 255).to(torch.uint8).cpu().numpy()) - orig_width, orig_height = img_pil.size - # 2. Apply resolution constraints (if config is not -1) - target_size = self.data_config.resolution.get( - self._cam_key_mapping[key], -1 - ) - if target_size != -1: - # Maintain aspect ratio logic - if orig_width > orig_height: # Landscape image - new_width = target_size - new_height = int(target_size * orig_height / orig_width) - else: # Portrait image - new_height = target_size - new_width = int(target_size * orig_width / orig_height) - img_pil = img_pil.resize((new_width, new_height)) - - # 3. Apply smart scaling (qwen logic) - current_width, current_height = img_pil.size - resized_height, resized_width = smart_resize( - current_height, - current_width, - factor=self.data_config.image_factor, - min_pixels=self.data_config.min_pixels, - max_pixels=self.data_config.max_pixels, - ) - resized_img = img_pil.resize((resized_width, resized_height)) - processed_frames.append(resized_img) - - return processed_frames, orig_height, orig_width, resized_height, resized_width - - def __getitem__(self, index): - data = self._dataset[index] - image_inputs, h, w, resize_h, resize_w = self._vision_preprocess(data) - agent_pos = data[self._state_key_mapping["state"]] - action = data[self._action_key_mapping["action"]] - frame_index = data["frame_index"] - instruction_info = {"instruction": data["task"]} - generate_subtask_ratio = self.data_config.generate_subtask_ratio - - complete_text, generate_subtask = get_wallx_normal_text( - instruction_info, - self.dataload_config.get("action_horizon", 33) - 1, - frame_index, - self.data_config.priority_order, - self._cam_key_mapping, - generate_subtask_ratio=generate_subtask_ratio, - ) - text = process_grounding_points( - complete_text, h, w, resize_h, resize_w, self.data_config.model_type - ) - result = { - "image_inputs": image_inputs, - "text": text, - "action": action, - "agent_pos": agent_pos, - "frame_index": frame_index, - } - - return result - - def __len__(self) -> int: - return len(self._dataset) - - def _eval(self): - self._dataset = self.val_dataset - - def _train(self): - self._dataset = self.train_dataset - - def get_train_dataloader(self): - """ - Get distributed training dataloader - - Args: - rank: Current process rank - world_size: Total number of processes - seed: Random seed for reproducibility - """ - self._train() - - batch_size = self.config.get("batch_size_per_gpu", 8) - num_workers = self.config.get("num_workers", 4) - - # Create distributed sampler - sampler = DistributedSampler( - self, - num_replicas=self.world_size, - rank=self.rank, - shuffle=True, - seed=self.seed, - drop_last=True, # Ensure all processes have same number of batches - ) - - dataloader = torch.utils.data.DataLoader( - self, - batch_size=batch_size, - sampler=sampler, # Use distributed sampler instead of shuffle=True - num_workers=num_workers, - collate_fn=DataCollator( - self.config, - self.dataload_config, - self.normalizer_action, - self.normalizer_propri, - self.lerobot_config, - ), - pin_memory=True, # Enable for GPU training - persistent_workers=num_workers > 0, # Only if num_workers > 0 - prefetch_factor=2, # Reduce memory usage - drop_last=True, # Avoid incomplete batches - ) - - return dataloader, sampler - - def get_val_dataloader(self): - """ - Get distributed evaluation dataloader (no shuffling for consistent evaluation) - """ - self._eval() - - batch_size = self.config.get( - "eval_batch_size_per_gpu", self.config.get("batch_size_per_gpu", 8) - ) - num_workers = self.config.get("num_workers", 4) - - # Create distributed sampler for evaluation (no shuffle) - sampler = DistributedSampler( - self, - num_replicas=self.world_size, - rank=self.rank, - shuffle=False, # No shuffling for evaluation - drop_last=False, # Keep all samples for evaluation - ) - - dataloader = torch.utils.data.DataLoader( - self, - batch_size=batch_size, - sampler=sampler, - num_workers=num_workers, - collate_fn=DataCollator( - self.config, self.dataload_config, self.norm_stats, self.lerobot_config - ), - pin_memory=True, - persistent_workers=num_workers > 0, - prefetch_factor=2, - drop_last=False, - ) - - return dataloader, sampler - - -class DataCollator: - # Class-level cache for processors to avoid reloading - _processor_cache = {} - _action_tokenizer_cache = {} - - def __init__( - self, - config, - dataload_config, - normalizer_action, - normalizer_propri, - lerobot_config, - ): - self.config = config - self.dataload_config = dataload_config - - self.normalizer_action = normalizer_action[0] - self.normalizer_propri = normalizer_propri - self.lerobot_config = lerobot_config - - self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) - self.dataset_name = self.config["data"]["lerobot_config"].get("repo_id", "") - self.dataset_name = [self.dataset_name] * self.config["batch_size_per_gpu"] - self.load_processor() - - def load_processor(self): - processor_path = self.config["pretrained_wallx_path"] - action_tokenizer_path = self.config.get("action_tokenizer_path", None) - - if ( - self.use_fast_tokenizer - and action_tokenizer_path not in self._action_tokenizer_cache - ): - self._action_tokenizer_cache[action_tokenizer_path] = ( - AutoProcessor.from_pretrained( - action_tokenizer_path, trust_remote_code=True - ) - ) - - # Use cached processors if available - if processor_path not in self._processor_cache: - processor = AutoProcessor.from_pretrained(processor_path, use_fast=True) - if self.config.get("padding_side", "left") == "left": - processor.tokenizer.padding_side = "left" - - new_tokens = ["<|propri|>", "<|action|>"] - processor.tokenizer.add_tokens(new_tokens) - if self.use_fast_tokenizer and self.config.get("model_type") == "qwen2_5": - action_tokenizer = self._action_tokenizer_cache[action_tokenizer_path] - new_tokens = [ - f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size) - ] - processor.tokenizer.add_tokens(new_tokens) - begin_idx_token = "<|action_token_0|>" - token_id = processor.tokenizer.convert_tokens_to_ids(begin_idx_token) - processor.tokenizer.init_kwargs["action_token_start_index"] = token_id - processor.tokenizer.init_kwargs["action_token_vocab_size"] = ( - action_tokenizer.vocab_size - ) - - self._processor_cache[processor_path] = processor - - self.processor = self._processor_cache[processor_path] - - if not self.use_fast_tokenizer: - self.train_action_tokenizer = None - else: - self.train_action_tokenizer = self._action_tokenizer_cache[ - action_tokenizer_path - ] - - @classmethod - def _normalize(cls, action, min_stat, delta): - """ - Normalize action data using min-max normalization. - """ - delta = torch.where(delta == 0, torch.ones_like(delta), delta) - x = (action - min_stat) / delta - x = x * 2 - 1 - x = torch.clamp(x, -1, 1) - return x - - def __call__(self, batch): - additional_inputs = {} - - for key in batch[0].keys(): - if key == "agent_pos": - agent_pos = torch.stack([item["agent_pos"] for item in batch]) - if agent_pos.dim() == 2: - agent_pos = agent_pos.unsqueeze(1) - agent_pos_mask = (~torch.isnan(agent_pos)).float() - # print("agent_pos_mask",agent_pos_mask.shape) - agent_pos.nan_to_num_(nan=0.0) - - # if agent_pos.shape[-1] != 20: - # agent_pos = torch.cat( - # [ - # agent_pos, - # torch.zeros( - # agent_pos.shape[0], - # agent_pos.shape[1], - # 20 - agent_pos.shape[-1], - # ), - # ], - # dim=-1, - # ) - # agent_pos_mask = torch.cat( - # [ - # agent_pos_mask, - # torch.zeros( - # agent_pos_mask.shape[0], - # agent_pos_mask.shape[1], - # 20 - agent_pos_mask.shape[-1], - # ), - # ], - # dim=-1, - # ) - agent_pos = self.normalizer_propri.normalize_data( - agent_pos, self.dataset_name - ) - additional_inputs["proprioception"] = agent_pos - additional_inputs["agent_pos_mask"] = agent_pos_mask - elif key == "action": - action = torch.stack([item["action"] for item in batch]) - if action.dim() == 2: - action = action.unsqueeze(1) - dof_mask = (~torch.isnan(action)).float() - action.nan_to_num_(nan=0.0) - - # if action.shape[-1] != 20: - # action = torch.cat( - # [ - # action, - # torch.zeros( - # action.shape[0], action.shape[1], 20 - action.shape[-1] - # ), - # ], - # dim=-1, - # ) - # dof_mask = torch.cat( - # [ - # dof_mask, - # torch.zeros( - # dof_mask.shape[0], - # dof_mask.shape[1], - # 20 - dof_mask.shape[-1], - # ), - # ], - # dim=-1, - # ) - action = self.normalizer_action.normalize_data( - action, self.dataset_name - ) - additional_inputs["action_chunk"] = action - additional_inputs["dof_mask"] = dof_mask - elif key == "image_inputs": - additional_inputs["image_inputs"] = [ - item["image_inputs"] for item in batch - ] - elif key == "text": - additional_inputs["text"] = [item["text"] for item in batch] - elif key == "frame_index": - additional_inputs["frame_index"] = torch.stack( - [item["frame_index"] for item in batch] - ) - else: - raise NotImplementedError( - f"{key} input not implemented in preprocesser" - ) - - additional_inputs["text"] = replace_action_token( - additional_inputs["text"], - additional_inputs["action_chunk"], - self.train_action_tokenizer if self.use_fast_tokenizer else None, - [self.lerobot_config["repo_id"]] * additional_inputs["text"].__len__(), - additional_inputs["dof_mask"], - ) - - inputs = preprocesser_call( - processor=self.processor, - text=additional_inputs.pop("text"), - images=additional_inputs.pop("image_inputs"), - videos=None, - padding=True, - truncation=True, - return_tensors="pt", - max_length=self.dataload_config.get("max_length", 768), - ) - - action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") - - # Gating token types - additional_inputs["moe_token_types"] = inputs.input_ids == action_token_id - - inputs.update(additional_inputs) - - inputs["dataset_names"] = [self.lerobot_config["repo_id"]] * inputs[ - "action_chunk" - ].shape[0] - - return inputs - - -def load_lerobot_data( - config, - lerobot_config, - normalizer_action, - normalizer_propri, - rank=0, - world_size=1, - seed=42, -): - """ - Load LeRobot dataset with distributed support - - Args: - config: Model configuration - rank: Current process rank (default: 0) - world_size: Total number of processes (default: 1) - seed: Random seed for reproducibility (default: 42) - - Returns: - dataset: Training dataset - train_num: Number of training samples per process - sampler: Distributed sampler (None if world_size=1) - """ - - # Set seed for reproducibility - torch.manual_seed(seed) - - dataload_config = get_data_configs(config["data"]) - - repo_id = lerobot_config.get("repo_id", None) - assert repo_id is not None, "repo id is required" - root = lerobot_config.get("root", None) - meta_info = LeRobotDatasetMetadata(repo_id, root=root) - dataset_fps = meta_info.fps - episodes_num = meta_info.total_episodes - - # norm_stats_path = config.get("norm_stats_path", None) - # assert ( - # norm_stats_path is not None - # ), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats" - # norm_stats = load_norm_stats(norm_stats_path, repo_id) - - delta_timestamps = { - # action chunk - KEY_MAPPINGS[repo_id]["action"]: [ - t / dataset_fps - for t in range(dataload_config.get("action_horizon", 33) - 1) - ], - } - batch_size = config.get("batch_size_per_gpu", 8) - episodes = np.arange(episodes_num).tolist() - - train_test_split = dataload_config.get("train_test_split", 0.95) - train_episodes = episodes[: int(episodes_num * train_test_split)] - test_episodes = episodes[int(episodes_num * train_test_split) :] - - train_dataset = LeRobotDataset( - repo_id, - root=root, - episodes=train_episodes, - delta_timestamps=delta_timestamps, - video_backend="pyav", - ) - - if rank == 0: - print(f"Selected train episodes: {train_dataset.episodes}") - print(f"Number of train episodes selected: {train_dataset.num_episodes}") - print(f"Number of train frames selected: {train_dataset.num_frames}") - print(f"Selected test episodes: {test_episodes}") - - dataset = PreprocessedDataset( - train_dataset, - config, - dataload_config, - normalizer_action, - normalizer_propri, - lerobot_config, - seed=seed, - rank=rank, - world_size=world_size, - ) - - # Calculate samples per process - if world_size > 1: - # With DistributedSampler, each process gets approximately len(dataset) // world_size samples - samples_per_process = len(dataset) // world_size - train_num = samples_per_process // batch_size - else: - train_num = len(dataset) // batch_size - - if rank == 0: - print("\n" + "=" * 50) - print("LeRobot Data Loading Configuration:") - print(f"✦ RANK: {rank}") - print(f"✦ WORLD SIZE: {world_size}") - print(f"✦ BATCH SIZE PER GPU: {batch_size}") - print(f"✦ REPO ID: {repo_id}") - print(f"✦ TOTAL DATASET SIZE: {len(dataset)}") - if world_size > 1: - print(f"✦ SAMPLES PER PROCESS: {samples_per_process}") - print(f"✦ BATCHES PER PROCESS: {train_num}") - print(f"✦ TOTAL BATCHES (ALL PROCESSES): {train_num * world_size}") - else: - print(f"✦ TOTAL BATCHES: {train_num}") - print(f"✦ SEED: {seed}") - print("=" * 50 + "\n") - - return dataset, train_num - - -def get_distributed_dataloader( - dataset, config, rank=0, world_size=1, seed=42, is_train=True -): - """ - Helper function to get distributed dataloader - - Args: - dataset: PreprocessedDataset instance - config: Configuration dict - rank: Current process rank - world_size: Total number of processes - seed: Random seed - is_train: Whether this is for training (affects shuffling) - - Returns: - dataloader: Distributed DataLoader - sampler: DistributedSampler - """ - if is_train: - return dataset.get_train_dataloader(rank=rank, world_size=world_size, seed=seed) - else: - return dataset.get_val_dataloader(rank=rank, world_size=world_size) - - -def get_data_configs(config): - default_data_config = { - "train_test_split": 0.95, - "split_seed": 42, - "batch_size": 8, - "action_horizon": 21, - "action_history_length": 0, - "image_horizon": 1, - "image_history_length": 0, - "left_padding": False, - "right_padding": False, - "return_first_obs": False, - "return_last_obs": False, - "randomize_obs_after": None, - "datasets": [], - "labeled_pathes": [], - } - data_config = default_data_config | config - data_config["action_horizon"] += 1 - - return data_config - - -class TestDataset(PreprocessedDataset): - def __init__( - self, - dataset, - config, - dataload_config, - normalizer_action, - normalizer_propri, - lerobot_config, - seed=42, - ): - super().__init__( - dataset, - config, - dataload_config, - normalizer_action, - normalizer_propri, - lerobot_config, - seed=seed, - rank=0, - world_size=1, - test_only=True, - ) - - def get_dataloader(self): - """ - Get distributed evaluation dataloader (no shuffling for consistent evaluation) - """ - - dataloader = torch.utils.data.DataLoader( - self, - batch_size=1, - collate_fn=DataCollator( - self.config, - self.dataload_config, - self.normalizer_action, - self.normalizer_propri, - self.lerobot_config, - ), - ) - - return dataloader - - -def load_test_dataset( - config, - lerobot_config, - normalizer_action, - normalizer_propri, - seed=42, - episode=0, -): - """ - Load test dataset - - Args: - config: Model configuration - seed: Random seed for reproducibility (default: 42) - - Returns: - dataset: Test dataset - """ - - # Set seed for reproducibility - torch.manual_seed(seed) - - repo_id = lerobot_config.get("repo_id", None) - assert repo_id is not None, "repo id is required" - root = lerobot_config.get("root", None) - meta_info = LeRobotDatasetMetadata(repo_id, root=root) - dataset_fps = meta_info.fps - dataload_config = get_data_configs(config["data"]) - - norm_stats_path = config.get("norm_stats_path", None) - assert ( - norm_stats_path is not None - ), "norm stats is required, please refer to 'wall-x/scripts/compute_norm_stats.py' to compute stats" - # norm_stats = load_norm_stats(norm_stats_path, repo_id) - - delta_timestamps = { - # action chunk - KEY_MAPPINGS[repo_id]["action"]: [ - t / dataset_fps - for t in range(dataload_config.get("action_horizon", 33) - 1) - ], - } - - dataset = LeRobotDataset( - repo_id, - episodes=[episode], - delta_timestamps=delta_timestamps, - video_backend="pyav", - root=root, - ) - - print(f"Selected episodes: {dataset.episodes}") - print(f"Number of episodes selected: {dataset.num_episodes}") - print(f"Number of frames selected: {dataset.num_frames}") - - dataset = TestDataset( - dataset, - config, - dataload_config, - normalizer_action, - normalizer_propri, - lerobot_config, - seed=seed, - ) - - return dataset diff --git a/wall_x/fusions/backend.py b/wall_x/fusions/backend.py deleted file mode 100644 index 3f9d2e25..00000000 --- a/wall_x/fusions/backend.py +++ /dev/null @@ -1,433 +0,0 @@ -""" -High-performance C++ backend interface for optimized matrix operations. - -This module provides Python bindings for custom CUDA kernels optimized for -transformer and MoE (Mixture of Experts) operations, including: -- Asymmetric dual expert operations -- Token permutation/unpermutation for MoE routing -- RoPE (Rotary Position Embedding) operations -""" - -import torch -from typing import Tuple, Optional -import wallx_csrc as backend - - -def _allocate_asymmetric_dual_outputs( - input_expert0: torch.Tensor, - input_expert1: torch.Tensor, - weight_expert0: torch.Tensor, - weight_expert1: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Allocate output tensors for asymmetric dual expert GEMM operations. - - This function handles the case where two experts may have different output - dimensions, which is common in heterogeneous MoE architectures. - - Args: - input_expert0 (torch.Tensor): Expert 0 input tensor of shape [m0, k] - input_expert1 (torch.Tensor): Expert 1 input tensor of shape [m1, k] - weight_expert0 (torch.Tensor): Expert 0 weight tensor of shape [k, n0] - weight_expert1 (torch.Tensor): Expert 1 weight tensor of shape [k, n1] - - Returns: - Tuple[torch.Tensor, torch.Tensor]: Pre-allocated output tensors - - output_expert0: Shape [m0, n0] - - output_expert1: Shape [m1, n1] - - Raises: - AssertionError: If tensor dimensions are incompatible - """ - # Validate input tensor dimensions - assert input_expert0.ndim == 2, "Expected 2D tensor for input_expert0" - assert input_expert1.ndim == 2, "Expected 2D tensor for input_expert1" - assert weight_expert0.ndim == 2, "Expected 2D tensor for weight_expert0" - assert weight_expert1.ndim == 2, "Expected 2D tensor for weight_expert1" - - # Verify dimension compatibility for matrix multiplication - assert input_expert0.size(1) == weight_expert0.size( - 0 - ), f"Input expert0 K dimension {input_expert0.size(1)} != weight expert0 K dimension {weight_expert0.size(0)}" - assert input_expert1.size(1) == weight_expert1.size( - 0 - ), f"Input expert1 K dimension {input_expert1.size(1)} != weight expert1 K dimension {weight_expert1.size(0)}" - - # Calculate output shapes: [m, k] × [k, n] = [m, n] - m0, n0 = input_expert0.size(0), weight_expert0.size(1) - m1, n1 = input_expert1.size(0), weight_expert1.size(1) - - # Allocate output tensors with matching device and dtype - output_expert0 = torch.empty( - m0, n0, device=input_expert0.device, dtype=input_expert0.dtype - ) - output_expert1 = torch.empty( - m1, n1, device=input_expert1.device, dtype=input_expert1.dtype - ) - - return output_expert0, output_expert1 - - -def asym_dual_gmm_separated( - input_expert0: torch.Tensor, - input_expert1: torch.Tensor, - weight_expert0: torch.Tensor, - weight_expert1: torch.Tensor, - output_expert0: Optional[torch.Tensor] = None, - output_expert1: Optional[torch.Tensor] = None, - trans_a: bool = False, - trans_b: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Asymmetric dual expert grouped GEMM with separated inputs and outputs. - - This is the recommended interface for maximum flexibility and performance when - dealing with two experts that may have different intermediate dimensions. - The operation is equivalent to: - output_expert0 = input_expert0 @ weight_expert0 - output_expert1 = input_expert1 @ weight_expert1 - But optimized as a single fused kernel call. - - Args: - input_expert0 (torch.Tensor): Expert 0 input tensor of shape [m0, k] - input_expert1 (torch.Tensor): Expert 1 input tensor of shape [m1, k] - weight_expert0 (torch.Tensor): Expert 0 weight tensor of shape [k, n0] - weight_expert1 (torch.Tensor): Expert 1 weight tensor of shape [k, n1] - Note: n0 can be different from n1 - output_expert0 (torch.Tensor, optional): Pre-allocated output for expert 0 [m0, n0] - output_expert1 (torch.Tensor, optional): Pre-allocated output for expert 1 [m1, n1] - trans_a (bool, optional): Whether to transpose input tensors. Defaults to False. - trans_b (bool, optional): Whether to transpose weight tensors. Defaults to False. - - Returns: - Tuple[torch.Tensor, torch.Tensor]: Output tensors (output_expert0, output_expert1) - - Example: - >>> # Two experts with different output dimensions - >>> input0 = torch.randn(512, 1024, device='cuda') # 512 tokens for expert 0 - >>> input1 = torch.randn(256, 1024, device='cuda') # 256 tokens for expert 1 - >>> weight0 = torch.randn(1024, 2048, device='cuda') # Expert 0: 1024->2048 - >>> weight1 = torch.randn(1024, 4096, device='cuda') # Expert 1: 1024->4096 - >>> out0, out1 = asym_dual_gmm_separated(input0, input1, weight0, weight1) - """ - # Allocate outputs if not provided - if output_expert0 is None or output_expert1 is None: - alloc_out0, alloc_out1 = _allocate_asymmetric_dual_outputs( - input_expert0, input_expert1, weight_expert0, weight_expert1 - ) - if output_expert0 is None: - output_expert0 = alloc_out0 - if output_expert1 is None: - output_expert1 = alloc_out1 - - # Call optimized C++ backend kernel - backend.asym_dual_gmm( - input_expert0, - input_expert1, - weight_expert0, - weight_expert1, - output_expert0, - output_expert1, - trans_a, - trans_b, - ) - - return output_expert0, output_expert1 - - -def permute( - input: torch.Tensor, - indices: torch.Tensor, - num_out_tokens: int, - workspace: torch.Tensor, - max_expanded_token_num: int, -) -> torch.Tensor: - """ - Permute input tokens according to expert assignment indices for MoE routing. - - This function reorders tokens based on their assigned experts to enable - efficient grouped processing. Used in the forward pass of MoE layers. - - Args: - input (torch.Tensor): Input tokens to permute - indices (torch.Tensor): Expert assignment indices for each token - num_out_tokens (int): Number of output tokens after expansion - workspace (torch.Tensor): Temporary workspace tensor for intermediate computations - max_expanded_token_num (int): Maximum number of tokens after top-k expansion - - Returns: - torch.Tensor: Permuted tokens grouped by expert assignment - - Note: - This is typically used with top-k expert selection where each token - can be routed to multiple experts. - """ - return backend.permute( - input, indices, num_out_tokens, workspace, max_expanded_token_num - ) - - -def unpermute( - input: torch.Tensor, - row_id_map: torch.Tensor, - prob: torch.Tensor, - max_tokens: int, - num_topK: int, -) -> torch.Tensor: - """ - Unpermute expert outputs back to original token order with probability weighting. - - This function reverses the permutation applied in the forward pass and combines - outputs from multiple experts using their routing probabilities. - - Args: - input (torch.Tensor): Permuted expert outputs to unpermute - row_id_map (torch.Tensor): Mapping from permuted positions to original positions - prob (torch.Tensor): Expert routing probabilities for weighted combination - max_tokens (int): Maximum number of tokens in the sequence - num_topK (int): Number of top experts selected per token - - Returns: - torch.Tensor: Unpermuted tokens in original order with expert outputs combined - - Note: - The output combines multiple expert predictions for each token using - the routing probabilities as weights. - """ - return backend.unpermute(input, row_id_map, prob, max_tokens, num_topK) - - -def unpermute_bwd( - input_bwd: torch.Tensor, - input_fwd: torch.Tensor, - row_id_map: torch.Tensor, - prob: Optional[torch.Tensor], -) -> torch.Tensor: - """ - Backward pass for unpermute operation with gradient flow. - - This function handles the backward pass through the unpermute operation, - ensuring proper gradient flow for training MoE models. - - Args: - input_bwd (torch.Tensor): Backward gradients from the next layer - input_fwd (torch.Tensor): Forward pass inputs (for gradient computation) - row_id_map (torch.Tensor): Row mapping used in forward unpermute - prob (torch.Tensor, optional): Expert probabilities. If None, uniform weights are used. - - Returns: - torch.Tensor: Gradients with respect to the input of unpermute forward pass - - Note: - If prob is None, uniform probabilities are assumed for gradient computation. - """ - # Handle case where probabilities are not provided - if prob is None: - prob = torch.ones( - [input_bwd.size(0), 1], dtype=torch.float32, device=input_bwd.device - ) - - return backend.unpermute_bwd(input_bwd, input_fwd, row_id_map, prob) - - -def rope( - q: torch.Tensor, - k: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - q_out: torch.Tensor, - k_out: torch.Tensor, - mrope_section_doubled: bool, -) -> None: - """ - Apply RoPE (Rotary Position Embedding) to query and key tensors. - - Applies rotary position embeddings to query and key tensors using precomputed - cosine and sine values. Supports both standard RoPE and multi-dimensional RoPE (mRoPE). - - Args: - q (torch.Tensor): Query tensor to apply RoPE to - k (torch.Tensor): Key tensor to apply RoPE to - cos (torch.Tensor): Precomputed cosine values for rotation - sin (torch.Tensor): Precomputed sine values for rotation - q_out (torch.Tensor): Output tensor for rotated queries (in-place operation supported) - k_out (torch.Tensor): Output tensor for rotated keys (in-place operation supported) - mrope_section_doubled (bool): Whether using multi-dimensional RoPE with doubled sections - - Note: - This function performs in-place operations if q_out and k_out point to the same - memory as q and k respectively. The rotation is applied using the standard - RoPE formulation with complex number rotation. - """ - return backend.rope(q, k, cos, sin, q_out, k_out, mrope_section_doubled) - - -def rope_bwd( - grad_q_out: torch.Tensor, - grad_k_out: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - grad_q: torch.Tensor, - grad_k: torch.Tensor, - mrope_section_doubled: bool, -) -> None: - """ - Backward pass for RoPE operation with gradient computation. - - Computes gradients with respect to the input query and key tensors - for the RoPE operation used in transformer attention mechanisms. - - Args: - grad_q_out (torch.Tensor): Gradient with respect to output queries - grad_k_out (torch.Tensor): Gradient with respect to output keys - q (torch.Tensor): Original query tensor from forward pass - k (torch.Tensor): Original key tensor from forward pass - cos (torch.Tensor): Cosine values used in forward pass - sin (torch.Tensor): Sine values used in forward pass - grad_q (torch.Tensor): Output tensor for query gradients - grad_k (torch.Tensor): Output tensor for key gradients - mrope_section_doubled (bool): Whether using multi-dimensional RoPE configuration - - Note: - This function computes the analytical gradient of the RoPE operation, - which involves the inverse rotation compared to the forward pass. - """ - return backend.rope_bwd( - grad_q_out, grad_k_out, q, k, cos, sin, grad_q, grad_k, mrope_section_doubled - ) - - -def get_rope_index( - input_ids: torch.Tensor, - image_grid_thw: Optional[torch.Tensor], - video_grid_thw: Optional[torch.Tensor], - second_per_grid_ts: Optional[torch.Tensor], - attention_mask: Optional[torch.Tensor], - spatial_merge_size: int, - image_token_id: int, - video_token_id: int, - vision_start_token_id: int, - tokens_per_second: float, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Generate position indices for multimodal RoPE (Rotary Position Embedding). - - This function computes 3D position indices for text, image, and video tokens - to enable proper spatial-temporal position encoding in multimodal transformers. - - Args: - input_ids (torch.Tensor): Input token IDs of shape [batch_size, seq_len] - image_grid_thw (torch.Tensor, optional): Image grid specifications of shape [num_images, 3] (T, H, W) - video_grid_thw (torch.Tensor, optional): Video grid specifications of shape [num_videos, 3] (T, H, W) - second_per_grid_ts (torch.Tensor, optional): Temporal scaling per video grid of shape [num_videos] - attention_mask (torch.Tensor, optional): Attention mask of shape [batch_size, seq_len] - spatial_merge_size (int): Spatial dimension merge factor for patch grouping - image_token_id (int): Token ID representing image patches - video_token_id (int): Token ID representing video frames - vision_start_token_id (int): Token ID marking vision sequence start - tokens_per_second (float): Temporal scaling factor for video sequences - - Returns: - Tuple[torch.Tensor, torch.Tensor]: A tuple containing: - - position_ids: 3D position indices of shape [3, batch_size, seq_len] - - mrope_deltas: Position deltas for multimodal RoPE of shape [batch_size, 1] - - Note: - When both image_grid_thw and video_grid_thw are None, returns standard - text-only position indices based on attention_mask or sequence order. - """ - return backend.rope_index( - input_ids, - image_grid_thw, - video_grid_thw, - second_per_grid_ts, - attention_mask, - spatial_merge_size, - image_token_id, - video_token_id, - vision_start_token_id, - tokens_per_second, - ) - - -def rot_pos_emb( - inv_freq: torch.Tensor, - grid_thw: torch.Tensor, - spatial_merge_size: int, -) -> torch.Tensor: - """ - Compute fused rotary position embeddings for multimodal grids. - - This function efficiently computes rotary position embeddings for spatial-temporal - grids using a fused CUDA kernel, supporting both int32 and int64 grid specifications. - - Args: - inv_freq (torch.Tensor): Inverse frequencies for RoPE of shape [dim/2] - Must be float32 dtype on CUDA device - grid_thw (torch.Tensor): Grid specifications of shape [num_grids, 3] (T, H, W) - Supports int32 or int64 dtype on CUDA device - spatial_merge_size (int): Merge factor for spatial dimensions (must be positive) - - Returns: - torch.Tensor: Computed rotary embeddings of shape [total_tokens, dim] - where total_tokens is determined by grid layouts and spatial_merge_size - - Example: - >>> inv_freq = torch.randn(64, device='cuda', dtype=torch.float32) # 128-dim model - >>> grids = torch.tensor([[8, 14, 14], [16, 7, 7]], device='cuda', dtype=torch.int32) - >>> embeddings = rot_pos_emb(inv_freq, grids, spatial_merge_size=2) - >>> print(embeddings.shape) # [computed_tokens, 128] - - Note: - The function automatically dispatches to int32 or int64 implementations - based on the dtype of grid_thw. Output is always float32. - """ - return backend.rot_pos_emb(inv_freq, grid_thw, spatial_merge_size) - - -def get_window_index( - grid_thw: torch.Tensor, - spatial_merge_size: int, - vit_merger_window_size: int, - patch_size: int, - spatial_merge_unit: int, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Generate window attention indices for Vision Transformer architectures. - - Computes window-based attention indices for hierarchical processing of vision - tokens, enabling efficient sliding window attention patterns in ViT models. - - Args: - grid_thw (torch.Tensor): Grid specifications of shape [num_grids, 3] (T, H, W) - Must be int32 dtype on CUDA device - spatial_merge_size (int): Spatial dimension merge factor - vit_merger_window_size (int): Size of attention windows for ViT processing - patch_size (int): Size of vision patches in pixels - spatial_merge_unit (int): Unit size for spatial merging operations - - Returns: - Tuple[torch.Tensor, torch.Tensor]: A tuple containing: - - window_indices: Flattened window indices of shape [total_elements] - - cu_window_seqlens: Cumulative window sequence lengths of shape [num_windows + 1] - - Example: - >>> grids = torch.tensor([[1, 14, 14]], device='cuda', dtype=torch.int32) - >>> indices, seqlens = get_window_index( - ... grids, spatial_merge_size=2, vit_merger_window_size=7, - ... patch_size=16, spatial_merge_unit=4 - ... ) - - Note: - Returns empty tensors if input grid is empty or no valid windows can be formed. - The cu_window_seqlens tensor enables efficient batched attention computation. - """ - return backend.get_window_index( - grid_thw, - spatial_merge_size, - vit_merger_window_size, - patch_size, - spatial_merge_unit, - ) diff --git a/wall_x/fusions/ops.py b/wall_x/fusions/ops.py deleted file mode 100644 index 3f993b8a..00000000 --- a/wall_x/fusions/ops.py +++ /dev/null @@ -1,742 +0,0 @@ -import torch -import warnings -from wall_x.fusions import backend - - -class AsymmetricDualExpertGemm(torch.autograd.Function): - @staticmethod - def forward( - ctx, input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False - ): - """ - Forward pass for asymmetric dual expert GEMM. - - Args: - input_expert0: Expert 0 input [m0, k] - input_expert1: Expert 1 input [m1, k] - weight_expert0: Expert 0 weight [k, n0] or [n0, k] if trans_b=True - weight_expert1: Expert 1 weight [k, n1] or [n1, k] if trans_b=True - trans_b: Whether to transpose the weight matrices - - Returns: - Tuple of (output_expert0, output_expert1) - """ - # Validate inputs - assert input_expert0.dim() == 2, "input_expert0 must be 2D" - assert input_expert1.dim() == 2, "input_expert1 must be 2D" - assert weight_expert0.dim() == 2, "weight_expert0 must be 2D" - assert weight_expert1.dim() == 2, "weight_expert1 must be 2D" - - # Dimension validation depends on trans_b - if trans_b: - assert input_expert0.size(1) == weight_expert0.size( - 1 - ), "Expert 0 dimension mismatch (trans_b=True)" - assert input_expert1.size(1) == weight_expert1.size( - 1 - ), "Expert 1 dimension mismatch (trans_b=True)" - else: - assert input_expert0.size(1) == weight_expert0.size( - 0 - ), "Expert 0 dimension mismatch (trans_b=False)" - assert input_expert1.size(1) == weight_expert1.size( - 0 - ), "Expert 1 dimension mismatch (trans_b=False)" - - # Save tensors and trans_b for backward pass - ctx.save_for_backward( - input_expert0, input_expert1, weight_expert0, weight_expert1 - ) - ctx.trans_b = trans_b - - # Allocate output tensors - m0 = input_expert0.size(0) - m1 = input_expert1.size(0) - n0 = weight_expert0.size(0) if trans_b else weight_expert0.size(1) - n1 = weight_expert1.size(0) if trans_b else weight_expert1.size(1) - - output_expert0 = torch.empty( - m0, n0, device=input_expert0.device, dtype=input_expert0.dtype - ) - output_expert1 = torch.empty( - m1, n1, device=input_expert1.device, dtype=input_expert1.dtype - ) - - # Call the backend C++ function - backend.asym_dual_gmm_separated( - input_expert0, - input_expert1, - weight_expert0, - weight_expert1, - output_expert0, - output_expert1, - trans_b=trans_b, - ) - - return output_expert0, output_expert1 - - @staticmethod - def backward(ctx, grad_output_expert0, grad_output_expert1): - """ - Optimized backward pass using specialized kernels. - Always computes all gradients to minimize kernel calls. - """ - grad_output_expert0 = grad_output_expert0.contiguous() - grad_output_expert1 = grad_output_expert1.contiguous() - - input_expert0, input_expert1, weight_expert0, weight_expert1 = ctx.saved_tensors - trans_b = ctx.trans_b - - # Always allocate all gradient tensors (no conditional computation) - grad_input_expert0 = torch.empty_like(input_expert0) - grad_input_expert1 = torch.empty_like(input_expert1) - grad_weight_expert0 = torch.empty_like(weight_expert0) - grad_weight_expert1 = torch.empty_like(weight_expert1) - - # Compute input gradients: grad_input = grad_output @ weight^T (if trans_b=False) - # = grad_output @ weight (if trans_b=True) - backend.asym_dual_gmm_separated( - grad_output_expert0, - grad_output_expert1, - weight_expert0, - weight_expert1, - grad_input_expert0, - grad_input_expert1, - trans_a=False, - trans_b=not trans_b, - ) - - # Compute weight gradients - if trans_b: - # When trans_b=True in forward: output = input @ weight^T - # So grad_weight^T = input^T @ grad_output - # Which means grad_weight = grad_output^T @ input - backend.asym_dual_gmm_separated( - grad_output_expert0, - grad_output_expert1, - input_expert0, - input_expert1, - grad_weight_expert0, - grad_weight_expert1, - trans_a=True, - trans_b=False, - ) - else: - # When trans_b=False in forward: output = input @ weight - # So grad_weight = input^T @ grad_output - backend.asym_dual_gmm_separated( - input_expert0, - input_expert1, - grad_output_expert0, - grad_output_expert1, - grad_weight_expert0, - grad_weight_expert1, - trans_a=True, - trans_b=False, - ) - - return ( - grad_input_expert0, - grad_input_expert1, - grad_weight_expert0, - grad_weight_expert1, - None, - ) - - -def asym_dual_gmm( - input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False -): - """ - Convenience function for asymmetric dual expert GEMM. - - Args: - input_expert0: Expert 0 input [m0, k] - input_expert1: Expert 1 input [m1, k] - weight_expert0: Expert 0 weight [k, n0] or [n0, k] if trans_b=True - weight_expert1: Expert 1 weight [k, n1] or [n1, k] if trans_b=True - trans_b: Whether to transpose the weight matrices - - Returns: - Tuple of (output_expert0, output_expert1) - """ - return AsymmetricDualExpertGemm.apply( - input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b - ) - - -################################################################################################ -## -## PermuteMoE topK -## -################################################################################################ - - -class PermuteMoE_topK(torch.autograd.Function): - - workspace_fw = None - dtype = None - max_expanded_token_num = 0 - - @staticmethod - def forward( - ctx, - input_act: torch.Tensor, - indices: torch.Tensor, - num_out_tokens: int, - max_token_num: int, - ): - """ - indices: for topK=1, indices in a 1-d tensor of shape [num_tokens], - otherwise, it's a 2-d tensor of shape [num_tokens, topK] - """ - # Empty input check - if not input_act.numel(): - return input_act, None - - # For top1 case, view the indices as 2D tensor to unify the shape for topk>=2 cases. - if indices.dim() == 1: - indices = indices.view(-1, 1) - - # Device check - if input_act.is_cpu: - raise RuntimeError( - "[Error] The input `input_act` of permute_topK op is on the device: CPU!" - ) - if indices.is_cpu: - warnings.warn( - "The input `indices` of permute_topK op is on the device: CPU!" - ) - - # Shape check - if input_act.size(0) != indices.size(0): - raise RuntimeError( - f"[Error] permute_topK op input `indices` shape mismatch! " - f"Expect {input_act.size(0)}, but got {indices.size(0)}." - ) - - # Data type check - if indices.dtype != torch.int32: - warnings.warn( - f"The data type of the input `indices` of permute_topK op is {indices.dtype}! " - "The recommended type is torch.int32." - ) - indices = indices.to(torch.int32) - - # Contiguous check - if not input_act.is_contiguous(): - warnings.warn("The input `input_act` of permute_topK op is discontiguous!") - input_act = input_act.contiguous() - if not indices.is_contiguous(): - warnings.warn("The input `indices` of permute_topK op is discontiguous!") - indices = indices.contiguous() - - num_topK = indices.size(1) - - input_max_expanded_token_num = max(max_token_num, input_act.size(0)) * num_topK - if PermuteMoE_topK.max_expanded_token_num < input_max_expanded_token_num: - PermuteMoE_topK.max_expanded_token_num = input_max_expanded_token_num - PermuteMoE_topK.workspace_fw = [] - - if PermuteMoE_topK.dtype != input_act.dtype: - PermuteMoE_topK.dtype = input_act.dtype - PermuteMoE_topK.workspace_fw = [] - - permuted_act, row_id_map, PermuteMoE_topK.workspace_fw = backend.permute( - input_act, - indices, - num_out_tokens, - PermuteMoE_topK.workspace_fw, - PermuteMoE_topK.max_expanded_token_num, - ) - - ctx.row_id_map = row_id_map - ctx.num_tokens = indices.size(0) - ctx.num_topK = num_topK - return permuted_act, row_id_map - - @staticmethod - def backward(ctx, permuted_act_grad, _): - # Empty input check - if not permuted_act_grad.numel(): - return permuted_act_grad, None, None, None - - if not permuted_act_grad.is_contiguous(): - permuted_act_grad = permuted_act_grad.contiguous() - - row_id_map = ctx.row_id_map - num_tokens = ctx.num_tokens - num_topK = ctx.num_topK - - unpermuted_act_grad = backend.unpermute( - permuted_act_grad, row_id_map, torch.tensor([]), num_tokens, num_topK - ) - return unpermuted_act_grad, None, None, None - - -################################################################################################ -## -## UnpermuteMoE topK -## -################################################################################################ - - -class UnpermuteMoE_topK(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - input_act: torch.Tensor, - row_id_map: torch.Tensor, - probs: torch.Tensor = None, - ): - # Empty input check - if not input_act.numel(): - ctx.probs = probs - return input_act - - # Device check - if input_act.is_cpu: - raise RuntimeError( - "[Error] The input `input_act` of unpermute_topK op is on the device: CPU!" - ) - if row_id_map.is_cpu: - warnings.warn( - "The input `row_id_map` of unpermute_topK op is on the device: CPU!" - ) - row_id_map = row_id_map.cuda() - if probs is not None and probs.is_cpu: - warnings.warn( - "The input `probs` of unpermute_topK op is on the device: CPU!" - ) - probs = probs.cuda() - - # Shape check - if probs is not None and row_id_map.size(0) != probs.size(0) * probs.size(1): - raise RuntimeError( - f"[Error] unpermute_topK op input `probs` shape mismatch! " - f"Expect {row_id_map.size(0)}, but got {probs.size(0) * probs.size(1)}." - ) - - # Data type check - if row_id_map.dtype != torch.int32: - warnings.warn( - f"The data type of the input `row_id_map` of unpermute_topK op is {row_id_map.dtype}! " - "The recommended type is torch.int32." - ) - row_id_map = row_id_map.to(torch.int32) - if probs is not None and probs.dtype != torch.float32: - warnings.warn( - f"The data type of the input `probs` of unpermute_topK op is {probs.dtype}! " - "The recommended type is torch.float32." - ) - probs = probs.to(torch.float32) - - # Contiguous check - if not input_act.is_contiguous(): - warnings.warn( - "The input `input_act` of unpermute_topK op is discontiguous!" - ) - input_act = input_act.contiguous() - if not row_id_map.is_contiguous(): - warnings.warn( - "The input `row_id_map` of unpermute_topK op is discontiguous!" - ) - row_id_map = row_id_map.contiguous() - if probs is not None and not probs.is_contiguous(): - warnings.warn("The input `probs` of unpermute_topK op is discontiguous!") - probs = probs.contiguous() - - num_tokens = probs.size(0) if probs is not None else input_act.size(0) - num_topK = probs.size(1) if probs is not None else 1 - - unpermuted_output = backend.unpermute( - input_act, - row_id_map, - probs if probs is not None else torch.tensor([]), - num_tokens, - num_topK, - ) - - ctx.save_for_backward(input_act, row_id_map, probs) - return unpermuted_output - - @staticmethod - def backward(ctx, unpermuted_act_grad): - # Empty input check - if not unpermuted_act_grad.numel(): - return unpermuted_act_grad, None, ctx.probs - - if not unpermuted_act_grad.is_contiguous(): - unpermuted_act_grad = unpermuted_act_grad.contiguous() - - input_act, row_id_map, probs = ctx.saved_tensors - - act_grad = None - if ctx.needs_input_grad[0]: - act_grad, prob_grad = backend.unpermute_bwd( - unpermuted_act_grad, input_act, row_id_map, probs - ) - - if not ctx.needs_input_grad[2]: - prob_grad = None - return act_grad, None, prob_grad - - -def permute(input_act, indices, num_out_tokens=None, max_token_num=0): - num_out_tokens = 0 if num_out_tokens is None else num_out_tokens - return PermuteMoE_topK.apply(input_act, indices, num_out_tokens, max_token_num) - - -def unpermute(input_act, row_id_map, probs=None): - return UnpermuteMoE_topK.apply(input_act, row_id_map, probs) - - -################################################################################################ -## -## mutlimodal RoPE -## -################################################################################################ - - -class MultimodalRoPE(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - mrope_section: list, - ): - - # Device check - if q.is_cpu: - raise RuntimeError( - "[Error] The input `q` of multimodal_rope op is on the device: CPU!" - ) - if k.is_cpu: - raise RuntimeError( - "[Error] The input `k` of multimodal_rope op is on the device: CPU!" - ) - if cos.is_cpu: - raise RuntimeError( - "[Error] The input `cos` of multimodal_rope op is on the device: CPU!" - ) - if sin.is_cpu: - raise RuntimeError( - "[Error] The input `sin` of multimodal_rope op is on the device: CPU!" - ) - if len(mrope_section) != 3: - raise RuntimeError( - "[Error] The input `mrope_section` of multimodal_rope op must be a list of 3 integers!" - ) - - # Contiguous check - if not q.is_contiguous(): - warnings.warn("The input `q` of multimodal_rope op is discontiguous!") - q = q.contiguous() - if not k.is_contiguous(): - warnings.warn("The input `k` of multimodal_rope op is discontiguous!") - k = k.contiguous() - if not cos.is_contiguous(): - warnings.warn("The input `cos` of multimodal_rope op is discontiguous!") - cos = cos.contiguous() - if not sin.is_contiguous(): - warnings.warn("The input `sin` of multimodal_rope op is discontiguous!") - sin = sin.contiguous() - - # Prepare mrope_section_doubled - mrope_section_doubled = [x * 2 for x in mrope_section] - - # Create output tensors - q_out = torch.empty_like(q) - k_out = torch.empty_like(k) - backend.rope(q, k, cos, sin, q_out, k_out, mrope_section_doubled) - - ctx.save_for_backward(q, k, cos, sin) - ctx.mrope_section_doubled = mrope_section_doubled - - return q_out, k_out - - @staticmethod - def backward(ctx, grad_q_out, grad_k_out): - - if not grad_q_out.is_contiguous(): - grad_q_out = grad_q_out.contiguous() - if not grad_k_out.is_contiguous(): - grad_k_out = grad_k_out.contiguous() - - q, k, cos, sin = ctx.saved_tensors - - grad_q = None - grad_k = None - if ctx.needs_input_grad[0]: - grad_q = torch.empty_like(q) - if ctx.needs_input_grad[1]: - grad_k = torch.empty_like(k) - - if grad_q is not None or grad_k is not None: - backend.rope_bwd( - grad_q_out, - grad_k_out, - q, - k, - cos, - sin, - grad_q if grad_q is not None else torch.empty_like(q), - grad_k if grad_k is not None else torch.empty_like(k), - ctx.mrope_section_doubled, - ) - - return grad_q, grad_k, None, None, None - - -def multimodal_rope(q, k, cos, sin, mrope_section): - return MultimodalRoPE.apply(q, k, cos, sin, mrope_section) - - -################################################################################################ -## -## RoPE Index 3D -## -################################################################################################ - - -def get_rope_index( - input_ids: torch.Tensor, - spatial_merge_size: int, - image_token_id: int, - video_token_id: int, - vision_start_token_id: int, - tokens_per_second: float, - image_grid_thw: torch.Tensor = None, - video_grid_thw: torch.Tensor = None, - second_per_grid_ts: torch.Tensor = None, - attention_mask: torch.Tensor = None, -): - """ - Generate 3D RoPE position indices for multimodal transformer inputs. - - Computes position indices for text, image, and video tokens to enable proper - spatial-temporal position encoding in multimodal transformers with RoPE. - - Args: - input_ids (torch.Tensor): Input token sequence of shape [batch_size, seq_len] - Must be LongTensor on CUDA device - spatial_merge_size (int): Spatial merge size for patch grouping (must be positive) - image_token_id (int): Token ID representing image patches - video_token_id (int): Token ID representing video frames - vision_start_token_id (int): Token ID marking start of vision sequences - tokens_per_second (float): Temporal scaling factor for video sequences (must be positive) - image_grid_thw (torch.Tensor, optional): Image grid dimensions of shape [num_images, 3] (T, H, W) - video_grid_thw (torch.Tensor, optional): Video grid dimensions of shape [num_videos, 3] (T, H, W) - second_per_grid_ts (torch.Tensor, optional): Video time intervals of shape [num_videos] - attention_mask (torch.Tensor, optional): Attention mask of shape [batch_size, seq_len] - - Returns: - Tuple[torch.Tensor, torch.Tensor]: A tuple containing: - - position_ids: 3D position indices of shape [3, batch_size, seq_len] - - mrope_position_deltas: mRoPE position deltas of shape [batch_size, 1] - - Raises: - TypeError: If input_ids is not a torch.Tensor - ValueError: If input dimensions are incorrect or tensors not on CUDA - """ - # Input validation - if not isinstance(input_ids, torch.Tensor): - raise TypeError("input_ids must be a torch.Tensor") - - if input_ids.dim() != 2: - raise ValueError("input_ids must be 2D tensor (batch_size, seq_len)") - - if not input_ids.is_cuda: - raise ValueError("input_ids must be on CUDA device") - - # Parameter validation - if not isinstance(spatial_merge_size, int) or spatial_merge_size <= 0: - raise ValueError( - f"spatial_merge_size must be positive integer, got {spatial_merge_size}" - ) - - if not isinstance(tokens_per_second, (int, float)) or tokens_per_second <= 0: - raise ValueError( - f"tokens_per_second must be positive number, got {tokens_per_second}" - ) - - return backend.get_rope_index( - input_ids, - image_grid_thw, - video_grid_thw, - second_per_grid_ts, - attention_mask, - spatial_merge_size, - image_token_id, - video_token_id, - vision_start_token_id, - float(tokens_per_second), - ) - - -################################################################################################ -## -## Fused Rotary Position Embedding -## -################################################################################################ - - -def rot_pos_emb( - inv_freq: torch.Tensor, - grid_thw: torch.Tensor, - spatial_merge_size: int, -) -> torch.Tensor: - """ - Compute fused rotary position embeddings using optimized CUDA kernel. - - This function fuses all rotary position embedding computations into a single - CUDA kernel for improved performance with spatial-temporal grids. - - Args: - inv_freq (torch.Tensor): Inverse frequencies tensor of shape [dim/2] - Contains precomputed 1.0 / (theta ** (torch.arange(0, dim, 2) / dim)) - Must be float32 on CUDA device - grid_thw (torch.Tensor): Grid dimensions tensor of shape [num_grids, 3] - Each row contains (T, H, W) for temporal, height, width dimensions - Supports int32 or int64 on CUDA device - spatial_merge_size (int): Spatial merge size for token grouping (must be positive) - - Returns: - torch.Tensor: Rotary position embeddings of shape [total_tokens, dim] - where dim = 2 * len(inv_freq) - First half contains h_pos frequencies, second half contains w_pos frequencies - - Raises: - TypeError: If inputs are not torch.Tensor or spatial_merge_size not int - ValueError: If tensor dimensions incorrect, not on CUDA, or devices mismatch - RuntimeError: If CUDA kernel execution fails - """ - # Type checking - if not isinstance(inv_freq, torch.Tensor): - raise TypeError(f"inv_freq must be a torch.Tensor, got {type(inv_freq)}") - - if not isinstance(grid_thw, torch.Tensor): - raise TypeError(f"grid_thw must be a torch.Tensor, got {type(grid_thw)}") - - # Dimension checking - if inv_freq.dim() != 1: - raise ValueError( - f"inv_freq must be 1-dimensional, got {inv_freq.dim()}D tensor" - ) - - if grid_thw.dim() != 2: - raise ValueError( - f"grid_thw must be 2-dimensional, got {grid_thw.dim()}D tensor" - ) - - if grid_thw.size(1) != 3: - raise ValueError( - f"grid_thw must have shape [num_grids, 3], got shape {list(grid_thw.shape)}" - ) - - # Device checking - if not inv_freq.is_cuda: - raise ValueError("inv_freq must be on CUDA device") - - if not grid_thw.is_cuda: - raise ValueError("grid_thw must be on CUDA device") - - # Ensure both tensors are on the same device - if inv_freq.device != grid_thw.device: - raise ValueError( - f"inv_freq and grid_thw must be on the same device, " - f"got {inv_freq.device} and {grid_thw.device}" - ) - - # Parameter validation - if not isinstance(spatial_merge_size, int): - raise TypeError( - f"spatial_merge_size must be an integer, got {type(spatial_merge_size)}" - ) - - if spatial_merge_size <= 0: - raise ValueError( - f"spatial_merge_size must be positive, got {spatial_merge_size}" - ) - - # Ensure inv_freq is float32 (the kernel expects float) - if inv_freq.dtype != torch.float32: - inv_freq = inv_freq.to(torch.float32) - - # Call the CUDA backend - try: - return backend.rot_pos_emb(inv_freq, grid_thw, spatial_merge_size) - except RuntimeError as e: - raise RuntimeError(f"CUDA kernel execution failed: {str(e)}") - - -################################################################################################ -## -## Fused Window Index Generation -## -################################################################################################ - - -def get_window_index( - grid_thw: torch.Tensor, - window_size: int, - spatial_merge_size: int, - patch_size: int, - spatial_merge_unit: int = 1, -): - """ - Generate window attention indices for Vision Transformer architectures. - - Computes window-based attention indices for hierarchical processing of vision - tokens, enabling efficient sliding window attention patterns in ViT models. - - Args: - grid_thw (torch.Tensor): Grid specifications of shape [num_grids, 3] (T, H, W) - Must be or will be converted to int32 on CUDA device - window_size (int): Window size for attention computation - spatial_merge_size (int): Spatial merge size for patch grouping - patch_size (int): Size of vision patches in pixels - spatial_merge_unit (int, optional): Spatial merging unit size. Defaults to 1. - - Returns: - Tuple[torch.Tensor, torch.Tensor]: A tuple containing: - - window_index: Window indices tensor of shape [total_elements] - - cu_window_seqlens: Cumulative window sequence lengths of shape [num_windows + 1] - - Raises: - AssertionError: If grid_thw dimensions are incorrect - - Note: - Returns empty tensors if input grid is empty or no valid windows can be formed. - The function automatically converts input to CUDA int32 if needed. - """ - # Input validation - assert ( - grid_thw.dim() == 2 and grid_thw.size(1) == 3 - ), f"grid_thw must have shape (num_grids, 3), got {grid_thw.shape}" - - # Ensure input is on CUDA and int32 type - if not grid_thw.is_cuda: - grid_thw = grid_thw.cuda() - - if grid_thw.dtype != torch.int32: - grid_thw = grid_thw.to(torch.int32) - - # Calculate vit_merger_window_size - vit_merger_window_size = window_size // spatial_merge_size // patch_size - - # Call CUDA backend - window_index, cu_window_seqlens = backend.get_window_index( - grid_thw, - spatial_merge_size, - vit_merger_window_size, - patch_size, - spatial_merge_unit, - ) - - return window_index, cu_window_seqlens diff --git a/wall_x/infer/env.py b/wall_x/infer/env.py deleted file mode 100644 index 35106f75..00000000 --- a/wall_x/infer/env.py +++ /dev/null @@ -1,143 +0,0 @@ -""" -Base Environment Class for Robot Control and Inference -""" - -from typing import Dict, Any, List -from abc import ABC, abstractmethod -import time -from wall_x.infer.infer_config import InferConfig -from wall_x.infer.utils import KeyboardThread -from wall_x.infer.logger import InferLogger - - -class BaseEnv(ABC): - def __init__(self, config: InferConfig): - self.config = config - self.logger = InferLogger.get_env_logger("Env") - - @abstractmethod - def get_observation(self) -> Dict[str, Any]: - raise NotImplementedError - - @abstractmethod - def apply_action(self, input: dict) -> None: - raise NotImplementedError - - @abstractmethod - def get_instruction(self) -> str: - raise NotImplementedError - - def reset(self) -> Dict[str, Any]: - raise NotImplementedError - - def stop(self) -> None: - raise NotImplementedError - - -class RealRobotEnv(BaseEnv): - def __init__( - self, config: InferConfig, instructions: List[str], enable_keyboard: bool = True - ): - """ - Args: - config: Inference configuration - instruction: Task instruction - """ - super().__init__(config) - self.instruction = "test" - self.model = self._register_model() - self.robot = self._register_robot() - - # Keyboard control - self.keyboard_thread = None - if enable_keyboard: - self.keyboard_thread = KeyboardThread() - - # Instruction list - self.instructions = instructions - self.instruction_index = 0 - - # def _register_model(self) -> WallxModelWrapper: - # return WallxModelWrapper(self.config) - - def _register_robot(self): - from wall_x.infer.robot import DesktopRobot, TurtleRobot - - if self.config.robot_type == "desktop": - return DesktopRobot(self.config) - elif self.config.robot_type == "turtle": - return TurtleRobot(self.config) - else: - raise ValueError(f"Invalid robot type: {self.config.robot_type}") - - def get_observation(self): - return self.robot.get_observation() - - def apply_action(self, input: dict): - self.robot.apply_action(input) - - def get_instruction(self) -> str: - """Return task instruction""" - return self.instructions[self.instruction_index] - - def reset(self): - self.robot.go_home() - - def listen_to_keyboard(self): - if self.keyboard_thread is not None: - if self.keyboard_thread.should_stop: - time.sleep(1) - return True - if self.keyboard_thread.should_reset: - self.reset() - self.keyboard_thread.should_reset = False - time.sleep(1) - return True - if self.keyboard_thread.new_instruction_index is not None: - new_index = self.keyboard_thread.new_instruction_index - # Check if index is valid - if 0 <= new_index < len(self.instructions): - self.instruction_index = new_index - self.logger.info( - f"[Keyboard] Instruction index switched to {new_index}: {self.instructions[new_index]}" - ) - else: - self.logger.info( - f"[Keyboard] Invalid instruction index {new_index}, valid range: 0-{len(self.instructions)-1}" - ) - # Reset flag - self.keyboard_thread.new_instruction_index = None - time.sleep(1) - return True - return False - - def run_infer_flow_action(self): - while True: - if self.listen_to_keyboard(): - continue - observation = self.get_observation() - instruction = self.get_instruction() - model_output = self.model.infer_flow_action(observation, instruction) - self.apply_action(model_output) - - def run_infer_flow_action_with_subtask(self, subtask_interval: int = 2): - step = 0 - subtask = "" - while True: - if self.listen_to_keyboard(): - continue - observation = self.get_observation() - instruction = self.get_instruction() - if step == 0 or step % subtask_interval == 0: - subtask = self.model.infer_subtask(observation, instruction) - model_output = self.model.infer_flow_action(observation, subtask) - self.apply_action(model_output) - - def run_infer_ar_action(self): - while True: - if self.listen_to_keyboard(): - continue - observation = self.get_observation() - instruction = self.get_instruction() - model_output = self.model.infer_ar_action(observation, instruction) - self.apply_action(model_output) diff --git a/wall_x/infer/env_libero.py b/wall_x/infer/env_libero.py deleted file mode 100644 index 07fac47a..00000000 --- a/wall_x/infer/env_libero.py +++ /dev/null @@ -1,744 +0,0 @@ -import os -import json -import numpy as np -from typing import Dict, Any, Tuple, List -from libero.libero import benchmark -from wall_x.infer.env import BaseEnv, InferConfig -from wall_x.serving.policy.wall_x_policy import WallXPolicy - -from wall_x.infer.base_dataclass import RobotStateActionData -from wall_x.infer.utils_libero import ( - get_libero_env, - get_libero_dummy_action, - get_libero_image, - get_libero_wrist_image, - quat2axisangle, - TaskSuite, - save_rollout_video, -) -from robosuite.wrappers import VisualizationWrapper - - -def _create_libero_env_standalone( - task_id: int, - task_suite_name: str, - model_family: str = "wallx", - resolution: int = 256, - seed: int = 7, -) -> Any: - """ - Standalone function to create a Libero environment, independent of LiberoRobotEnv instance. - Used for creating environments in subprocess during multi-batch inference, - avoiding serialization of large objects containing the model. - - Args: - task_id: Task ID - task_suite_name: Task suite name - model_family: Model family - resolution: Resolution - seed: Random seed - - Returns: - Environment instance - """ - from libero.libero import benchmark - - # Get task suite and task - benchmark_dict = benchmark.get_benchmark_dict() - task_suite = benchmark_dict[task_suite_name]() - task = task_suite.get_task(task_id) - - # Create environment - env, _ = get_libero_env( - task, - model_family=model_family, - resolution=resolution, - seed=seed, - ) - - # Wrap environment - env.env = VisualizationWrapper(env.env) - env.env.set_visualization_setting(setting="grippers", visible=False) - - return env - - -class LiberoRobotEnv(BaseEnv): - def __init__( - self, - config: InferConfig, - task_suite_name: str = TaskSuite.LIBERO_SPATIAL, - initial_states_path: str = "DEFAULT", - rollout_dir: str = "./rollouts", - model_family: str = "wallx", - resolution: int = 256, - seed: int = 7, - ): - - super().__init__(config) - self.logger.info( - f"Initializing LiberoRobotEnv (Stateless), task suite: {task_suite_name}" - ) - - self.model = self._register_model() - self.model_family = model_family - self.resolution = resolution - self.seed = seed - - self.logger.info("Importing Libero and related utils...") - self.RobotStateActionData = RobotStateActionData - - self.rollout_dir = os.path.join(rollout_dir, task_suite_name) - os.makedirs(self.rollout_dir, exist_ok=True) - - if save_rollout_video is not None: - self.save_rollout_video = save_rollout_video - else: - self.save_rollout_video = None - self.logger.warning("save_rollout_video not found, video saving disabled.") - - self.task_suite_name = task_suite_name - benchmark_dict = benchmark.get_benchmark_dict() - self.task_suite = benchmark_dict[self.task_suite_name]() - self.num_tasks = self.task_suite.n_tasks - - self.initial_states_path = initial_states_path - self.all_initial_states = None - if self.initial_states_path != "DEFAULT": - try: - with open(self.initial_states_path, "r") as f: - self.all_initial_states = json.load(f) - self.logger.info( - f"Loaded custom initial states from {self.initial_states_path}" - ) - except Exception as e: - self.logger.error(f"Failed to load initial states file: {e}") - raise - - def _register_model(self) -> WallXPolicy: - - return WallXPolicy( - model_path=self.config.model_path, - train_config=self.config.train_config, - action_tokenizer_path=self.config.action_tokenizer_path, - action_dim=self.config.action_dim, - agent_pos_dim=self.config.action_dim, - pred_horizon=self.config.pred_horizon, - camera_key=self.config.cam_names, - predict_mode=self.config.predict_mode, - ) - - def get_instruction(self, task_desc: str) -> str: - return task_desc - - def get_observation(self, raw_obs: Dict[str, Any]) -> Dict[str, Any]: - if raw_obs is None: - raise ValueError("Raw observation is None") - - data_obj = self.RobotStateActionData(config=self.config) - - pos = raw_obs["robot0_eef_pos"] - rot = quat2axisangle(raw_obs["robot0_eef_quat"]) - grip = raw_obs["robot0_gripper_qpos"][0:1] - - data_obj.save_state_data_with_key(pos[None], "follow_right_ee_cartesian_pos") - data_obj.save_state_data_with_key(rot[None], "follow_right_ee_rotation") - data_obj.save_state_data_with_key(grip[None], "follow_right_gripper") - data_obj.dof_mask = self._get_dof_mask() - - face_view = get_libero_image(raw_obs) - right_wrist_view = get_libero_wrist_image(raw_obs) - - return { - "robot_state_action_data": data_obj, - "face_view": face_view, - "right_wrist_view": right_wrist_view, - } - - def apply_action( - self, input_data: Dict[str, Any], env: Any = None, replay_images: list = None - ) -> bool: - if env is None: - raise ValueError( - "In Stateless mode, apply_action must be called with explicit 'env' parameter" - ) - - action_data = input_data["robot_state_action_data"] - right_arm_traj = self._get_right_arm_action(action_data) - while ( - right_arm_traj is not None - and right_arm_traj.ndim > 2 - and right_arm_traj.shape[0] == 1 - ): - right_arm_traj = right_arm_traj.squeeze(0) - - done = False - t = 0 - - try: - for i in range(len(right_arm_traj)): - if done: - break - - action_7d = right_arm_traj[i] - obs, reward, done, info = env.step(action_7d) - t += 1 - - if obs is not None and replay_images is not None: - replay_images.append(get_libero_image(obs)) - - input_data["_last_obs"] = obs - - except Exception as e: - self.logger.error(f"Env step error: {e}") - return False # Error treated as failure - - return done, t - - def apply_action_batch( - self, - vec_env: Any, - trajectories: List[np.ndarray], - active_indices: List[int], - status_list: List[Dict[str, Any]], - model_outputs: List[Dict[str, Any]], - ) -> None: - """ - Execute action trajectories in parallel batch. - - Uses SubprocVectorEnv to execute actions in parallel for all active environments. - Integrates vec_env.step(batch_actions, id=still_active) in this function. - """ - if not trajectories: - return - - max_traj_len = max(len(traj) for traj in trajectories) - if max_traj_len == 0: - return - - for step_idx in range(max_traj_len): - # Check if there are still active environments - still_active = [ - idx - for idx in active_indices - if (not status_list[idx]["done"]) - and status_list[idx]["count"] > 0 - and step_idx < len(trajectories[active_indices.index(idx)]) - ] - if not still_active: - break - - # Build batch actions (only includes actions for still_active environments) - batch_actions = [] - for idx in still_active: - traj_idx = active_indices.index(idx) - action_7d = trajectories[traj_idx][step_idx] - # Ensure action_7d is numpy array or list - if isinstance(action_7d, np.ndarray): - batch_actions.append(action_7d) - else: - batch_actions.append(np.array(action_7d)) - - # Convert to numpy array with shape (batch_size, action_dim) - batch_actions = np.array(batch_actions) - - # Execute step in parallel (only for still_active environments) - obs_list, reward_list, done_list, info_list = vec_env.step( - batch_actions, id=still_active - ) - - # Process returned results - if obs_list.dtype == object: - obs_list = list(obs_list) - else: - obs_list = [obs_list[i] for i in range(len(obs_list))] - done_list = [bool(done_list[i]) for i in range(len(done_list))] - - # Update each environment's status - for i, idx in enumerate(still_active): - st = status_list[idx] - obs = obs_list[i] - done = done_list[i] - - if obs is not None: - st["current_obs"] = obs - if st["replay_images"] is not None: - st["replay_images"].append(get_libero_image(obs)) - - st["count"] -= 1 - st["success"] = done - st["done"] = done or st["count"] <= 0 - - # Update model_output's _last_obs - model_outputs[active_indices.index(idx)]["_last_obs"] = obs - - def get_task_info(self, task_id: int) -> Tuple[str, Any]: - """ - Get task information (task description and initial states) without creating environment. - - Returns: - Tuple[str, Any]: (task_desc, default_initial_states) - """ - if task_id < 0 or task_id >= self.num_tasks: - raise ValueError(f"Invalid task ID: {task_id}") - - task = self.task_suite.get_task(task_id) - task_desc = task.language - default_initial_states = self.task_suite.get_task_init_states(task_id) - - return task_desc, default_initial_states - - def create_env_for_task(self, task_id: int) -> Tuple[Any, str, Any]: - if task_id < 0 or task_id >= self.num_tasks: - raise ValueError(f"Invalid task ID: {task_id}") - - task = self.task_suite.get_task(task_id) - default_initial_states = self.task_suite.get_task_init_states(task_id) - - env, task_desc = get_libero_env( - task, - model_family=self.model_family, - resolution=self.resolution, - seed=self.seed, - ) - - env.env = VisualizationWrapper(env.env) - env.env.set_visualization_setting(setting="grippers", visible=False) - - return env, task_desc, default_initial_states - - def _get_initial_state_for_episode( - self, task_desc: str, default_states: Any, episode_idx: int - ) -> np.ndarray: - if self.initial_states_path == "DEFAULT": - if default_states is None: - raise ValueError("Default states missing") - return default_states[episode_idx] - else: - if self.all_initial_states is None: - raise ValueError("Custom states not loaded") - initial_states_task_key = task_desc.replace(" ", "_") - episode_key = f"demo_{episode_idx}" - if not self.all_initial_states[initial_states_task_key][episode_key][ - "success" - ]: - raise ValueError(f"Expert demo failed for {episode_key}") - return np.array( - self.all_initial_states[initial_states_task_key][episode_key][ - "initial_state" - ] - ) - - def reset_env( - self, env: Any, task_desc: str, default_states: Any, episode_idx: int - ) -> Any: - try: - if episode_idx >= 0: - state = self._get_initial_state_for_episode( - task_desc, default_states, episode_idx - ) - obs = env.set_init_state(state) - return obs - else: - return env.reset() - except Exception as e: - self.logger.error(f"Reset failed: {e}, falling back to default reset") - return env.reset() - - def _get_dof_mask(self): - dof_config = self.config.train_config["dof_config"] - total_dof = sum(dof_config.values()) - dof_mask = np.ones((1, self.config.action_horizon, total_dof)) - mask_keys = [ - "follow_left_ee_cartesian_pos", - "follow_left_ee_rotation", - "follow_left_gripper", - "head_actions", - "height", - "velocity_decomposed", - ] - start_idx = 0 - for key, dof_size in dof_config.items(): - if key in mask_keys: - dof_mask[:, :, start_idx : start_idx + dof_size] = 0 - start_idx += dof_size - return dof_mask - - def _get_right_arm_action( - self, robot_state_action_data: RobotStateActionData - ) -> np.ndarray: - right_ee_cartesian_pos = robot_state_action_data.data[ - "action_right_ee_cartesian_pos" - ] - right_ee_rotation = robot_state_action_data.data["action_right_ee_rotation"] - right_gripper = robot_state_action_data.data["action_right_gripper"] - return np.concatenate( - [right_ee_cartesian_pos, right_ee_rotation, right_gripper], axis=1 - ) - - def _get_left_arm_action( - self, robot_state_action_data: RobotStateActionData - ) -> np.ndarray: - left_ee_cartesian_pos = robot_state_action_data.data[ - "action_left_ee_cartesian_pos" - ] - left_ee_rotation = robot_state_action_data.data["action_left_ee_rotation"] - left_gripper = robot_state_action_data.data["action_left_gripper"] - return np.concatenate( - [left_ee_cartesian_pos, left_ee_rotation, left_gripper], axis=1 - ) - - def _save_rollout( - self, - replay_images: List[np.ndarray], - success: bool, - task_id: int, - task_desc: str, - episode_idx: int, - ): - if not self.save_rollout_video or not replay_images: - return - try: - task_name_safe = task_desc.replace(" ", "_").replace(".", "") - filename = f"{episode_idx}{'_SUCCESS' if success else '_FAILURE'}--_{task_name_safe}.mp4" - self.save_rollout_video( - self.rollout_dir, - replay_images, - filename, - success=success, - task_description=task_desc, - log_file=None, - model_family=self.model_family, - ) - self.logger.info(f"Saved video: {filename}") - except Exception as e: - self.logger.error(f"Save video failed: {e}") - - def run_infer_flow_action( - self, - env: Any, - task_id: int, - task_desc: str, - default_initial_states: Any, - episode_idx: int, - max_infer_times: int = 5, - num_steps_wait: int = 10, - ) -> bool: - replay_images = [] - num_steps = 0 - done = False - count = max_infer_times - - current_obs = self.reset_env( - env, task_desc, default_initial_states, episode_idx - ) - if current_obs is None: - return False - - while num_steps < num_steps_wait: - obs, reward, done, info = env.step( - get_libero_dummy_action(self.model_family) - ) - num_steps += 1 - if obs is not None: - current_obs = obs - - while not done and count > 0: - try: - model_input = self.get_observation(current_obs) - instruction = self.get_instruction(task_desc) - model_input["prompt"] = instruction - model_input["dataset_names"] = "libero_all" - - state = np.concatenate( - [ - model_input["robot_state_action_data"].data[ - "state_right_ee_cartesian_pos" - ], - model_input["robot_state_action_data"].data[ - "state_right_ee_rotation" - ], - model_input["robot_state_action_data"].data[ - "state_right_gripper" - ], - ], - axis=-1, - ) - - model_input["state"] = state - model_output = self.model.infer(model_input) - - model_output["robot_state_action_data"] = model_input[ - "robot_state_action_data" - ] - model_output["robot_state_action_data"].save_action_data( - model_output["predict_action"] - ) - - model_output["_last_obs"] = None - - done, delta_t = self.apply_action( - model_output, env=env, replay_images=replay_images - ) - - if model_output.get("_last_obs") is not None: - current_obs = model_output["_last_obs"] - - count -= delta_t - - except Exception as e: - self.logger.error(f"Episode Error: {e}") - import traceback - - traceback.print_exc() - break - - success = done - if count <= 0 and not done: - self.logger.warning( - f"Timeout: reached {max_infer_times} steps without success." - ) - success = False - - self._save_rollout(replay_images, success, task_id, task_desc, episode_idx) - return success - - def run_infer_flow_action_batch( - self, - vec_env: Any, - task_ids: List[int] = None, - task_descs: List[str] = None, - default_initial_states_list: List[Any] = None, - episode_indices: List[int] = None, - max_infer_times: int = 5, - num_steps_wait: int = 10, - ) -> List[bool]: - """ - Support batch inference: model inference in parallel (batch), environment execution in parallel (SubprocVectorEnv). - - Uses SubprocVectorEnv to run environments in subprocess during multi-batch inference, all environments execute actions in parallel. - - Returns a list of success flags for each sample. - """ - if vec_env is None: - raise ValueError("vec_env must be specified") - batch_size = len(vec_env) - if task_ids is not None: - assert len(task_ids) == batch_size, "task_ids length must match envs" - if episode_indices is not None: - assert ( - len(episode_indices) == batch_size - ), "episode_indices length must match envs" - - status_list = [] - for i in range(batch_size): - status_list.append( - { - "vec_env": vec_env, - "env_id": i, # Index in vec_env - "task_desc": task_descs[i], - "replay_images": [], - "num_steps": 0, - "done": False, # Whether episode has ended - "success": False, # Whether successfully completed - "count": max_infer_times, - "current_obs": None, - "default_states": None, - } - ) - - # Initialize/reset: Use SubprocVectorEnv to batch set initial states - init_states_to_set = [] - for i in range(batch_size): - task_desc = task_descs[i] - default_states = default_initial_states_list[i] - status_list[i]["default_states"] = default_states - ep_i = episode_indices[i] - init_state = self._get_initial_state_for_episode( - task_desc, default_states, ep_i - ) - init_states_to_set.append(init_state) - - # Batch set initial states - try: - obs_list = vec_env.set_init_state(init_states_to_set) - if obs_list.dtype == object: - obs_list = list(obs_list) - else: - obs_list = [obs_list[i] for i in range(len(obs_list))] - - for i, obs in enumerate(obs_list): - if obs is None: - raise ValueError( - f"Reset environment returned None, task_id: {task_ids[i]}, episode_idx: {episode_indices[i]}" - ) - status_list[i]["current_obs"] = obs - status_list[i]["done"] = False - status_list[i]["success"] = False - status_list[i]["count"] = max_infer_times - except Exception as e: - self.logger.error(f"Failed to batch set initial states: {e}") - raise - - # Warmup steps (batch execution) - dummy_action = get_libero_dummy_action(self.model_family) - dummy_actions = np.array([dummy_action] * batch_size) - for _ in range(num_steps_wait): - obs_list, _, done_list, _ = vec_env.step(dummy_actions) - # Update current_obs - if obs_list.dtype == object: - obs_list = list(obs_list) - else: - obs_list = [obs_list[i] for i in range(len(obs_list))] - for i, obs in enumerate(obs_list): - if obs is not None: - status_list[i]["current_obs"] = obs - - # Main loop: model parallel inference, environment parallel execution (SubprocVectorEnv) - while any((not st["done"]) and st["count"] > 0 for st in status_list): - active_indices = [ - idx - for idx, st in enumerate(status_list) - if (not st["done"]) and st["count"] > 0 - ] - if not active_indices: - break - print(f"Batch infer loop, active indices: {active_indices}") - - observations = [] - instructions = [] - for idx in active_indices: - st = status_list[idx] - observations.append(self.get_observation(st["current_obs"])) - instructions.append(self.get_instruction(st["task_desc"])) - - # Model batch inference - model_outputs = self.model.infer_flow_action_batch( - observations, instructions - ) - # Extract action trajectories for all active environments - trajectories = [] - for out in model_outputs: - action_data = out["robot_state_action_data"] - right_arm_traj = self._get_right_arm_action(action_data) - while ( - right_arm_traj is not None - and right_arm_traj.ndim > 2 - and right_arm_traj.shape[0] == 1 - ): - right_arm_traj = right_arm_traj.squeeze(0) - if right_arm_traj is None or len(right_arm_traj) == 0: - # If trajectory is empty, create an empty trajectory - right_arm_traj = np.array([]).reshape(0, 7) - trajectories.append(right_arm_traj) - - # Use apply_action_batch to execute action trajectories in parallel - try: - self.apply_action_batch( - vec_env=vec_env, - trajectories=trajectories, - active_indices=active_indices, - status_list=status_list, - model_outputs=model_outputs, - ) - except Exception as e: - self.logger.error(f"Batch parallel action error: {e}") - # Mark all active environments as failed - for idx in active_indices: - status_list[idx]["done"] = True - status_list[idx]["success"] = False - - # Save replay and results - success_list = [] - for i, st in enumerate(status_list): - success = st.get("success", False) - if st["count"] <= 0 and not st["success"]: - self.logger.warning( - f"Batch timeout: reached {max_infer_times} steps without success (idx {i})." - ) - st["replay_images"] = st.get("replay_images", []) - tid_i = task_ids[i] - epi_i = episode_indices[i] - self._save_rollout( - st["replay_images"], - success, - tid_i, - st["task_desc"], - epi_i, - ) - success_list.append(success) - - return success_list - - def run_infer_ar_action( - self, - env: Any, - task_id: int, - task_desc: str, - default_initial_states: Any, - episode_idx: int, - max_infer_times: int = 10, - num_steps_wait: int = 10, - ) -> bool: - replay_images = [] - num_steps = 0 - done = False - count = max_infer_times - - current_obs = self.reset_env( - env, task_desc, default_initial_states, episode_idx - ) - if current_obs is None: - self.logger.error("Environment reset returned None.") - return False - - while num_steps < num_steps_wait: - obs, reward, done, info = env.step( - get_libero_dummy_action(self.model_family) - ) - num_steps += 1 - if obs is not None: - current_obs = obs - - while not done and count > 0: - try: - model_input = self.get_observation(current_obs) - instruction = self.get_instruction(task_desc) - - model_output = self.model.infer_ar_action(model_input, instruction) - - model_output["_last_obs"] = None - - done, delta_t = self.apply_action( - model_output, env=env, replay_images=replay_images - ) - - if model_output.get("_last_obs") is not None: - current_obs = model_output["_last_obs"] - else: - if not done: - self.logger.warning( - "Did not receive new observation after apply_action, but episode is not done." - ) - - count -= delta_t - - except Exception as e: - self.logger.error( - f"AR Episode Run Error (Task {task_id}, Ep {episode_idx}): {e}" - ) - import traceback - - traceback.print_exc() - break - - success = done - if count <= 0 and not done: - self.logger.warning( - f"Timeout: AR policy reached {max_infer_times} steps without success." - ) - success = False - - self._save_rollout(replay_images, success, task_id, task_desc, episode_idx) - - return success diff --git a/wall_x/infer/infer_config.py b/wall_x/infer/infer_config.py deleted file mode 100644 index bf7dc823..00000000 --- a/wall_x/infer/infer_config.py +++ /dev/null @@ -1,587 +0,0 @@ -import yaml -import os -from wall_x.model.model_utils import update_model_config - -# from x2robot_dataset.configs.config import X2RDataConfig - -import json -from typing import List, Dict, Optional, Any -from dataclasses import dataclass, field - - -@dataclass -class X2RDataConfig: - """ - Unified X2Robot data configuration class (reorganized by README's 5 modules): - 1) Data I/O and caching - 2) Visual input and sampling (image/camera) - 3) Action and time series - 4) Instruction and multimodal - 5) Data cleaning and alignment (validation/augmentation/framework constraints) - """ - - # ---------------------------------------------------------------------- - # 1) Data I/O and caching - # ---------------------------------------------------------------------- - cache_dir: str = "~/.cache/dataset_cache" - dataset_config_path: Optional[str] = None - use_cache: bool = True - check_mode: bool = True - preload_size: int = 128 - buffer_size: int = 20000 - batch_size: int = 32 - train_test_split: float = 0.9 - seed: int = 42 - episode_chunk_size: int = ( - 500 # Commonly used on VG side (number of frames for episode chunking) - ) - - # ---------------------------------------------------------------------- - # 2) Visual input and sampling (image/camera) - # ---------------------------------------------------------------------- - # Camera mapping - cam_mapping: Dict[str, str] = field( - default_factory=lambda: { - "faceImg": "face_view", - "leftImg": "left_wrist_view", - "rightImg": "right_wrist_view", - } - ) - # Image and augmentation - resolution: Dict[str, int] = field( - default_factory=lambda: { - "face_view": -1, - "left_wrist_view": 128, - "right_wrist_view": 128, - } - ) - cam_augmentation_list: List[str] = field(default_factory=list) - - # Image time series (history/future) - image_horizon: int = 1 - image_history_length: int = 0 - image_history_interval: int = 1 - future_image_length: int = 0 - future_image_interval: int = 1 - future_image_indices: Optional[List[int]] = ( - None # If provided, length must equal image_horizon - ) - - # Smart scaling - max_pixels: int = field( - default_factory=lambda: 1280 * 28 * 28 - ) # Will be replaced with MAX_PIXELS in __post_init__ - min_pixels: int = field( - default_factory=lambda: 4 * 28 * 28 - ) # Will be replaced with MIN_PIXELS in __post_init__ - image_factor: int = 28 # Will be replaced with IMAGE_FACTOR in __post_init__ - - # ---------------------------------------------------------------------- - # 3) Action and time series - # ---------------------------------------------------------------------- - predict_action_keys: List[str] = field(default_factory=list) - obs_action_keys: List[str] = field(default_factory=list) - - # Action window - action_horizon: int = 21 - action_history_length: int = 0 - action_horizon_flow: int = 32 - action_horizon_ar: int = 0 - - # Padding strategy - left_padding: bool = True - right_padding: bool = True - - # Dimension configuration - dof_config: Dict[str, int] = field(default_factory=dict) # Input degrees of freedom - agent_pos_config: Dict[str, int] = field( - default_factory=dict - ) # Output degrees of freedom - - # State augmentation - state_augmentation_ratio: float = 1.0 # Ratio of augmented states - state_augmentation_prob: float = ( - 0.1 # Random dimension masking probability for state string - ) - state_drop_prob: float = 0.0 # Probability of dropping entire state - - # ---------------------------------------------------------------------- - # 4) Instruction and multimodal - # ---------------------------------------------------------------------- - default_instruction: str = "" - instruction_path: Optional[str] = None - instruction_key: Optional[List[Dict]] = None - - multimodal_chunk_size: int = 500 - generate_subtask_ratio: float = 0.0 - cot_ratio: float = 0.0 - multimodal_data_ratio: float = ( - 0.25 # Multimodal data ratio per batch in VLA dataset - ) - instruction_key_prob: Optional[Dict[str, float]] = None - trunc_action_with_instruction: bool = True - use_embodied_system_prompt_ratio: float = 0.0 - - # ---------------------------------------------------------------------- - # 5) Data cleaning and alignment (validation/augmentation/framework constraints) - # ---------------------------------------------------------------------- - filter_angle_outliers: bool = False - trim_stationary: bool = False - use_state_string_representation: bool = False - pad_prefix_to_same_length: bool = False - put_ar_predict_in_postfix: bool = ( - False # Whether to put ar prediction in postfix, set to True in prediction mode, False in training - ) - pad_to_128_multiple: bool = ( - False # Triton Attention requirement (deprecated, always set to False) - ) - max_seqlen: int = 768 - model_type: Optional[str] = None # qwen2_5, qwen2 - model_config_path: Optional[str] = ( - None # Model config path (used to derive PaddingSide) - ) - low_dim_obs_horizon: int = 1 # To be deprecated - - # ---------------------------------------------------------------------- - # Validation and post-processing - # ---------------------------------------------------------------------- - def __post_init__(self): - # TODO: Determine VGA model type validation here - # assert self.model_type in ["qwen2_5", "qwen3"], f"Unsupported model type: {self.model_type}" - - if self.model_type == "qwen2_5": - self.max_pixels = 16384 * 28 * 28 - self.min_pixels = 4 * 28 * 28 - self.image_factor = 28 - elif self.model_type == "qwen3": - self.max_pixels = 16384 * 32 * 32 - self.min_pixels = 4 * 32 * 32 - self.image_factor = 32 - - # Future image indices validation - if ( - self.future_image_indices - and len(self.future_image_indices) != self.image_horizon - ): - raise ValueError( - f"future_image_indices length must equal image_horizon: " - f"{len(self.future_image_indices)} != {self.image_horizon}" - ) - - # Auto-derive action window - if self.action_horizon == 0: - self.action_horizon = max(self.action_horizon_flow, self.action_horizon_ar) - - # Auto-derive action keys - if not self.obs_action_keys: - self.obs_action_keys = list(self.agent_pos_config.keys()) - if not self.predict_action_keys: - self.predict_action_keys = list(self.dof_config.keys()) - - # Derive PaddingSide - # @Ryan: Only FlashAttention can use RightPadding, other AttnImpl use LeftPadding - if self.model_config_path is not None: - with open(self.model_config_path, "r", encoding="utf-8") as f: - cfg = json.load(f) - - attn_impl = cfg["_attn_implementation"] - - if attn_impl == "flash_attention_2": - self.padding_side = "right" - else: - self.padding_side = "left" - - # Convenience properties - @property - def use_6D_rotation(self) -> bool: - """Whether to use 6D rotation (auto-determined from predict_action_keys)""" - if hasattr(self, "_use_6D_rotation"): - return self._use_6D_rotation - self._use_6D_rotation = any("6D" in key for key in self.predict_action_keys) - return self._use_6D_rotation - - @property - def use_relative_action(self) -> bool: - """Whether to use relative action (auto-determined from predict_action_keys)""" - if hasattr(self, "_use_relative_action"): - return self._use_relative_action - self._use_relative_action = any( - "relative" in key for key in self.predict_action_keys - ) - return self._use_relative_action - - # ---------------------------------------------------------------------- - # YAML initialization - # ---------------------------------------------------------------------- - @classmethod - def from_yaml_dict(cls, yaml_dict: Dict[str, Any]) -> "X2RDataConfig": - """ - Create config object from YAML config dict. Prioritizes data sub-config, then top-level fields. - """ - data_config = yaml_dict.get("data", {}) - params: Dict[str, Any] = {} - - # 1) Data I/O and caching - params.update( - { - "cache_dir": data_config.get( - "cache_dir", yaml_dict.get("cache_dir", "~/.cache/dataset_cache") - ), - "dataset_config_path": data_config.get( - "dataset_config_path", yaml_dict.get("dataset_config_path", None) - ), - "use_cache": data_config.get( - "use_cache", yaml_dict.get("use_cache", True) - ), - "check_mode": data_config.get( - "check_mode", yaml_dict.get("check_mode", True) - ), - "preload_size": data_config.get( - "preload_size", yaml_dict.get("preload_size", 128) - ), - "buffer_size": data_config.get( - "buffer_size", yaml_dict.get("buffer_size", 20000) - ), - "batch_size": data_config.get( - "batch_size", - yaml_dict.get( - "batch_size_per_gpu", yaml_dict.get("batch_size", 32) - ), - ), - "train_test_split": data_config.get("train_test_split", 0.9), - "seed": yaml_dict.get("seed", 42), - "episode_chunk_size": data_config.get("episode_chunk_size", 500), - } - ) - - # 2) Visual input and sampling (image/camera) - params.update( - { - "cam_mapping": data_config.get( - "cam_mapping", - { - "faceImg": "face_view", - "leftImg": "left_wrist_view", - "rightImg": "right_wrist_view", - }, - ), - "resolution": data_config.get( - "resolution", - {"face_view": -1, "left_wrist_view": 128, "right_wrist_view": 128}, - ), - "cam_augmentation_list": data_config.get("cam_augmentation_list", []), - "image_horizon": data_config.get("image_horizon", 1), - "image_history_length": data_config.get("image_history_length", 0), - "image_history_interval": data_config.get("image_history_interval", 1), - "future_image_length": data_config.get("future_image_length", 0), - "future_image_interval": data_config.get("future_image_interval", 1), - "future_image_indices": data_config.get("future_image_indices", None), - "max_pixels": data_config.get("max_pixels", 1280 * 28 * 28), - "min_pixels": data_config.get("min_pixels", 4 * 28 * 28), - "image_factor": data_config.get("image_factor", 28), - } - ) - - # 3) Action and time series - params.update( - { - "predict_action_keys": data_config.get("predict_action_keys", []), - "obs_action_keys": data_config.get("obs_action_keys", []), - "action_horizon": data_config.get("action_horizon", 0), - "action_history_length": data_config.get("action_history_length", 0), - "action_horizon_flow": data_config.get( - "action_horizon_flow", yaml_dict.get("action_horizon_flow", 32) - ), - "action_horizon_ar": data_config.get("action_horizon_ar", 0), - "left_padding": data_config.get("left_padding", True), - "right_padding": data_config.get("right_padding", True), - "dof_config": yaml_dict.get( - "dof_config", data_config.get("dof_config", {}) - ), - "agent_pos_config": yaml_dict.get( - "agent_pos_config", data_config.get("agent_pos_config", {}) - ), - "state_augmentation_prob": data_config.get( - "state_augmentation_prob", 0.05 - ), - "state_drop_prob": data_config.get("state_drop_prob", 0.0), - } - ) - - # 4) Instruction and multimodal - params.update( - { - "default_instruction": data_config.get("default_instruction", ""), - "instruction_path": data_config.get("instruction_path", None), - "instruction_key": data_config.get("instruction_key", None), - "multimodal_chunk_size": data_config.get("multimodal_chunk_size", 500), - "generate_subtask_ratio": data_config.get( - "generate_subtask_ratio", 0.0 - ), - "cot_ratio": data_config.get("cot_ratio", 0.0), - "multimodal_data_ratio": data_config.get("multimodal_data_ratio", 0.25), - "instruction_key_prob": data_config.get("instruction_key_prob", None), - "trunc_action_with_instruction": data_config.get( - "trunc_action_with_instruction", True - ), - "use_embodied_system_prompt_ratio": data_config.get( - "use_embodied_system_prompt_ratio", - yaml_dict.get("use_embodied_system_prompt_ratio", 0.0), - ), - } - ) - - # 5) Data cleaning and alignment (validation/augmentation/framework constraints) - params.update( - { - "filter_angle_outliers": data_config.get( - "filter_angle_outliers", False - ), - "trim_stationary": data_config.get("trim_stationary", False), - "use_state_string_representation": data_config.get( - "use_state_string_representation", - yaml_dict.get("use_state_string_representation", False), - ), - "pad_prefix_to_same_length": data_config.get( - "pad_prefix_to_same_length", False - ), - "put_ar_predict_in_postfix": data_config.get( - "put_ar_predict_in_postfix", False - ), - # "pad_to_128_multiple": data_config.get("pad_to_128_multiple", True), - "padding_side": data_config.get("padding_side", "left"), - "max_seqlen": yaml_dict.get("max_seqlen", 768), - "model_type": yaml_dict.get("model_type", "qwen2_5"), - "model_config_path": yaml_dict.get("qwen_vl_act_config_path", None), - "low_dim_obs_horizon": data_config.get("low_dim_obs_horizon", 1), - } - ) - - # Only keep valid fields defined in dataclass - valid_fields = {f.name for f in cls.__dataclass_fields__.values()} - filtered = {k: v for k, v in params.items() if k in valid_fields} - return cls(**filtered) - - # ---------------------------------------------------------------------- - # Dict-style access (for compatibility with existing calls) - # ---------------------------------------------------------------------- - def __getitem__(self, key: str): - try: - return getattr(self, key) - except AttributeError: - raise KeyError(f"'{key}' not found in {self.__class__.__name__}") - - def __setitem__(self, key: str, value): - setattr(self, key, value) - - def __contains__(self, key: str) -> bool: - return hasattr(self, key) - - def keys(self): - return self.__dict__.keys() - - def values(self): - return self.__dict__.values() - - def items(self): - return self.__dict__.items() - - -class InferConfig: - def __init__( - self, - checkpoint_path: str | None = None, - train_config_path: str | None = None, - robot_host: str = "0.0.0.0", - robot_port: int = 33723, - robot_id: str = "10053", - robot_type: str = "desktop", # ["desktop", "turtle"] - robot_action_start_ratio: float = 0, # Action execution start ratio - robot_action_end_ratio: float = 0.8, # Action execution end ratio - robot_action_interpolate_multiplier: int = 70, # Action interpolation - robot_use_joint_angle_control: bool = False, # Use joint control (model must be joint prediction model) - turtle_as_desktop: bool = False, # Use turtle body for desktop operation, fixed chassis head movement, head camera, and chassis height - action_horizon: int = 10, # Please correctly fill in the model's horizon - action_dim: int | None = None, - model_device: str = "cuda:0", - num_inference_timesteps: int = 10, - norm_key: str = "x2_normal", - cam_names: list[str] = ["face_view", "right_wrist_view"], - ): - # Private attribute for storing path - assert checkpoint_path is not None - self._checkpoint_path = checkpoint_path - if os.path.exists(os.path.join(checkpoint_path, "normalizer_action.pth")): - self.normalizer_action_path = os.path.join( - checkpoint_path, "normalizer_action.pth" - ) - if os.path.exists(os.path.join(checkpoint_path, "normalizer_propri.pth")): - self.normalizer_propri_path = os.path.join( - checkpoint_path, "normalizer_propri.pth" - ) - - self.model_path = checkpoint_path - self.action_tokenizer_path = "/x2robot_v2/Models/fast/" - - # Other configuration attributes - self.robot_host = robot_host - self.robot_port = robot_port - self.robot_type = robot_type # ["desktop", "turtle"] - self.robot_id = robot_id - self.robot_action_start_ratio = robot_action_start_ratio - self.robot_action_end_ratio = robot_action_end_ratio - self.robot_action_interpolate_multiplier = robot_action_interpolate_multiplier - self.robot_use_joint_angle_control = ( - robot_use_joint_angle_control # Use joint angle control - ) - self.turtle_as_desktop = turtle_as_desktop - - self._action_horizon = ( - action_horizon # Default controlled by train config's flow action horizon - ) - self._action_dim = action_dim # Default determined by train config's dof config - - self.action_dim = action_dim - self.pred_horizon = action_horizon - self.predict_mode = "diffusion" - self.camera_key = cam_names - - self.model_device = model_device - self.num_inference_timesteps = ( - num_inference_timesteps # flow matching related config - ) - - # Initialize config objects - self.train_config: dict = {} - self.model_config = None - self.data_config = None - self.norm_key = norm_key - self.cam_names = cam_names - # Load all configs - self._load_all_configs(train_config_path) - - @property - def checkpoint_path(self) -> str | None: - return self._checkpoint_path - - @checkpoint_path.setter - def checkpoint_path(self, value: str | None): - """When checkpoint_path is updated, reload all configs""" - if self._checkpoint_path != value: - self._checkpoint_path = value - self._load_all_configs() - - @property - def action_horizon(self) -> int: - return self._action_horizon - - @action_horizon.setter - def action_horizon(self, value: int): - self._action_horizon = value - - @property - def action_dim(self) -> int | None: - return self._action_dim - - @action_dim.setter - def action_dim(self, value: int | None): - self._action_dim = value - - def _load_all_configs(self, train_config_path=None): - """Unified entry point for loading all configs""" - self._load_train_config(train_config_path) - self._load_model_config() - self._load_data_config() - - # Update action_horizon and action_dim (if needed) - if self._action_horizon is None: - self._action_horizon = self.train_config.get("data", {}).get( - "action_horizon_flow", 32 - ) - assert self._action_horizon is not None and self._action_horizon > 0 - - if self._action_dim is None: - self._action_dim = sum(self.train_config.get("dof_config", {}).values()) - - def _load_train_config(self, train_config_path): - if train_config_path is None: - train_config_path = os.path.join(self._checkpoint_path, "config.yml") - with open(train_config_path, "r") as f: - self.train_config = yaml.load(f, Loader=yaml.FullLoader) - - ckpt_dir = self._checkpoint_path - preprocessor_file = os.path.join(ckpt_dir, "preprocessor_config.json") - if os.path.exists(preprocessor_file): - print(f"[LoadConfig] Found {preprocessor_file}, override processor_path.") - self.train_config["processor_path"] = ckpt_dir - - tokenizer_file = os.path.join(ckpt_dir, "tokenizer.json") - tokenizer_config_file = os.path.join(ckpt_dir, "tokenizer_config.json") - if "action_tokenizer_path" in self.train_config and not os.path.exists( - self.train_config["action_tokenizer_path"] - ): - if os.path.exists(tokenizer_file) and os.path.exists(tokenizer_config_file): - print( - f"[LoadConfig] Found tokenizer files in {ckpt_dir}, override action_tokenizer_path." - ) - self.train_config["action_tokenizer_path"] = ckpt_dir - else: - print("[LoadConfig] Cannot load action tokenizer! ") - - def _load_model_config(self): - ckpt_config_path = os.path.join(self._checkpoint_path, "config.json") - resolved_cfg_path = None - - if os.path.exists(ckpt_config_path): - # Prefer checkpoint config - resolved_cfg_path = ckpt_config_path - print(f"[LoadModelConfig] Using checkpoint config.json: {ckpt_config_path}") - else: - # Fallback to original config path - fallback_cfg = self.train_config.get("qwen_vl_act_config_path", None) - if fallback_cfg is not None: - resolved_cfg_path = fallback_cfg - print(f"[LoadModelConfig] Using fallback act config: {fallback_cfg}") - - if resolved_cfg_path is None or (not os.path.exists(resolved_cfg_path)): - raise ValueError( - f"[LoadModelConfig] Cannot load model config! " - f"Checked:\n" - f" - Checkpoint config.json: {ckpt_config_path}\n" - f" - Fallback path: {self.train_config.get('qwen_vl_act_config_path', None)}" - ) - - # Save back to config for consistency - self.train_config["qwen_vl_act_config_path"] = resolved_cfg_path - - model_type = self.train_config["model_type"] - if model_type == "qwen2_5": - from wall_x.model.qwen2_5_based import Qwen2_5_VLConfig - - ConfigClass = Qwen2_5_VLConfig - - # elif model_type == "qwen3": - # from wall_x.model.qwen3_based import Qwen3VLConfig - - # ConfigClass = Qwen3VLConfig - - else: - raise ValueError(f"[LoadModelConfig] Unsupported model type: {model_type}") - - print(f"[LoadModelConfig] Loading model config from: {resolved_cfg_path}") - self.model_config = ConfigClass.from_pretrained(resolved_cfg_path) - - self.model_config = update_model_config(self.train_config, self.model_config) - - self.model_config._attn_implementation = "sdpa" - self.model_config.vision_config._attn_implementation = "flash_attention_2" - - print("[LoadModelConfig] Model config loaded and updated successfully.") - - def _load_data_config(self): - self.data_config = X2RDataConfig.from_yaml_dict(self.train_config) - - -if __name__ == "__main__": - config = InferConfig() - print(config.train_config) diff --git a/wall_x/infer/utils_libero.py b/wall_x/infer/utils_libero.py deleted file mode 100644 index b6fd5cb3..00000000 --- a/wall_x/infer/utils_libero.py +++ /dev/null @@ -1,323 +0,0 @@ -"""Utils for evaluating policies in LIBERO simulation environments.""" - -import math -import os -from enum import Enum -import imageio -import numpy as np -import matplotlib.pyplot as plt -import torch -from transformers import BatchFeature -import random -import time - -from libero.libero import get_libero_path -from libero.libero.envs import OffScreenRenderEnv - - -# Define task suite constants -class TaskSuite(str, Enum): - LIBERO_SPATIAL = "libero_spatial" - LIBERO_OBJECT = "libero_object" - LIBERO_GOAL = "libero_goal" - LIBERO_10 = "libero_10" - LIBERO_90 = "libero_90" - - -# Define max steps for each task suite -TASK_MAX_STEPS = { - TaskSuite.LIBERO_SPATIAL: 220, # longest training demo has 193 steps - TaskSuite.LIBERO_OBJECT: 280, # longest training demo has 254 steps - TaskSuite.LIBERO_GOAL: 300, # longest training demo has 270 steps - TaskSuite.LIBERO_10: 520, # longest training demo has 505 steps - TaskSuite.LIBERO_90: 400, # longest training demo has 373 steps -} - - -# Initialize important constants -ACTION_DIM = 7 -DATE = time.strftime("%Y_%m_%d") -DATE_TIME = time.strftime("%Y_%m_%d-%H_%M_%S") - -# Configure NumPy print settings -np.set_printoptions(formatter={"float": lambda x: "{0:0.3f}".format(x)}) - - -def set_seed_everywhere(seed: int) -> None: - """ - Set random seed for all random number generators for reproducibility. - - Args: - seed: The random seed to use - """ - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - np.random.seed(seed) - random.seed(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False - os.environ["PYTHONHASHSEED"] = str(seed) - - -def normalize_gripper_action(action: np.ndarray, binarize: bool = True) -> np.ndarray: - """ - Normalize gripper action from [0,1] to [-1,+1] range. - - This is necessary for some environments because the dataset wrapper - standardizes gripper actions to [0,1]. Note that unlike the other action - dimensions, the gripper action is not normalized to [-1,+1] by default. - - Normalization formula: y = 2 * (x - orig_low) / (orig_high - orig_low) - 1 - - Args: - action: Action array with gripper action in the last dimension - binarize: Whether to binarize gripper action to -1 or +1 - - Returns: - np.ndarray: Action array with normalized gripper action - """ - # Create a copy to avoid modifying the original - normalized_action = action.copy() - - # Normalize the last action dimension to [-1,+1] - orig_low, orig_high = 0.0, 1.0 - normalized_action[..., -1] = ( - 2 * (normalized_action[..., -1] - orig_low) / (orig_high - orig_low) - 1 - ) - - if binarize: - # Binarize to -1 or +1 - normalized_action[..., -1] = np.sign(normalized_action[..., -1]) - - return normalized_action - - -def invert_gripper_action(action: np.ndarray) -> np.ndarray: - """ - Flip the sign of the gripper action (last dimension of action vector). - - This is necessary for environments where -1 = open, +1 = close, since - the RLDS dataloader aligns gripper actions such that 0 = close, 1 = open. - - Args: - action: Action array with gripper action in the last dimension - - Returns: - np.ndarray: Action array with inverted gripper action - """ - # Create a copy to avoid modifying the original - inverted_action = action.copy() - - # Invert the gripper action - inverted_action[..., -1] *= -1.0 - - return inverted_action - - -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 - - -def get_libero_env(task, model_family, resolution=256, seed=7): - """Initializes and returns the LIBERO environment, along with the task description.""" - task_description = task.language - task_bddl_file = os.path.join( - get_libero_path("bddl_files"), task.problem_folder, task.bddl_file - ) - env_args = { - "bddl_file_name": task_bddl_file, - "camera_heights": resolution, - "camera_widths": resolution, - } - env = OffScreenRenderEnv(**env_args) - env.seed( - seed - ) # IMPORTANT: seed seems to affect object positions even when using fixed initial state - return env, task_description - - -def get_libero_dummy_action(model_family: str): - """Get dummy/no-op action, used to roll out the simulation while the robot does nothing.""" - return [0, 0, 0, 0, 0, 0, -1] - - -def get_libero_image(obs): - """Extracts third-person image from observations and preprocesses it.""" - img = obs["agentview_image"] - img = img[::-1, ::-1] # IMPORTANT: rotate 180 degrees to match train preprocessing - return img - - -def get_libero_wrist_image(obs): - """Extracts wrist camera image from observations and preprocesses it.""" - img = obs["robot0_eye_in_hand_image"] - img = img[::-1, ::-1] # IMPORTANT: rotate 180 degrees to match train preprocessing - return img - - -def save_rollout_video( - rollout_dir, - rollout_images, - idx, - success, - task_description, - log_file=None, - model_family="openvla_oft", -): - """Saves an MP4 replay of an episode.""" - processed_task_description = ( - task_description.lower() - .replace(" ", "_") - .replace("\n", "_") - .replace(".", "_")[:50] - ) - mp4_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}.mp4" - video_writer = imageio.get_writer(mp4_path, fps=30) - for img in rollout_images: - video_writer.append_data(img) - video_writer.close() - print(f"Saved rollout MP4 at path {mp4_path}") - if log_file is not None: - log_file.write(f"Saved rollout MP4 at path {mp4_path}\n") - return mp4_path - - -def save_rollout_data( - rollout_dir, - rollout_data, - idx, - success, - task_description, - log_file=None, - model_family="openvla_oft", -): - """ - Saves an NPY file of the rollout data. - """ - - # Process task description to make it suitable for filename - processed_task_description = ( - task_description.lower() - .replace(" ", "_") - .replace("\n", "_") - .replace(".", "_")[:50] - ) - - # Build .npy file path - npy_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--action.npy" - - # Save rollout_data as .npy file - np.save(npy_path, rollout_data) - print(f"Saved rollout data at path {npy_path}") - - fig, axes = plt.subplots(nrows=1, ncols=7, figsize=(20, 3)) - - titles = ["x", "y", "z", "roll", "pitch", "yaw", "grasp"] - - for i in range(rollout_data.shape[1]): - ax = axes[i] # Select the i-th subplot - ax.plot(rollout_data[:, i], label=f"Feature {i+1}") # Plot line chart - ax.set_title(titles[i]) # Set subplot title - ax.set_xlabel("Time in one episode") # Set x-axis label - - axes[-1].legend(["predicted action"], loc="upper right") - - plt.tight_layout() - png_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--action.png" - plt.savefig(png_path, dpi=300) # Save image as PNG file - - # If log file is provided, record the save path - if log_file is not None: - log_file.write(f"Saved rollout data at path {npy_path}\n") - - return npy_path - - -def save_rollout_observation( - rollout_dir, - rollout_data, - idx, - success, - task_description, - log_file=None, - model_family="openvla_oft", -): - """ - Saves an NPY file of the rollout data. - """ - - # Process task description to make it suitable for filename - processed_task_description = ( - task_description.lower() - .replace(" ", "_") - .replace("\n", "_") - .replace(".", "_")[:50] - ) - - # Build .npy file path - npy_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--observation.npy" - - # Save rollout_data as .npy file - np.save(npy_path, rollout_data) - print(f"Saved rollout data at path {npy_path}") - - if rollout_data.shape[1] == 7: - fig, axes = plt.subplots(nrows=1, ncols=7, figsize=(20, 3)) - titles = ["x", "y", "z", "roll", "pitch", "yaw", "grasp"] - else: - fig, axes = plt.subplots(nrows=1, ncols=8, figsize=(20, 3)) - titles = ["x", "y", "z", "roll", "pitch", "yaw", "-", "grasp"] - - for i in range(rollout_data.shape[1]): - ax = axes[i] # Select the i-th subplot - ax.plot(rollout_data[:, i], label=f"Feature {i+1}") # Plot line chart - ax.set_title(titles[i]) # Set subplot title - ax.set_xlabel("Time in one episode") # Set x-axis label - - axes[-1].legend(["predicted action"], loc="upper right") - - plt.tight_layout() - png_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--observation.png" - plt.savefig(png_path, dpi=300) # Save image as PNG file - - # If log file is provided, record the save path - if log_file is not None: - log_file.write(f"Saved rollout data at path {npy_path}\n") - - return npy_path - - -def quat2axisangle(quat): - """ - Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55 - - Converts quaternion to axis-angle format. - Returns a unit vector direction scaled by its angle in radians. - - Args: - quat (np.array): (x,y,z,w) vec4 float angles - - Returns: - np.array: (ax,ay,az) axis-angle exponential coordinates - """ - # clip quaternion - if quat[3] > 1.0: - quat[3] = 1.0 - elif quat[3] < -1.0: - quat[3] = -1.0 - - den = np.sqrt(1.0 - quat[3] * quat[3]) - if math.isclose(den, 0.0): - # This is (close to) a zero degree rotation, immediately return - return np.zeros(3) - - return (quat[:3] * 2.0 * math.acos(quat[3])) / den diff --git a/wall_x/model/action_head.py b/wall_x/model/action_head.py deleted file mode 100644 index 2f3ab733..00000000 --- a/wall_x/model/action_head.py +++ /dev/null @@ -1,808 +0,0 @@ -import torch - -import torch.nn as nn - -from typing import Union -import math - -from diffusers.schedulers.scheduling_ddpm import DDPMScheduler -from torch.distributions import Beta - - -def print_rank_last(message): - """If distributed is initialized, print only on last rank.""" - if torch.distributed.is_initialized(): - if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1): - print(message, flush=True) - else: - print(message, flush=True) - - -class Normalizer(nn.Module): - @classmethod - def from_ckpt(cls, ckpt_path): - instance = cls.__new__(cls) - nn.Module.__init__(instance) - - instance.min = nn.ParameterDict() - instance.delta = nn.ParameterDict() - instance.min_key = "min" - instance.delta_key = "delta" - - ckpt = torch.load(ckpt_path, map_location="cpu") - - for key, value in ckpt.items(): - # Parse key: "min.robot_name" -> prefix="min", name="robot_name" - try: - prefix, name = key.split(".", 1) - if hasattr(instance, prefix): - getattr(instance, prefix)[name] = nn.Parameter( - value, requires_grad=False - ) - print("prefix", prefix) - print("name", name) - except ValueError: - continue - - return instance - - def __init__( - self, action_statistic_dof, dof_config, min_key="min", delta_key="delta" - ): - super(Normalizer, self).__init__() - - self.min_key = min_key - self.delta_key = delta_key - - action_statistic = {} - for robot_name in action_statistic_dof.keys(): - action_statistic[robot_name] = {} - all_dof_min = [] - all_dof_delta = [] - for k in dof_config: - if k in action_statistic_dof[robot_name]: - if ( - min_key in action_statistic_dof[robot_name][k] - and delta_key in action_statistic_dof[robot_name][k] - ): - all_dof_min.extend(action_statistic_dof[robot_name][k][min_key]) - all_dof_delta.extend( - action_statistic_dof[robot_name][k][delta_key] - ) - else: - if robot_name == "x2_normal" or "libero" in robot_name: - print_rank_last( - f"Normalizer (Warning): min_key {min_key} or delta_key {delta_key} " - ) - print_rank_last( - f"not in action_statistic_dof[{robot_name}][{k}], use default min 0.0 and delta 1.0" - ) - all_dof_min.extend([0.0] * dof_config[k]) - all_dof_delta.extend([1.0] * dof_config[k]) - else: - if robot_name == "x2_normal" or "libero" in robot_name: - print_rank_last( - f"Normalizer (Warning): Action {k} not in action_statistic_dof for {robot_name}, use default min 0.0 and delta 1.0" - ) - all_dof_min.extend([0.0] * dof_config[k]) - all_dof_delta.extend([1.0] * dof_config[k]) - all_dof_min = torch.tensor(all_dof_min) - all_dof_delta = torch.tensor(all_dof_delta) - action_statistic[robot_name][min_key] = all_dof_min - action_statistic[robot_name][delta_key] = all_dof_delta - - self.min = nn.ParameterDict( - { - k: nn.Parameter(action_statistic[k][min_key], requires_grad=False) - for k in action_statistic.keys() - } - ) - self.delta = nn.ParameterDict( - { - k: nn.Parameter(action_statistic[k][delta_key], requires_grad=False) - for k in action_statistic.keys() - } - ) - - for k, v in action_statistic.items(): - print_rank_last( - f"Normalizer: {k} min {action_statistic[k][min_key]} delta {action_statistic[k][delta_key]}" - ) - - def normalize_data(self, xs, dataset_names): - new_xs = [] - dataset_names = [name for name in dataset_names if name != "x2_multimodal"] - for x, dataset_name in zip(xs, dataset_names): - x = (x - self.min[dataset_name]) / (self.delta[dataset_name]) - x = x * 2 - 1 - x = torch.clamp(x, -1, 1) - new_xs.append(x) - new_xs = torch.stack(new_xs) - return new_xs - - def unnormalize_data(self, xs, dataset_names, dof_mask=None): - new_xs = [] - dataset_names = [name for name in dataset_names if name != "x2_multimodal"] - dof_mask = dof_mask if dof_mask is not None else [None] * len(xs) - for x, dataset_name, mask in zip(xs, dataset_names, dof_mask): - x = (x + 1) / 2 - if mask is not None: - mask = mask[0].bool() - action_space_delta = self.delta[dataset_name][mask] - action_space_min = self.min[dataset_name][mask] - else: - action_space_delta = self.delta[dataset_name] - action_space_min = self.min[dataset_name] - x = x * action_space_delta + action_space_min - new_xs.append(x) - new_xs = torch.stack(new_xs) - return new_xs - - -class SinusoidalPosEmb(nn.Module): - def __init__(self, dim: int, min_period: float = 4e-3, max_period: float = 4.0): - super().__init__() - self.dim = dim - if dim % 2 != 0: - raise ValueError(f"embedding_dim ({dim}) must be divisible by 2") - self.min_period = min_period - self.max_period = max_period - - def forward(self, x): - device = x.device - half_dim = self.dim // 2 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp( - torch.arange(half_dim, device=device, dtype=torch.float32) * -emb - ) - emb = x[:, None] * emb[None, :] - emb = torch.cat((emb.sin(), emb.cos()), dim=-1) - return emb - - -class Downsample1d(nn.Module): - def __init__(self, dim): - super().__init__() - self.conv = nn.Conv1d(dim, dim, 3, 2, 1) - - def forward(self, x): - return self.conv(x) - - -class Upsample1d(nn.Module): - def __init__(self, dim): - super().__init__() - self.conv = nn.ConvTranspose1d(dim, dim, 4, 2, 1) - - def forward(self, x): - return self.conv(x) - - -class Conv1dBlock(nn.Module): - """ - Conv1d --> GroupNorm --> Mish - """ - - def __init__(self, inp_channels, out_channels, kernel_size, n_groups=8): - super().__init__() - - self.block = nn.Sequential( - nn.Conv1d( - inp_channels, out_channels, kernel_size, padding=kernel_size // 2 - ), - nn.GroupNorm(n_groups, out_channels), - nn.Mish(), - ) - - def forward(self, x): - return self.block(x) - - -class ConditionalResidualBlock1D(nn.Module): - def __init__(self, in_channels, out_channels, cond_dim, kernel_size=3, n_groups=8): - super().__init__() - - self.blocks = nn.ModuleList( - [ - Conv1dBlock(in_channels, out_channels, kernel_size, n_groups=n_groups), - Conv1dBlock(out_channels, out_channels, kernel_size, n_groups=n_groups), - ] - ) - - # FiLM modulation https://arxiv.org/abs/1709.07871 - # predicts per-channel scale and bias - cond_channels = out_channels * 2 - self.out_channels = out_channels - self.cond_encoder = nn.Sequential( - nn.Mish(), - nn.Linear(cond_dim, cond_channels), - nn.Dropout(0.1), - nn.Unflatten(-1, (-1, 1)), - ) - - # make sure dimensions compatible - self.residual_conv = ( - nn.Conv1d(in_channels, out_channels, 1) - if in_channels != out_channels - else nn.Identity() - ) - - def forward(self, x, cond): - """ - x : [ batch_size x in_channels x horizon ] - cond : [ batch_size x cond_dim] - - returns: - out : [ batch_size x out_channels x horizon ] - """ - out = self.blocks[0](x) - embed = self.cond_encoder(cond) - - embed = embed.reshape(embed.shape[0], 2, self.out_channels, 1) - scale = embed[:, 0, ...] - bias = embed[:, 1, ...] - out = scale * out + bias - - out = self.blocks[1](out) - out = out + self.residual_conv(x) - return out - - -class ConditionalUnet1D(nn.Module): - def __init__( - self, - input_dim, - global_cond_dim, - diffusion_step_embed_dim=256, - down_dims=[256, 512, 1024], - # down_dims=[512, 1024, 2048], - kernel_size=5, - n_groups=8, - ): - """ - input_dim: Dim of actions. - global_cond_dim: Dim of global conditioning applied with FiLM - in addition to diffusion step embedding. This is usually obs_horizon * obs_dim - diffusion_step_embed_dim: Size of positional encoding for diffusion iteration k - down_dims: Channel size for each UNet level. - The length of this array determines numebr of levels. - kernel_size: Conv kernel size - n_groups: Number of groups for GroupNorm - """ - - super().__init__() - all_dims = [input_dim] + list(down_dims) - start_dim = down_dims[0] - - dsed = diffusion_step_embed_dim - diffusion_step_encoder = nn.Sequential( - SinusoidalPosEmb(dsed), - nn.Linear(dsed, dsed * 4), - nn.Mish(), - nn.Linear(dsed * 4, dsed), - ) - cond_dim = dsed + global_cond_dim - - in_out = list(zip(all_dims[:-1], all_dims[1:])) - mid_dim = all_dims[-1] - self.mid_modules = nn.ModuleList( - [ - ConditionalResidualBlock1D( - mid_dim, - mid_dim, - cond_dim=cond_dim, - kernel_size=kernel_size, - n_groups=n_groups, - ), - ConditionalResidualBlock1D( - mid_dim, - mid_dim, - cond_dim=cond_dim, - kernel_size=kernel_size, - n_groups=n_groups, - ), - ] - ) - - down_modules = nn.ModuleList([]) - for ind, (dim_in, dim_out) in enumerate(in_out): - is_last = ind >= (len(in_out) - 1) - down_modules.append( - nn.ModuleList( - [ - ConditionalResidualBlock1D( - dim_in, - dim_out, - cond_dim=cond_dim, - kernel_size=kernel_size, - n_groups=n_groups, - ), - ConditionalResidualBlock1D( - dim_out, - dim_out, - cond_dim=cond_dim, - kernel_size=kernel_size, - n_groups=n_groups, - ), - Downsample1d(dim_out) if not is_last else nn.Identity(), - ] - ) - ) - - up_modules = nn.ModuleList([]) - for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): - is_last = ind >= (len(in_out) - 1) - up_modules.append( - nn.ModuleList( - [ - ConditionalResidualBlock1D( - dim_out * 2, - dim_in, - cond_dim=cond_dim, - kernel_size=kernel_size, - n_groups=n_groups, - ), - ConditionalResidualBlock1D( - dim_in, - dim_in, - cond_dim=cond_dim, - kernel_size=kernel_size, - n_groups=n_groups, - ), - Upsample1d(dim_in) if not is_last else nn.Identity(), - ] - ) - ) - - final_conv = nn.Sequential( - Conv1dBlock(start_dim, start_dim, kernel_size=kernel_size), - nn.Conv1d(start_dim, input_dim, 1), - ) - - self.diffusion_step_encoder = diffusion_step_encoder - self.up_modules = up_modules - self.down_modules = down_modules - self.final_conv = final_conv - - # print("number of parameters: {:e}".format( - # sum(p.numel() for p in self.parameters())) - # ) - - def forward( - self, - sample: torch.Tensor, - timestep: Union[torch.Tensor, float, int], - global_cond=None, - ): - """ - x: (B,T,input_dim) - timestep: (B,) or int, diffusion step - global_cond: (B,global_cond_dim) - output: (B,T,input_dim) - """ - # (B,T,C) - sample = sample.moveaxis(-1, -2) - # (B,C,T) - - # 1. time - timesteps = timestep - if not torch.is_tensor(timesteps): - # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can - timesteps = torch.tensor( - [timesteps], dtype=torch.long, device=sample.device - ) - elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0: - timesteps = timesteps[None].to(sample.device) - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timesteps = timesteps.expand(sample.shape[0]) - - global_feature = self.diffusion_step_encoder(timesteps) - - if global_cond is not None: - global_feature = torch.cat([global_feature, global_cond], axis=-1) - x = sample - h = [] - for idx, (resnet, resnet2, downsample) in enumerate(self.down_modules): - x = resnet(x, global_feature) - x = resnet2(x, global_feature) - h.append(x) - x = downsample(x) # bs, 2048, 5 - - for mid_module in self.mid_modules: - x = mid_module(x, global_feature) # bs, 2048, 5 - - for idx, (resnet, resnet2, upsample) in enumerate(self.up_modules): - x = torch.cat((x, h.pop()), dim=1) - x = resnet(x, global_feature) - x = resnet2(x, global_feature) - x = upsample(x) # bs, 512, 20 - - x = self.final_conv(x) # bs, 14, 20 - - # (B,C,T) - x = x.moveaxis(-1, -2) - # (B,T,C) - return x - - -class DP_Action_head(nn.Module): - def __init__( - self, - action_dim=14, - transformer_dim=896, - global_cond_dim=1806, - load_pretrained=True, - ): - super().__init__() - self.action_dim = action_dim - self.transformer_dim = transformer_dim - self.load_pretrained = load_pretrained - self.noise_scheduler = DDPMScheduler( - num_train_timesteps=132, - beta_schedule="squaredcos_cap_v2", - clip_sample=True, - prediction_type="epsilon", - ) - if self.load_pretrained: - self.global_cond_dim = global_cond_dim - self.condition_proj = nn.Sequential( - nn.Linear(self.transformer_dim, 2 * self.transformer_dim), - nn.ReLU(), - nn.Linear(2 * self.transformer_dim, 2 * self.global_cond_dim), - nn.ReLU(), - nn.Linear(2 * self.global_cond_dim, self.global_cond_dim), - ) - - self.noise_pred_net = ConditionalUnet1D( - input_dim=self.action_dim, - # down_dims=[256,512,1024], - down_dims=[512, 1024, 2048], - global_cond_dim=self.global_cond_dim, - ) - - # load pretrained model - action_pretrained_path = "/x2robot/liangyuxin/workspace/DiffusionPolicy/big_mix_0718_mn/30_noise_pred_net.pth" - print("load noise_pred_net from:", action_pretrained_path, flush=True) - self.noise_pred_net.load_state_dict(torch.load(action_pretrained_path)) - else: - self.noise_pred_net = ConditionalUnet1D( - input_dim=self.action_dim, - down_dims=[256, 512, 1024], - global_cond_dim=self.transformer_dim, - ) - - def forward(self, naction, condition, sample_times): - bs = naction.shape[0] - noise_shape = ( - naction.shape[0] * sample_times, - naction.shape[1], - naction.shape[2], - ) - noise = torch.randn(noise_shape, device=naction.device) - naction = ( - naction.unsqueeze(1) - .repeat(1, sample_times, 1, 1) - .reshape(bs * sample_times, naction.shape[1], naction.shape[2]) - ) - - timesteps = torch.randint( - 0, - self.noise_scheduler.config.num_train_timesteps, - (bs * sample_times,), - device=naction.device, - ).long() - condition = condition.to(self.condition_proj[0].weight.data.dtype) - if self.load_pretrained: - condition = self.condition_proj(condition) - - noisy_actions = self.noise_scheduler.add_noise(naction, noise, timesteps) - noise_pred = self.noise_pred_net( - noisy_actions, timesteps, global_cond=condition - ) - return noise, noise_pred - - @torch.no_grad() - def predict(self, condition, naction=None): - bs = condition.shape[0] - condition = condition.to(self.condition_proj[0].weight.data.dtype) - if self.load_pretrained: - condition = self.condition_proj(condition) - - if naction is not None: - noise_shape = (naction.shape[0], naction.shape[1], naction.shape[2]) - else: - noise_shape = (bs, 16, self.action_dim) # tobe parameterized - noise = torch.randn(noise_shape, device=condition.device) - naction_pred = noise - # init scheduler - self.noise_scheduler.set_timesteps( - self.noise_scheduler.config.num_train_timesteps - ) - - for k in self.noise_scheduler.timesteps: - # predict noise - noise_pred = self.noise_pred_net( - sample=naction_pred, timestep=k, global_cond=condition - ) - - # inverse diffusion step (remove noise) - naction_pred = self.noise_scheduler.step( - model_output=noise_pred, timestep=k, sample=naction_pred - ).prev_sample - - return naction, naction_pred - - -class ActionProcessor(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.dof_config = config.dof_config - self.agent_pos_config = config.agent_pos_config - self.action_dim = sum([v for k, v in self.dof_config.items()]) - self.propri_dim = sum([v for k, v in self.agent_pos_config.items()]) - - print_rank_last( - f"self.dof_config: {self.dof_config}; action_dim: {self.action_dim}; self.agent_pos_config: {self.agent_pos_config}; propri_dim: {self.propri_dim}" - ) - - self.action_hidden_size = config.action_hidden_size - self.state_hidden_size = config.state_hidden_size - self.hidden_size = config.hidden_size - - if not self.config.use_state_string_representation: - if self.config.proj_with_mask: - self.propri_proj = nn.Linear( - self.propri_dim * 2, self.state_hidden_size, bias=False - ) - else: - self.propri_proj = nn.Linear( - self.propri_dim, self.state_hidden_size, bias=False - ) - - # noise scheduler configing - if getattr(self.config, "use_flow_action_expert", True): - noise_scheduler_config = config.noise_scheduler - self.beta_alpha = noise_scheduler_config.get("beta_alpha", 1.5) - self.beta_beta = noise_scheduler_config.get("beta_beta", 1.0) - self.s = noise_scheduler_config.get("s", 0.999) - alpha_tensor = torch.tensor(self.beta_alpha, dtype=torch.float32).to("cuda") - beta_tensor = torch.tensor(self.beta_beta, dtype=torch.float32).to("cuda") - self.beta_dist = Beta(alpha_tensor, beta_tensor) - self.time_embed = SinusoidalPosEmb(self.action_hidden_size) - - # project to hidden space - if self.config.proj_with_mask: - self.w1 = nn.Linear( - self.action_dim * 2, self.action_hidden_size, bias=False - ) - else: - self.w1 = nn.Linear( - self.action_dim, self.action_hidden_size, bias=False - ) - if not self.config.use_adarms: - self.w2 = nn.Linear( - self.action_hidden_size * 2, self.action_hidden_size, bias=False - ) - self.w3 = nn.Linear( - self.action_hidden_size, self.action_hidden_size, bias=False - ) - self.act_fn = nn.SiLU() - else: - self.time_mlp_in = nn.Linear( - self.action_hidden_size, self.action_hidden_size - ) - self.time_mlp_out = nn.Linear( - self.action_hidden_size, self.action_hidden_size - ) - self.act_fn = nn.SiLU() - - # project back to action space - self.action_proj_back = nn.Linear( - self.action_hidden_size, self.action_dim, bias=False - ) - self.mse_loss = nn.MSELoss(reduction="none") - - def set_normalizer(self, normalizer_action, normalizer_propri): - self.normalizer_action = normalizer_action - self.normalizer_propri = normalizer_propri - - # dataset_name = self.config["data"]["lerobot_config"]["repo_id"] - # print("normalizer_propri min", self.normalizer_propri.min.__getattr__(dataset_name), flush=True) - # print("normalizer_propri delta", self.normalizer_propri.delta.__getattr__(dataset_name), flush=True) - # print("normalizer_action min", self.normalizer_action.min.__getattr__(dataset_name), flush=True) - # print("normalizer_action delta", self.normalizer_action.delta.__getattr__(dataset_name), flush=True) - - def sample_time(self, batch_size, device, dtype): - """ - Sampling Time Step - Generates random numbers in the range [0, 1] using a Beta distribution, and then scales them. - - Parameters: - batch_size (int): Batch size - device: Device type - dtype: Data type - - Returns: - torch.Tensor: Sampled time steps, with shape [batch_size] - """ - sample = self.beta_dist.sample([batch_size]).to(device=device, dtype=dtype) - time = (1 - sample) * self.s - return time - - def proprioception_proj( - self, proprioception, dataset_names=None, dof_mask=None, use_history=False - ): - """ - proprioception: [batch_size, 1, action_dim] - dataset_names: [batch_size] - dof_mask: [batch_size, action_dim] - """ - proprioception = proprioception.to(device=self.propri_proj.weight.device).to( - dtype=self.propri_proj.weight.dtype - ) - if dof_mask is not None: - if self.config.proj_with_mask: - proprioception = torch.cat( - [proprioception, dof_mask], dim=-1 - ) # .unsqueeze(1) - proprioception = proprioception.to(device=self.propri_proj.weight.device).to( - dtype=self.propri_proj.weight.dtype - ) - proprio_embed = self.propri_proj( - proprioception - ) # [batch_size, 1, state_hidden_size] - - if self.state_hidden_size < self.hidden_size: - # padding to hidden size - padding_size = self.hidden_size - self.state_hidden_size - padding = torch.zeros( - (proprio_embed.shape[0], 1, padding_size), - device=proprio_embed.device, - dtype=proprio_embed.dtype, - ) - proprio_embed = torch.cat([proprio_embed, padding], dim=-1) - - return proprio_embed # [batch_size, 1, hidden_size] - - def forward(self, action_chunk, dataset_names, dof_mask=None): - """ - Parameters: - action_chunk (torch.Tensor): Action sequence, shape [batch_size, action_chunk_len, action_dim] - dataset_names: [batch_size] - dof_mask: [batch_size, action_dim] - - Returns: - torch.Tensor: Processed action representation, shape [batch_size, seq_len, hidden_size] - """ - with torch.autocast("cuda", dtype=torch.float32): - action_chunk = action_chunk.to(dtype=torch.float32) - batch_size = action_chunk.shape[0] - device = action_chunk.device - dtype = action_chunk.dtype - - # 1. add noise to action_chunk - noise = torch.randn_like(action_chunk) - time = self.sample_time(batch_size, device, dtype) - time_expanded = time.unsqueeze(-1).unsqueeze(-1) - noisy_action = (1 - time_expanded) * noise + time_expanded * action_chunk - flow = action_chunk - noise - - # 2. sinusoidal positional encoding for timesteps - time_embed = self.time_embed(time).to(torch.float32) - - self.noise = noise - self.noisy_action = noisy_action # for new x-pred - - # 3.action_chunk_nosiy + t_pos_emb -> MLP_act_chunk -> action_chunk_nosiy_emb_with_t (dim=trans * chunk) - if dof_mask is not None: - noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) - - noisy_action = noisy_action.to(dtype=self.w1.weight.dtype) - action_embed = self.w1(noisy_action) - - self.time_expanded = time_expanded # for new x-pred - - if not self.config.use_adarms: - time_embed = ( - time_embed.unsqueeze(1) - .repeat(1, action_embed.shape[1], 1) - .to(dtype=self.w2.weight.dtype) - ) - concat_embed = torch.cat([action_embed, time_embed], dim=-1) - concat_embed = self.w2(concat_embed) - action_time_embed = self.w3(self.act_fn(concat_embed)) - adarms_cond = None - else: - time_embed = self.time_mlp_in(time_embed) - time_embed = self.act_fn(time_embed) - time_embed = self.time_mlp_out(time_embed) - time_embed = self.act_fn(time_embed) - action_time_embed = action_embed - adarms_cond = time_embed - - if self.action_hidden_size < self.hidden_size: - # padding to hidden size - padding_size = self.hidden_size - self.action_hidden_size - padding = torch.zeros( - ( - action_time_embed.shape[0], - action_time_embed.shape[1], - padding_size, - ), - device=action_time_embed.device, - dtype=action_time_embed.dtype, - ) - action_time_embed = torch.cat([action_time_embed, padding], dim=-1) - - return action_time_embed, flow, adarms_cond - - def step(self, timestep, noisy_action, dof_mask=None): - # noisy_action: bs, pred_horizon, action_dim - # timestep: bs - with torch.autocast("cuda", dtype=torch.float32): - if dof_mask is not None and self.config.proj_with_mask: - if dof_mask.shape[1] == 1: - dof_mask = dof_mask.unsqueeze(1).repeat(1, noisy_action.shape[1], 1) - noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) - - noisy_action = noisy_action.to(dtype=self.w1.weight.dtype) - time_embed = self.time_embed(timestep).to(torch.float32) # bs,hidden_size - action_embed = self.w1(noisy_action) - - if not self.config.use_adarms: - time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1) - time_embed = time_embed.to(device=noisy_action.device).to( - dtype=noisy_action.dtype - ) - concat_embed = torch.cat([action_embed, time_embed], dim=-1) - concat_embed = self.w2(concat_embed) - embed = self.w3(self.act_fn(concat_embed)) # is this right? - adarms_cond = None - else: - time_embed = time_embed.to(dtype=self.time_mlp_in.weight.dtype) - time_embed = self.time_mlp_in(time_embed) - time_embed = self.act_fn(time_embed) - time_embed = self.time_mlp_out(time_embed) - time_embed = self.act_fn(time_embed) - embed = action_embed - adarms_cond = time_embed - - if self.action_hidden_size < self.hidden_size: - # padding to hidden size - padding_size = self.hidden_size - self.action_hidden_size - padding = torch.zeros( - (embed.shape[0], embed.shape[1], padding_size), - device=embed.device, - dtype=embed.dtype, - ) - embed = torch.cat([embed, padding], dim=-1) - - return embed, adarms_cond - - def flow_loss( - self, - action_hidden_states, - flow, - action_chunk, - dof_mask=None, - flow_loss_mask=None, - ): - with torch.autocast("cuda", dtype=torch.float32): - action_pred = self.action_proj_back( - action_hidden_states[:, : self.action_hidden_size] - ) - v_pred = action_pred - loss = self.mse_loss(v_pred, flow) - if dof_mask is not None: - dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1]) - loss = loss * dof_mask - - if flow_loss_mask is not None: - flow_loss_mask = ( - flow_loss_mask.unsqueeze(-1) - .reshape(-1, 1) - .expand(-1, loss.shape[-1]) - ) - loss = loss * flow_loss_mask - return loss diff --git a/wall_x/model/core/__init__.py b/wall_x/model/core/__init__.py new file mode 100644 index 00000000..de478ad3 --- /dev/null +++ b/wall_x/model/core/__init__.py @@ -0,0 +1 @@ +"""Core building blocks shared across model implementations.""" diff --git a/wall_x/model/core/action/__init__.py b/wall_x/model/core/action/__init__.py new file mode 100644 index 00000000..aa2a816f --- /dev/null +++ b/wall_x/model/core/action/__init__.py @@ -0,0 +1,10 @@ +"""Action model components: normalizer, processor, head, MoE.""" + +from wall_x.model.core.action.normalizer import ( + Normalizer, + create_normalizers, + normalize_data_with_virtual_tail, + unnormalize_data_with_virtual_tail, +) +from wall_x.model.core.action.processor import ActionProcessor +from wall_x.model.core.action.head import SinusoidalPosEmb diff --git a/wall_x/model/core/action/head.py b/wall_x/model/core/action/head.py new file mode 100644 index 00000000..30a1229c --- /dev/null +++ b/wall_x/model/core/action/head.py @@ -0,0 +1,27 @@ +"""Action head helpers used by the VLA action processor.""" + +import math + +import torch +import torch.nn as nn + + +class SinusoidalPosEmb(nn.Module): + """Sinusoidal timestep embedding for action flow timesteps.""" + + def __init__(self, dim: int, min_period: float = 4e-3, max_period: float = 4.0): + super().__init__() + if dim % 2 != 0: + raise ValueError(f"embedding_dim ({dim}) must be divisible by 2") + self.dim = dim + self.min_period = min_period + self.max_period = max_period + + def forward(self, x): + half_dim = self.dim // 2 + exponent = math.log(10000) / (half_dim - 1) + frequencies = torch.exp( + torch.arange(half_dim, device=x.device, dtype=torch.float32) * -exponent + ) + emb = x[:, None] * frequencies[None, :] + return torch.cat((emb.sin(), emb.cos()), dim=-1) diff --git a/wall_x/model/core/action/moe.py b/wall_x/model/core/action/moe.py new file mode 100644 index 00000000..cffe2bff --- /dev/null +++ b/wall_x/model/core/action/moe.py @@ -0,0 +1,127 @@ +import torch +import torch.nn as nn +import torch.utils.checkpoint as cp + +from transformers.activations import ACT2FN + +from wall_x.model.core.ops import unpermute, permute + + +class TokenTypeRouter(nn.Module): + def __init__(self, num_experts: int): + super().__init__() + self.num_experts = num_experts + + def forward(self, token_types: torch.Tensor) -> torch.Tensor: + """ + Route tokens to experts based on token_type. + + Args: + token_types (torch.Tensor): Tensor of shape (batch_size, seq_length) containing each token type. + + Returns: + experts_indices (torch.Tensor): Tensor of shape (batch_size, seq_length) containing each assigned expert index. + """ + # Simple rule: assign by token_type modulo the expert count + experts_indices = token_types % self.num_experts + return experts_indices + + +class BlockSparseMLP(nn.Module): + def __init__(self, config, use_selective_recompute: bool = False): + super().__init__() + self.hidden_size = config["hidden_size"] + self.intermediate_size = config["intermediate_size"] + self.hidden_act = config["hidden_act"] + + self.use_selective_recompute = use_selective_recompute + + self.gate_up_proj = nn.Linear( + self.hidden_size, 2 * self.intermediate_size, bias=False + ) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + + self.act_fn = ACT2FN[self.hidden_act] + + # FIXME: The full MLP is recomputed for now; recomputing only activations can be optimized later. + def _full_mlp(self, hidden_state): + gate_up_out = self.gate_up_proj(hidden_state) + gate_out, up_out = gate_up_out.split( + [self.intermediate_size, self.intermediate_size], dim=-1 + ) + + act_out = self.act_fn(gate_out) * up_out + return self.down_proj(act_out) + + def forward(self, hidden_state): + if self.use_selective_recompute: + # Checkpoint-recompute the whole expert MLP + return cp.checkpoint( + self._full_mlp, + hidden_state, + use_reentrant=False, + ) + else: + return self._full_mlp(hidden_state) + + +class SparseMoeBlock(nn.Module): + def __init__(self, config, num_experts: int, use_selective_recompute: bool = False): + super().__init__() + self.num_experts = num_experts + self.use_selective_recompute = use_selective_recompute + + # Pass use_selective_recompute to each expert + self.experts = nn.ModuleList( + [ + BlockSparseMLP( + config.experts[i], use_selective_recompute=use_selective_recompute + ) + for i in range(num_experts) + ] + ) + + if not hasattr(config, "dim_inputs") or not config.dim_inputs: + raise ValueError("config.dim_inputs must be set") + + self.dim_inputs = config.dim_inputs + self.permuted = config.mot_opt + + def forward( + self, + hidden_states: torch.Tensor, + experts_indices: torch.Tensor, + start_indices: torch.Tensor, + end_indices: torch.Tensor, + ) -> torch.Tensor: + + if self.permuted: + permuted_inputs = hidden_states + else: + batch_size, seq_length, hidden_dim = hidden_states.shape + + flat_hidden = hidden_states.reshape(-1, hidden_dim) + experts_indices = experts_indices.reshape(-1) + probs = torch.ones_like(experts_indices, dtype=torch.float32).reshape(-1, 1) + permuted_inputs, row_id_map = permute(flat_hidden, experts_indices) + + # buffer + final_output = torch.zeros_like(permuted_inputs) + + # Expert forward, including selective recompute + for expert_idx, expert in enumerate(self.experts): + start, end = start_indices[expert_idx], end_indices[expert_idx] + if start == end: + continue + + dim_input = self.dim_inputs[expert_idx] + expert_input = permuted_inputs[start:end, :dim_input] + + partial_output = expert(expert_input) + final_output[start:end, :dim_input] = partial_output[:, :dim_input] + + if self.permuted: + return final_output + else: + final_output = unpermute(final_output, row_id_map, probs) + return final_output.reshape(batch_size, seq_length, hidden_dim) diff --git a/wall_x/model/core/action/normalizer.py b/wall_x/model/core/action/normalizer.py new file mode 100644 index 00000000..ca7b3341 --- /dev/null +++ b/wall_x/model/core/action/normalizer.py @@ -0,0 +1,447 @@ +import json +import logging + +import torch +import torch.nn as nn + +from wall_x.utils.constant import _ACTION_KEY_FULL_MAPPING as _MODEL_KEY_TO_RAW_KEY + +logger = logging.getLogger(__name__) + + +def resolve_normalizer_dataset_names( + dataset_names, + normalizer, + batch_size, + *, + source="dataset_names", + allow_skip_names=None, +): + """Validate dataset names before indexing a normalizer ParameterDict. + + The normalizer key is data-dependent. Callers should pass the + dataset names carried by the batch/model inputs, matching the qact/VLA + path. This helper intentionally only validates presence, batch-size + alignment, and key membership; it does not remap unknown names or fall + back to a default key, since silent fallback can apply the wrong action + scale without failing loudly. Pass ``allow_skip_names`` explicitly only + for names that should bypass normalizer lookup, such as multimodal-only + rows that do not carry action values. + """ + if dataset_names is None: + raise KeyError(f"Missing {source}; cannot choose normalizer key.") + + if isinstance(dataset_names, str): + names = [dataset_names] + elif isinstance(dataset_names, (list, tuple)): + names = [str(name) for name in dataset_names] + else: + raise TypeError( + f"{source} must be a string or sequence of strings, got " + f"{type(dataset_names).__name__}" + ) + + if len(names) != batch_size: + raise ValueError( + f"{source} count ({len(names)}) does not match batch size " + f"({batch_size}). names={names}" + ) + + available = list(getattr(normalizer, "delta", {}).keys()) + if not available: + raise KeyError("Normalizer has no registered dataset keys.") + + allowed_skip = set(allow_skip_names or ()) + missing = sorted( + {name for name in names if name not in available and name not in allowed_skip} + ) + if missing: + raise KeyError( + f"{source} contains keys not present in normalizer: {missing}. " + f"available={available}" + ) + + return names + + +def _normalizer_dataset_name_list(dataset_names): + if dataset_names is None: + return [] + if isinstance(dataset_names, str): + return [dataset_names] + return list(dataset_names) + + +def _normalizer_width(normalizer, dataset_names): + names = _normalizer_dataset_name_list(dataset_names) + if not names: + return None + if any(name not in normalizer.delta for name in names): + return None + widths = [int(normalizer.delta[name].shape[0]) for name in names] + if not widths or len(set(widths)) != 1: + return None + return widths[0] + + +def _has_uniform_virtual_tail_mask(dof_mask, width): + if dof_mask is None: + return False + mask = dof_mask.reshape(-1, dof_mask.shape[-1]).bool() + if mask.shape[-1] <= width: + return False + prefix_active = mask[:, :width].all(dim=1) + tail_inactive = (~mask[:, width:]).all(dim=1) + return bool((prefix_active & tail_inactive).all().item()) + + +def normalize_data_with_virtual_tail(normalizer, values, dataset_names, dof_mask): + """Normalize real prefix dims when model tensors include a virtual tail.""" + width = _normalizer_width(normalizer, dataset_names) + if width is None or values.shape[-1] == width: + return normalizer.normalize_data(values, dataset_names) + if values.shape[-1] < width or not _has_uniform_virtual_tail_mask(dof_mask, width): + return normalizer.normalize_data(values, dataset_names) + + out = values.clone() + out[..., :width] = normalizer.normalize_data(values[..., :width], dataset_names) + out[..., width:] = 0 + return out + + +def unnormalize_data_with_virtual_tail(normalizer, values, dataset_names, dof_mask): + """Unnormalize real prefix dims when model tensors include a virtual tail.""" + width = _normalizer_width(normalizer, dataset_names) + if width is None or values.shape[-1] == width: + return normalizer.unnormalize_data(values, dataset_names, None) + if values.shape[-1] < width or not _has_uniform_virtual_tail_mask(dof_mask, width): + return normalizer.unnormalize_data(values, dataset_names, dof_mask) + + out = values.clone() + out[..., :width] = normalizer.unnormalize_data( + values[..., :width], dataset_names, None + ) + out[..., width:] = 0 + return out + + +def print_rank_last(message): + """If distributed is initialized, log only on last rank.""" + if torch.distributed.is_initialized(): + if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1): + logger.info(message) + else: + logger.info(message) + + +class Normalizer(nn.Module): + @classmethod + def from_ckpt(cls, ckpt_path): + instance = cls.__new__(cls) + nn.Module.__init__(instance) + + instance.min = nn.ParameterDict() + instance.delta = nn.ParameterDict() + instance.min_key = "min" + instance.delta_key = "delta" + + ckpt = torch.load(ckpt_path, map_location="cpu") + + for key, value in ckpt.items(): + # parse key: "min.robot_name" -> prefix="min", name="robot_name" + try: + prefix, name = key.split(".", 1) + if hasattr(instance, prefix): + getattr(instance, prefix)[name] = nn.Parameter( + value, requires_grad=False + ) + except ValueError: + continue + + return instance + + @classmethod + def from_lerobot_norm_stats( + cls, + action_stats, + dataset_name, + ): + + # Create instance without calling __init__ + instance = cls.__new__(cls) + nn.Module.__init__(instance) + + # Set containers + instance.min = nn.ParameterDict() + instance.delta = nn.ParameterDict() + + # Fill dataset entry + instance.min[dataset_name] = nn.Parameter(action_stats.min, requires_grad=False) + instance.delta[dataset_name] = nn.Parameter( + action_stats.delta, requires_grad=False + ) + + # Record keys + instance.min_key = "min" + instance.delta_key = "delta" + + return instance + + def __init__( + self, + action_statistic_dof, + dof_config, + min_key="min", + delta_key="delta", + name="normalizer", + ): + super(Normalizer, self).__init__() + + self.min_key = min_key + self.delta_key = delta_key + + action_statistic = {} + normalizer_missing_information = [] + for robot_name in action_statistic_dof.keys(): + action_statistic[robot_name] = {} + all_dof_min = [] + all_dof_delta = [] + for k in dof_config: + if ( + k not in action_statistic_dof[robot_name] + and k.replace("master_", "follow_", 1) + in action_statistic_dof[robot_name] + ): + k = k.replace("master_", "follow_", 1) + if k in action_statistic_dof[robot_name]: + if ( + min_key in action_statistic_dof[robot_name][k] + and delta_key in action_statistic_dof[robot_name][k] + ): + all_dof_min.extend(action_statistic_dof[robot_name][k][min_key]) + all_dof_delta.extend( + action_statistic_dof[robot_name][k][delta_key] + ) + else: + normalizer_missing_information.append( + f"Normalizer (Warning): min_key {min_key} or delta_key {delta_key} not in action_statistic_dof[{robot_name}][{k}], use default min 0.0 and delta 1.0" + ) + all_dof_min.extend([0.0] * dof_config[k]) + all_dof_delta.extend([1.0] * dof_config[k]) + else: + # k is a model key; action_statistic_dof may store raw keys — try fallback lookup + raw_key = _MODEL_KEY_TO_RAW_KEY.get(k) + if ( + raw_key is not None + and raw_key in action_statistic_dof[robot_name] + ): + stat = action_statistic_dof[robot_name][raw_key] + if min_key in stat and delta_key in stat: + all_dof_min.extend(stat[min_key]) + all_dof_delta.extend(stat[delta_key]) + else: + normalizer_missing_information.append( + f"Normalizer (Warning): Action {k} not in action_statistic_dof for {robot_name}, use default min 0.0 and delta 1.0" + ) + all_dof_min.extend([0.0] * dof_config[k]) + all_dof_delta.extend([1.0] * dof_config[k]) + else: + normalizer_missing_information.append( + f"Normalizer (Warning): Action {k} not in action_statistic_dof for {robot_name}, use default min 0.0 and delta 1.0" + ) + all_dof_min.extend([0.0] * dof_config[k]) + all_dof_delta.extend([1.0] * dof_config[k]) + all_dof_min = torch.tensor(all_dof_min) + all_dof_delta = torch.tensor(all_dof_delta) + action_statistic[robot_name][min_key] = all_dof_min + action_statistic[robot_name][delta_key] = all_dof_delta + + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == ( + torch.distributed.get_world_size() - 1 + ): + # export normalizer_missing_information to file + with open(f"normalizer_missing_information_{name}.txt", "w") as f: + for info in normalizer_missing_information: + f.write(info + "\n") + logger.info( + "Normalizer missing information saved to normalizer_missing_information_%s.txt", + name, + ) + + self.min = nn.ParameterDict( + { + k: nn.Parameter(action_statistic[k][min_key], requires_grad=False) + for k in action_statistic.keys() + } + ) + self.delta = nn.ParameterDict( + { + k: nn.Parameter(action_statistic[k][delta_key], requires_grad=False) + for k in action_statistic.keys() + } + ) + + def normalize_data(self, xs, dataset_names): + new_xs = [] + dataset_names = [name for name in dataset_names if name != "x2_multimodal"] + for x, dataset_name in zip(xs, dataset_names): + # if dataset_name == "ex_normal": + # dataset_name = "x2_normal" + x = (x - self.min[dataset_name]) / (self.delta[dataset_name]) + x = x * 2 - 1 + x = torch.clamp(x, -1, 1) + new_xs.append(x) + new_xs = torch.stack(new_xs) + return new_xs + + def unnormalize_data(self, xs, dataset_names, dof_mask=None): + new_xs = [] + dataset_names = [name for name in dataset_names if name != "x2_multimodal"] + dof_mask = dof_mask if dof_mask is not None else [None] * len(xs) + for x, dataset_name, mask in zip(xs, dataset_names, dof_mask): + x = (x + 1) / 2 + if mask is not None: + mask = mask[0].bool() + action_space_delta = self.delta[dataset_name][mask] + action_space_min = self.min[dataset_name][mask] + else: + action_space_delta = self.delta[dataset_name] + action_space_min = self.min[dataset_name] + x = x * action_space_delta + action_space_min + new_xs.append(x) + new_xs = torch.stack(new_xs) + return new_xs + + +def _pad_lerobot_stats(stats, target_dim, label): + """Tail-pad flat LeRobot stats to the model's configured action/state dim.""" + target_dim = int(target_dim) + current_dim = int(stats.min.numel()) + if current_dim == target_dim: + return stats + if current_dim > target_dim: + raise ValueError( + f"LeRobot {label} norm stats dim ({current_dim}) exceeds configured " + f"{label} dim ({target_dim})" + ) + + pad_dim = target_dim - current_dim + min_stat = torch.cat( + [stats.min.flatten(), torch.zeros(pad_dim, dtype=stats.min.dtype)] + ) + delta = torch.cat( + [stats.delta.flatten(), torch.ones(pad_dim, dtype=stats.delta.dtype)] + ) + max_stat = torch.cat( + [stats.max.flatten(), torch.ones(pad_dim, dtype=stats.max.dtype)] + ) + logger.info( + "Padded LeRobot %s normalizer stats from %d to %d dims", + label, + current_dim, + target_dim, + ) + return type(stats)(min=min_stat, max=max_stat, delta=delta) + + +def pad_normalizer_to_dim(normalizer, target_dim, label): + """Tail-pad every dataset entry to ``target_dim`` (virtual ``action_padding`` tail).""" + target_dim = int(target_dim) + for name in list(normalizer.min.keys()): + current_dim = int(normalizer.min[name].numel()) + if current_dim == target_dim: + continue + if current_dim > target_dim: + raise ValueError( + f"Normalizer {label}[{name!r}] dim ({current_dim}) exceeds configured " + f"{label} dim ({target_dim})" + ) + + pad_dim = target_dim - current_dim + dtype = normalizer.min[name].dtype + normalizer.min[name] = nn.Parameter( + torch.cat( + [ + normalizer.min[name].flatten(), + torch.zeros(pad_dim, dtype=dtype), + ] + ), + requires_grad=False, + ) + normalizer.delta[name] = nn.Parameter( + torch.cat( + [ + normalizer.delta[name].flatten(), + torch.ones(pad_dim, dtype=dtype), + ] + ), + requires_grad=False, + ) + logger.info( + "Padded %s normalizer[%r] from %d to %d dims", + label, + name, + current_dim, + target_dim, + ) + + +def create_normalizers_from_lerobot_norm_stats( + norm_stats, dataset_name, action_dim, propri_dim +): + """Create model normalizers from LeRobot flat action/state norm stats.""" + action_stats = _pad_lerobot_stats(norm_stats["action"], action_dim, "action") + propri_stats = _pad_lerobot_stats(norm_stats["state"], propri_dim, "state") + return ( + Normalizer.from_lerobot_norm_stats(action_stats, dataset_name), + Normalizer.from_lerobot_norm_stats(propri_stats, dataset_name), + ) + + +def create_normalizers(config, action_statistic_dof=None): + """Create action and proprioception normalizers from explicit stats. + + Normalization statistics must come from the config, checkpoint, or dataset. + Public Wall-X builds intentionally do not bundle private default stats. + """ + + if action_statistic_dof is None: + custom_path = config.get("customized_action_statistic_dof") + if custom_path: + with open(custom_path, "r") as f: + action_statistic_dof = json.load(f) + else: + raise ValueError( + "create_normalizers requires action statistics. Provide " + "`customized_action_statistic_dof` in the config or pass " + "`action_statistic_dof` from the checkpoint/dataset." + ) + + min_key = config.get("min_key", "min") + delta_key = config.get("delta_key", "delta") + + # Required keys for constructing action/proprio normalizers. + missing_required = [] + if config.get("dof_config") is None: + missing_required.append("dof_config") + if config.get("agent_pos_config") is None: + missing_required.append("agent_pos_config") + if missing_required: + raise KeyError( + "create_normalizers requires non-None config keys: " + + ", ".join(missing_required) + ) + + normalizer_action = Normalizer( + action_statistic_dof, + config["dof_config"], + min_key=min_key, + delta_key=delta_key, + ) + normalizer_propri = Normalizer( + action_statistic_dof, + config["agent_pos_config"], + min_key=min_key, + delta_key=delta_key, + ) + return normalizer_action, normalizer_propri, action_statistic_dof diff --git a/wall_x/model/core/action/processor.py b/wall_x/model/core/action/processor.py new file mode 100644 index 00000000..358352ff --- /dev/null +++ b/wall_x/model/core/action/processor.py @@ -0,0 +1,317 @@ +import torch +import torch.nn as nn + +from wall_x.model.core.action.head import SinusoidalPosEmb +from wall_x.model.core.action.normalizer import print_rank_last + + +class ActionProcessor(nn.Module): + """ + Action processor + + Main responsibilities: + 1. Add noise to action sequences + 2. Generate time encodings + 3. Project actions into the model hidden space + + Uses a Beta distribution to control noise scheduling and provide flexible noise injection. + """ + + def __init__(self, config): + """ + Args: + config: Configuration object containing: + - action_dim: action-space dimension + - hidden_size: model hidden size + - noise_scheduler: noise scheduler configuration + """ + super().__init__() + self.config = config + self.dof_config = config.dof_config + self.agent_pos_config = config.agent_pos_config + self.action_dim = sum([v for k, v in self.dof_config.items()]) + self.propri_dim = sum([v for k, v in self.agent_pos_config.items()]) + + print_rank_last( + f"self.dof_config: {self.dof_config}; action_dim: {self.action_dim}; self.agent_pos_config: {self.agent_pos_config}; propri_dim: {self.propri_dim}" + ) + + self.action_hidden_size = config.action_hidden_size + self.state_hidden_size = config.state_hidden_size + self.hidden_size = getattr(config, "hidden_size", config.dim_inputs[0]) + + if not self.config.use_state_string_representation: + if self.config.proj_with_mask: + self.propri_proj = nn.Linear( + self.propri_dim * 2, self.state_hidden_size, bias=False + ) + else: + self.propri_proj = nn.Linear( + self.propri_dim, self.state_hidden_size, bias=False + ) + + # noise scheduler configing + if getattr(self.config, "use_flow_action_expert", True): + noise_scheduler_config = config.noise_scheduler + self.s = noise_scheduler_config.get("s", 0.999) + self.time_shift = noise_scheduler_config.get( + "time_shift", 1.0 + ) # time shift factor + self.time_embed = SinusoidalPosEmb(self.action_hidden_size) + + # project to hidden space + if self.config.proj_with_mask: + self.w1 = nn.Linear( + self.action_dim * 2, self.action_hidden_size, bias=False + ) + else: + self.w1 = nn.Linear( + self.action_dim, self.action_hidden_size, bias=False + ) + if not self.config.use_adarms: + self.w2 = nn.Linear( + self.action_hidden_size * 2, self.action_hidden_size, bias=False + ) + self.w3 = nn.Linear( + self.action_hidden_size, self.action_hidden_size, bias=False + ) + self.act_fn = nn.SiLU() + else: + self.time_mlp_in = nn.Linear( + self.action_hidden_size, self.action_hidden_size + ) + self.time_mlp_out = nn.Linear( + self.action_hidden_size, self.action_hidden_size + ) + self.act_fn = nn.SiLU() + + # project back to action space + self.action_proj_back = nn.Linear( + self.action_hidden_size, self.action_dim, bias=False + ) + self.mse_loss = nn.MSELoss(reduction="none") + + def set_normalizer(self, normalizer_action, normalizer_propri): + self.normalizer_action = normalizer_action + self.normalizer_propri = normalizer_propri + + def get_inference_times(self, num_steps, device, dtype): + """ + Get inference timesteps + + Apply time shift and scaling + + Args: + num_steps (int): number of inference steps + device: Device type + dtype: dtype + + Returns: + torch.Tensor: inference timestep sequence + """ + times = torch.linspace(0.0, 1.0, num_steps + 1, device=device, dtype=dtype) + if self.time_shift != 1.0: + times = (self.time_shift * times) / (1 + (self.time_shift - 1) * times) + times = times * self.s + return times + + def proprioception_proj( + self, proprioception, dataset_names=None, dof_mask=None, use_history=False + ): + """ + Args: + proprioception: [batch_size, 1, action_dim] + dataset_names: [batch_size] + dof_mask: [batch_size, action_dim] + """ + with torch.autocast("cuda", dtype=torch.float32): + proprioception = proprioception.to( + device=self.propri_proj.weight.device + ).to(dtype=self.propri_proj.weight.dtype) + if dof_mask is not None: + if self.config.proj_with_mask: + proprioception = torch.cat( + [proprioception, dof_mask], dim=-1 + ) # .unsqueeze(1) + proprioception = proprioception.to( + device=self.propri_proj.weight.device + ).to(dtype=self.propri_proj.weight.dtype) + proprio_embed = self.propri_proj( + proprioception + ) # [batch_size, 1, state_hidden_size] + if self.state_hidden_size < self.hidden_size: + # padding to hidden size + padding_size = self.hidden_size - self.state_hidden_size + padding = torch.zeros( + (proprio_embed.shape[0], 1, padding_size), + device=proprio_embed.device, + dtype=proprio_embed.dtype, + ) + proprio_embed = torch.cat([proprio_embed, padding], dim=-1) + return proprio_embed # [batch_size, 1, hidden_size] + + def forward(self, action_chunk, dataset_names, sample_time, dof_mask=None): + """ + Args: + action_chunk (torch.Tensor): action sequence with shape [batch_size, action_chunk_len, action_dim] + dataset_names: [batch_size] + dof_mask: [batch_size, action_dim] + + Returns: + torch.Tensor: processed action representation with shape [batch_size, seq_len, hidden_size] + """ + with torch.autocast("cuda", dtype=torch.float32): + action_chunk = action_chunk.to(dtype=torch.float32) + + # 1. add noise to action_chunk + noise = torch.randn_like(action_chunk) + time_expanded = sample_time.unsqueeze(-1).unsqueeze(-1) + noisy_action = ( + 1 - time_expanded + ) * noise + time_expanded * action_chunk # denoise from 0 to 1; integration does not need a negative sign + flow = action_chunk - noise # used to compute loss + + # 2. sinusoidal positional encoding for timesteps + time_embed = self.time_embed(sample_time).to(torch.float32) + + self.noise = noise + self.noisy_action = noisy_action # for new x-pred + # 3.action_chunk_nosiy + t_pos_emb -> MLP_act_chunk -> action_chunk_nosiy_emb_with_t (dim=trans * chunk) + if dof_mask is not None: + noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) + + noisy_action = noisy_action.to(dtype=self.w1.weight.dtype) + action_embed = self.w1(noisy_action) + + self.time_expanded = time_expanded # for new x-pred + + if not self.config.use_adarms: + time_embed = ( + time_embed.unsqueeze(1) + .repeat(1, action_embed.shape[1], 1) + .to(dtype=self.w2.weight.dtype) + ) + concat_embed = torch.cat([action_embed, time_embed], dim=-1) + concat_embed = self.w2(concat_embed) + action_time_embed = self.w3(self.act_fn(concat_embed)) + adarms_cond = None + else: + time_embed = self.time_mlp_in(time_embed) + time_embed = self.act_fn(time_embed) + time_embed = self.time_mlp_out(time_embed) + time_embed = self.act_fn(time_embed) + action_time_embed = action_embed + adarms_cond = time_embed + + if self.action_hidden_size < self.hidden_size: + # padding to hidden size + padding_size = self.hidden_size - self.action_hidden_size + padding = torch.zeros( + ( + action_time_embed.shape[0], + action_time_embed.shape[1], + padding_size, + ), + device=action_time_embed.device, + dtype=action_time_embed.dtype, + ) + action_time_embed = torch.cat([action_time_embed, padding], dim=-1) + + return action_time_embed, flow, adarms_cond + + def step(self, timestep, noisy_action, dof_mask=None): + # noisy_action: bs, pred_horizon, action_dim + # timestep: bs + with torch.autocast("cuda", dtype=torch.float32): + if dof_mask is not None and self.config.proj_with_mask: + noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) + + noisy_action = noisy_action.to(dtype=self.w1.weight.dtype) + time_embed = self.time_embed(timestep).to(torch.float32) # bs,hidden_size + action_embed = self.w1(noisy_action) + + if not self.config.use_adarms: + time_embed = time_embed.unsqueeze(1).repeat(1, action_embed.shape[1], 1) + time_embed = time_embed.to(device=noisy_action.device).to( + dtype=noisy_action.dtype + ) + concat_embed = torch.cat([action_embed, time_embed], dim=-1) + concat_embed = self.w2(concat_embed) + embed = self.w3(self.act_fn(concat_embed)) # is this right? + adarms_cond = None + else: + time_embed = time_embed.to(dtype=self.time_mlp_in.weight.dtype) + time_embed = self.time_mlp_in(time_embed) + time_embed = self.act_fn(time_embed) + time_embed = self.time_mlp_out(time_embed) + time_embed = self.act_fn(time_embed) + embed = action_embed + adarms_cond = time_embed + + if self.action_hidden_size < self.hidden_size: + # padding to hidden size + padding_size = self.hidden_size - self.action_hidden_size + padding = torch.zeros( + (embed.shape[0], embed.shape[1], padding_size), + device=embed.device, + dtype=embed.dtype, + ) + embed = torch.cat([embed, padding], dim=-1) + + return embed, adarms_cond + + def flow_loss( + self, + action_hidden_states, + flow, + action_chunk, + dof_mask=None, + flow_loss_mask=None, + ): + with torch.autocast("cuda", dtype=torch.float32): + action_pred = self.action_proj_back( + action_hidden_states[:, : self.action_hidden_size] + ) + + if getattr(self.config, "use_x_pred", False): + noisy_action_flat = self.noisy_action.reshape( + -1, self.noisy_action.shape[-1] + ) + time_expanded_flat = self.time_expanded.expand( + -1, self.noisy_action.shape[1], -1 + ).reshape(-1, 1) + v_pred = (action_pred - noisy_action_flat) / torch.clamp( + 1 - time_expanded_flat, min=0.05 + ) + x_pred = action_pred + else: + v_pred = action_pred + time_expanded_flat = self.time_expanded.expand( + -1, self.noisy_action.shape[1], -1 + ).reshape(-1, 1) + x_pred = (1 - time_expanded_flat) * v_pred + self.noisy_action.reshape( + -1, self.noisy_action.shape[-1] + ) + + if getattr(self.config, "use_x_loss", False): + loss = self.mse_loss( + x_pred, + action_chunk.reshape(-1, action_chunk.shape[-1]).to( + dtype=x_pred.dtype + ), + ) + else: + loss = self.mse_loss(v_pred, flow) + + if dof_mask is not None: + dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1]) + loss = loss * dof_mask + + if flow_loss_mask is not None: + flow_loss_mask = ( + flow_loss_mask.unsqueeze(-1) + .reshape(-1, 1) + .expand(-1, loss.shape[-1]) + ) + loss = loss * flow_loss_mask + return loss diff --git a/wall_x/model/core/attention/__init__.py b/wall_x/model/core/attention/__init__.py new file mode 100644 index 00000000..a2397889 --- /dev/null +++ b/wall_x/model/core/attention/__init__.py @@ -0,0 +1,9 @@ +"""Attention mechanisms: joint attention, VLA attention, mask builders, backend selector.""" + +from wall_x.model.core.attention.mask import ( + find_first_last_ones, + update_position_ids, + update_joint_attention_mask_2d, + update_joint_attention_flash_mask, +) +from wall_x.model.core.attention.selector import AttentionsSelectorMixin diff --git a/wall_x/model/joint_attention.py b/wall_x/model/core/attention/joint.py similarity index 56% rename from wall_x/model/joint_attention.py rename to wall_x/model/core/attention/joint.py index 5ba95466..779cf06b 100644 --- a/wall_x/model/joint_attention.py +++ b/wall_x/model/core/attention/joint.py @@ -2,42 +2,37 @@ import torch.nn as nn from typing import Optional, Tuple +from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2_5_VLRotaryEmbedding, + repeat_kv, +) + from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig from transformers.cache_utils import Cache from transformers.utils import logging -from wall_x.fusions import ops -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import ( - apply_multimodal_rotary_pos_emb, -) + +from wall_x.model.core.attention.mask import find_first_last_ones + +from wall_x.model.core.ops import unpermute, permute, m_rope + from flash_attn import flash_attn_func + from transformers.modeling_flash_attention_utils import ( is_flash_attn_greater_or_equal_2_10, ) -from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( - Qwen2_5_VLRotaryEmbedding, - repeat_kv, -) - -logger = logging.get_logger(__name__) +try: + from flash_mask.flash_mask_interface import flash_mask_attn_func +except ImportError: + flash_mask_attn_func = None -# def rotate_half(x): -# x1 = x[..., : x.shape[-1] // 2] -# x2 = x[..., x.shape[-1] // 2 :] -# return torch.cat((-x2, x1), dim=-1) +try: + from flash_mask.flash_mask_interface import flashmask_attn_func_stop_gradient +except ImportError: + flashmask_attn_func_stop_gradient = None -# def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=2): -# mrope_section = mrope_section * 2 -# cos_split = torch.cat( -# [m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1 -# ).unsqueeze(unsqueeze_dim) -# sin_split = torch.cat( -# [m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1 -# ).unsqueeze(unsqueeze_dim) -# q_embed = (q * cos_split) + (rotate_half(q) * sin_split) -# k_embed = (k * cos_split) + (rotate_half(k) * sin_split) -# return q_embed, k_embed +logger = logging.get_logger(__name__) class JointQwen2VLAttention(nn.Module): @@ -52,7 +47,7 @@ def __init__(self, config: Qwen2_5_VLConfig, layer_idx: Optional[int] = None): "when creating this class." ) if not hasattr(config, "dim_inputs") or not config.dim_inputs: - raise ValueError("Configuration must contain a valid dim_inputs") + raise ValueError("config.dim_inputs must be set") self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads @@ -69,30 +64,18 @@ def __init__(self, config: Qwen2_5_VLConfig, layer_idx: Optional[int] = None): self.dim_inputs = config.dim_inputs # Tuple[int, ...] - if config.model_type == "qwen2_5_vl": - bias_qkv = True - else: - bias_qkv = False + if config.model_type != "qwen2_5_vl": + raise NotImplementedError(f"Unsupported model type: {config.model_type}") + bias_qkv = True - self.q_proj_experts = nn.ModuleList( - [ - nn.Linear(dim_input, self.num_heads * self.head_dim, bias=bias_qkv) - for dim_input in self.dim_inputs - ] + qkv_out_features = ( + self.num_heads * self.head_dim + + 2 * self.num_key_value_heads * self.head_dim ) - self.k_proj_experts = nn.ModuleList( - [ - nn.Linear( - dim_input, self.num_key_value_heads * self.head_dim, bias=bias_qkv - ) - for dim_input in self.dim_inputs - ] - ) - self.v_proj_experts = nn.ModuleList( + + self.qkv_proj_experts = nn.ModuleList( [ - nn.Linear( - dim_input, self.num_key_value_heads * self.head_dim, bias=bias_qkv - ) + nn.Linear(dim_input, qkv_out_features, bias=bias_qkv) for dim_input in self.dim_inputs ] ) @@ -103,11 +86,7 @@ def __init__(self, config: Qwen2_5_VLConfig, layer_idx: Optional[int] = None): ] ) - # Rotary embedding init - if config.model_type == "qwen2_5_vl": - self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) - else: - raise NotImplementedError(f"Unsupported model type: {config.model_type}") + self.rotary_emb = Qwen2_5_VLRotaryEmbedding(config=config) def repeat_kv(self, hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ @@ -129,6 +108,10 @@ def repeat_kv(self, hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: return hidden_states.reshape(batch, slen, num_key_value_heads * n_rep, head_dim) + @property + def _projection_dtype(self): + return self.qkv_proj_experts[0].weight.dtype + def forward( self, hidden_states: torch.Tensor, @@ -147,10 +130,10 @@ def forward( orig_shape: Optional[Tuple[int]] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if token_types is None: - raise ValueError("token_types can not be None") + raise ValueError("token_types must not be empty") if token_types.max() >= len(self.dim_inputs): raise ValueError( - f"token_types contains invalid expert indices: {token_types.max()}" + f"token_types contains an invalid expert index: {token_types.max()}" ) if self.config.mot_opt: @@ -195,24 +178,42 @@ def forward( key_states, value_states, self.layer_idx, cache_kwargs ) else: - past_key_states, past_value_states = past_key_value[self.layer_idx] + # Compatible across transformers versions: + # v5.x: DynamicCache uses .layers[idx].keys/.values + # v4.x: DynamicCache uses .key_cache[idx]/.value_cache[idx] + # old: Cache object is subscriptable, returns (key, value) tuple + if hasattr(past_key_value, "layers"): + past_key_states = past_key_value.layers[self.layer_idx].keys + past_value_states = past_key_value.layers[self.layer_idx].values + elif hasattr(past_key_value, "key_cache"): + past_key_states = past_key_value.key_cache[self.layer_idx] + past_value_states = past_key_value.value_cache[self.layer_idx] + else: + past_key_states, past_value_states = past_key_value[self.layer_idx] key_states = torch.cat([past_key_states, key_states], dim=-2) value_states = torch.cat([past_value_states, value_states], dim=-2) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) + target_dtype = self._projection_dtype + if ( + query_states.dtype != target_dtype + or key_states.dtype != target_dtype + or value_states.dtype != target_dtype + ): + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + causal_mask = attention_mask if attention_mask is not None: - # Ensure that the attention_mask correctly matches across the head dimension. if len(attention_mask.shape) == 2: # [batch_size, seq_len] - # Expanded to a causal mask format of [batch_size, 1, seq_len, seq_len] bsz, seq_len = attention_mask.shape causal_mask = attention_mask.view(bsz, 1, 1, seq_len).expand( bsz, 1, seq_len, seq_len ) elif len(attention_mask.shape) == 3: # [batch_size, seq_len, seq_len] - # add head dimension: [batch_size, 1, seq_len, seq_len] causal_mask = attention_mask.unsqueeze(1) elif ( len(attention_mask.shape) == 4 @@ -220,10 +221,9 @@ def forward( causal_mask = attention_mask else: raise ValueError( - f"Unsupported attention_mask dim: {attention_mask.shape}" + f"Unsupported attention_mask shape: {attention_mask.shape}" ) - # convert the attention mask to bool type causal_mask = causal_mask.to(torch.bool) # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, @@ -261,6 +261,8 @@ def forward( attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(bsz, q_len, -1) + if attn_output.dtype != target_dtype: + attn_output = attn_output.to(target_dtype) if self.config.mot_opt: output = self._generate_output_mot_opt( @@ -269,7 +271,32 @@ def forward( else: output = self._generate_output(attn_output, masks) - return output, None, past_key_value + attention_map = None + if output_attentions: + with torch.no_grad(): + action_token_num = int((token_types > 0).sum()) + action_query_states = query_states[:, :, -action_token_num:] + scale = 1.0 / torch.sqrt( + torch.tensor( + self.head_dim, device=hidden_states.device, dtype=torch.float32 + ) + ) + attention_score = ( + torch.matmul(action_query_states, key_states.transpose(-2, -1)) + * scale + ) + mask = causal_mask[:, :, -action_token_num:].expand( + -1, attention_score.shape[1], -1, -1 + ) # Mask only queries used for actions + if mask.dtype != attention_score.dtype: + mask = mask.to(dtype=attention_score.dtype) + attention_score = attention_score.masked_fill( + ~(mask.bool()), float("-inf") + ) + attention_score = torch.softmax(attention_score, dim=-1) + attention_map = attention_score[0].mean(0) + + return output, attention_map, past_key_value def _generate_qkv(self, hidden_states, masks): bsz, q_len, _ = hidden_states.size() @@ -300,29 +327,24 @@ def _generate_qkv(self, hidden_states, masks): ) # for expert_idx in range(len(self.dim_inputs)): - for expert_idx, (q_proj, k_proj, v_proj, mask) in enumerate( - zip(self.q_proj_experts, self.k_proj_experts, self.v_proj_experts, masks) + for expert_idx, (qkv_proj, mask) in enumerate( + zip(self.qkv_proj_experts, masks) ): if not mask.any(): continue dim_input = self.dim_inputs[expert_idx] selected_hidden = hidden_states[mask].clone() - - q_out = q_proj(selected_hidden[:, :dim_input]).view( - -1, self.num_heads, self.head_dim + if selected_hidden.dtype != qkv_proj.weight.dtype: + selected_hidden = selected_hidden.to(qkv_proj.weight.dtype) + qkv_out = qkv_proj(selected_hidden[:, :dim_input]).view( + -1, self.num_heads + 2 * self.num_key_value_heads, self.head_dim ) - k_out = k_proj(selected_hidden[:, :dim_input]).view( - -1, self.num_key_value_heads, self.head_dim - ) - v_out = v_proj(selected_hidden[:, :dim_input]).view( - -1, self.num_key_value_heads, self.head_dim - ) - - if self.config.model_type == "qwen3_vl_text": - q_out = self.q_norms[expert_idx](q_out)[0] - k_out = self.k_norms[expert_idx](k_out)[0] - + q_out = qkv_out[:, : self.num_heads, :] + k_out = qkv_out[ + :, self.num_heads : self.num_heads + self.num_key_value_heads, : + ] + v_out = qkv_out[:, self.num_heads + self.num_key_value_heads :, :] query_states[mask] = q_out key_states[mask] = k_out value_states[mask] = v_out @@ -377,26 +399,22 @@ def _generate_qkv_mot_opt( ) # === Each expert processes its own token slice === - for expert_idx, (q_proj, k_proj, v_proj) in enumerate( - zip(self.q_proj_experts, self.k_proj_experts, self.v_proj_experts) - ): + for expert_idx, qkv_proj in enumerate(self.qkv_proj_experts): start, end = start_indices[expert_idx], end_indices[expert_idx] if start == end: continue dim_input = self.dim_inputs[expert_idx] expert_input = hidden_states[start:end, :dim_input] + if expert_input.dtype != qkv_proj.weight.dtype: + expert_input = expert_input.to(qkv_proj.weight.dtype) # Compute Q/K/V - q_out = q_proj(expert_input) - k_out = k_proj(expert_input) - v_out = v_proj(expert_input) - - if getattr(self.config, "model_type", None) == "qwen3_vl_text": - q_out = self.q_norms[expert_idx](q_out) - q_out = q_out[0] if isinstance(q_out, (tuple, list)) else q_out - k_out = self.k_norms[expert_idx](k_out) - k_out = k_out[0] if isinstance(k_out, (tuple, list)) else k_out + qkv_out = qkv_proj(expert_input) + kv_dim = self.num_key_value_heads * self.head_dim + q_out, k_out, v_out = torch.split( + qkv_out, [self.num_heads * self.head_dim, kv_dim, kv_dim], dim=-1 + ) q_buffer[start:end] = q_out k_buffer[start:end] = k_out @@ -404,9 +422,9 @@ def _generate_qkv_mot_opt( # === Restore tokens to the original order === # unpermute (using the same unpermute operation) - q_unpermuted = ops.unpermute(q_buffer, row_id_map, probs) - k_unpermuted = ops.unpermute(k_buffer, row_id_map, probs) - v_unpermuted = ops.unpermute(v_buffer, row_id_map, probs) + q_unpermuted = unpermute(q_buffer, row_id_map, probs) + k_unpermuted = unpermute(k_buffer, row_id_map, probs) + v_unpermuted = unpermute(v_buffer, row_id_map, probs) # === Reshape to final form === query_states = q_unpermuted.view( @@ -424,19 +442,14 @@ def _generate_qkv_mot_opt( def _apply_rotary_pos_embed( self, query_states, key_states, cos, sin, unsqueeze_dim=1 ): - if self.config.model_type == "qwen2_5_vl": - query_states, key_states = apply_multimodal_rotary_pos_emb( - query_states.contiguous(), - key_states.contiguous(), - cos.contiguous(), - sin.contiguous(), - self.rope_scaling["mrope_section"], - unsqueeze_dim, - ) - else: - raise NotImplementedError( - f"Unsupported model type: {self.config.model_type}" - ) + del unsqueeze_dim + query_states, key_states = m_rope( + query_states.contiguous(), + key_states.contiguous(), + cos[..., : (cos.size(3) // 2)].contiguous().float(), + sin[..., : (sin.size(3) // 2)].contiguous().float(), + self.rope_scaling["mrope_section"], + ) return query_states, key_states def _generate_output(self, attn_output, masks): @@ -451,16 +464,18 @@ def _generate_output(self, attn_output, masks): continue dim_input = self.dim_inputs[expert_idx] - # Obtain all necessary indexes in a single operation. - mask_indices = mask.nonzero(as_tuple=False) + mask_indices = mask.nonzero( + as_tuple=False + ) # more efficient index retrieval if mask_indices.numel() == 0: continue batch_indices = mask_indices[:, 0] seq_indices = mask_indices[:, 1] - # Use advanced indexing directly to avoid intermediate tensors. selected_attn_output = attn_output[batch_indices, seq_indices] + if selected_attn_output.dtype != o_proj.weight.dtype: + selected_attn_output = selected_attn_output.to(o_proj.weight.dtype) projected_output = o_proj(selected_attn_output) output[batch_indices, seq_indices, :dim_input] = projected_output @@ -497,7 +512,7 @@ def _generate_output_mot_opt( # === 1. Flatten and reorder by expert assignment === flat_attn_output = attn_output.view(-1, hidden_dim) # [B*S, H] flat_expert_indices = experts_indices.reshape(-1) # [B*S] - permuted_inputs, _ = ops.permute(flat_attn_output, flat_expert_indices) + permuted_inputs, _ = permute(flat_attn_output, flat_expert_indices) total_tokens = permuted_inputs.shape[0] # === 2. Initialize output buffer (still in permuted token space) === @@ -513,6 +528,8 @@ def _generate_output_mot_opt( dim_input = self.dim_inputs[expert_idx] expert_input = permuted_inputs[start:end] # [N_e, dim_input] + if expert_input.dtype != o_proj.weight.dtype: + expert_input = expert_input.to(o_proj.weight.dtype) expert_output = o_proj(expert_input) # [N_e, hidden_dim] # Write results into the buffer (overwrite only valid dimension region) @@ -523,13 +540,14 @@ def _generate_output_mot_opt( class JointQwen2VLFlashAttention(JointQwen2VLAttention): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: Optional[int] = None): + super().__init__(config, layer_idx) # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + self.deterministic = config.attn_deterministic def forward( self, @@ -552,7 +570,10 @@ def forward( ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if token_types is None: - raise ValueError("token_types cannot be empty") + raise ValueError("token_types must not be empty") + # This check will lead to cudastreamsync. + # if token_types.max() >= len(self.dim_inputs): + # raise ValueError(f"token_types contains an invalid expert index: {token_types.max()}") if self.config.mot_opt: bsz, q_len, _ = orig_shape @@ -611,7 +632,7 @@ def forward( elif hasattr(self.config, "_pre_quantization_dtype"): target_dtype = self.config._pre_quantization_dtype else: - target_dtype = self.q_proj.weight.dtype + target_dtype = self._projection_dtype logger.warning_once( f"The input hidden states seems to be silently casted in float32, this might be related to" @@ -630,6 +651,7 @@ def forward( dropout_rate, softmax_scale=None, causal=self.is_causal, + deterministic=self.deterministic, ) attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() @@ -644,8 +666,265 @@ def forward( return output, None, past_key_value +class JointQwen2VLFlashMaskAttention(JointQwen2VLAttention): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + token_types: Optional[torch.LongTensor] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + probs: Optional[torch.Tensor] = None, + row_id_map: Optional[torch.Tensor] = None, + orig_shape: Optional[Tuple[int]] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if token_types is None: + raise ValueError("token_types must not be empty") + if token_types.max() >= len(self.dim_inputs): + raise ValueError( + f"token_types contains an invalid expert index: {token_types.max()}" + ) + + if self.config.mot_opt: + bsz, q_len, _ = orig_shape + query_states, key_states, value_states = self._generate_qkv_mot_opt( + hidden_states, + token_types, + start_indices, + end_indices, + probs, + row_id_map, + bsz, + q_len, + ) + else: + bsz, q_len, _ = hidden_states.size() + masks = [ + (token_types == expert_idx) + for expert_idx in range(len(self.dim_inputs)) + ] + query_states, key_states, value_states = self._generate_qkv( + hidden_states, masks + ) + + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + query_states, key_states = self._apply_rotary_pos_embed( + query_states, key_states, cos, sin, unsqueeze_dim=2 + ) + + if past_key_value is not None: + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = self.repeat_kv(key_states, self.num_key_value_groups) + value_states = self.repeat_kv(value_states, self.num_key_value_groups) + # dropout_rate = 0.0 if not self.training else self.attention_dropout + + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self._projection_dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + # Expand the attention_mask head dimension from 1 to num_heads + if attention_mask is not None and attention_mask.shape[1] == 1: + attention_mask = attention_mask.expand( + -1, self.num_heads, -1, -1 + ).contiguous() + + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + attn_output = flash_mask_attn_func( + query_states, + key_states, + value_states, + startend_row_indices=attention_mask, + causal=False, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + + if self.config.mot_opt: + output = self._generate_output_mot_opt( + attn_output, token_types, start_indices, end_indices + ) + else: + output = self._generate_output(attn_output, masks) + + return output, None, past_key_value + + +class JointQwen2VLFlashMaskAttention_KI(JointQwen2VLAttention): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + token_types: Optional[torch.LongTensor] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + probs: Optional[torch.Tensor] = None, + row_id_map: Optional[torch.Tensor] = None, + orig_shape: Optional[Tuple[int]] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if token_types is None: + raise ValueError("token_types must not be empty") + if token_types.max() >= len(self.dim_inputs): + raise ValueError( + f"token_types contains an invalid expert index: {token_types.max()}" + ) + + if self.config.mot_opt: + bsz, q_len, _ = orig_shape + query_states, key_states, value_states = self._generate_qkv_mot_opt( + hidden_states, + token_types, + start_indices, + end_indices, + probs, + row_id_map, + bsz, + q_len, + ) + else: + bsz, q_len, _ = hidden_states.size() + masks = [ + (token_types == expert_idx) + for expert_idx in range(len(self.dim_inputs)) + ] + query_states, key_states, value_states = self._generate_qkv( + hidden_states, masks + ) + + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + query_states, key_states = self._apply_rotary_pos_embed( + query_states, key_states, cos, sin, unsqueeze_dim=2 + ) + + if past_key_value is not None: + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + # repeat k/v heads if n_kv_heads < n_heads + key_states = self.repeat_kv(key_states, self.num_key_value_groups) + value_states = self.repeat_kv(value_states, self.num_key_value_groups) + # dropout_rate = 0.0 if not self.training else self.attention_dropout + + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self._projection_dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + # Expand the attention_mask head dimension from 1 to num_heads + if attention_mask is not None and attention_mask.shape[1] == 1: + attention_mask = attention_mask.expand( + -1, self.num_heads, -1, -1 + ).contiguous() + + # has_moe1_token = token_types.any(dim=1) # Check whether each row has nonzero values + # moe0_seq_len = (token_types != 0).int().argmax(dim=1) + # moe0_seq_len = torch.where(has_moe1_token, moe0_seq_len, q_len) # Set to q_len when there are no tokens + flow_mask = token_types == 1 + start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) + + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + + attn_output = flashmask_attn_func_stop_gradient( + query_states, + key_states, + value_states, + start_flow_pos, + startend_row_indices=attention_mask, + causal=False, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + + if self.config.mot_opt: + output = self._generate_output_mot_opt( + attn_output, token_types, start_indices, end_indices + ) + else: + output = self._generate_output(attn_output, masks) + + return output, None, past_key_value + + JOINT_QWEN_ATTENTION_CLASSES = { "eager": JointQwen2VLAttention, "flash_attention_2": JointQwen2VLFlashAttention, + # "flash_attention_2_ki": JointQwen2VLFlashAttention_KI, + # "flash_attention_2_triton": JointQwen2VLFlashAttention_Triton, "sdpa": JointQwen2VLAttention, + "flash_mask": JointQwen2VLFlashMaskAttention, + "flash_mask_ki": JointQwen2VLFlashMaskAttention_KI, } diff --git a/wall_x/model/core/attention/mask.py b/wall_x/model/core/attention/mask.py new file mode 100644 index 00000000..ec446083 --- /dev/null +++ b/wall_x/model/core/attention/mask.py @@ -0,0 +1,259 @@ +import torch + + +def find_first_last_ones(tensor): + """ + Input: tensor of shape (bs, seq_len) containing 0s and 1s + Output: (first_indices, last_indices), each of shape (bs,) + first_indices[i] is the first index of 1 in batch i, or -1 if none exists + last_indices[i] is the last index of 1 in batch i, or -1 if none exists + """ + bs, seq_len = tensor.shape + masks = tensor == 1 + has_ones = masks.any(dim=1) + + first = torch.full((bs,), -1, dtype=torch.long, device=tensor.device) + last = first.clone() + + # Compute the first index of 1 + first[has_ones] = torch.argmax(masks[has_ones].float(), dim=1) + + # Compute the last index of 1 + flipped_masks = masks.flip(dims=[1]) + last_argmax = torch.argmax(flipped_masks[has_ones].float(), dim=1) + last[has_ones] = seq_len - 1 - last_argmax + + return first, last + + +def update_position_ids(position_ids, moe_token_types, positional_masks): + """Extracted from ActionModelMixMin._update_position_ids (was @staticmethod).""" + if positional_masks is None or "ar_predict_token_positions" not in positional_masks: + return position_ids + + new_position_ids = position_ids.clone() + ar_predict_token_positions = positional_masks["ar_predict_token_positions"] + flow_mask = moe_token_types == 1 + + start_ar_pos, end_ar_pos = find_first_last_ones(ar_predict_token_positions) + start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) + + for bs_i in range(position_ids.shape[1]): + if start_ar_pos[bs_i] != -1 and end_ar_pos[bs_i] != -1: + start_ar_ids = new_position_ids[:, bs_i, start_ar_pos[bs_i]] + start_flow_ids = new_position_ids[:, bs_i, start_flow_pos[bs_i]] + diff = start_flow_ids - start_ar_ids + new_position_ids[:, bs_i, start_flow_pos[bs_i] :] = position_ids[ + :, bs_i, start_flow_pos[bs_i] : + ] - diff.unsqueeze(-1) + + return new_position_ids + + +def update_joint_attention_mask_2d( + attention_mask, + moe_token_types, + positional_masks, + causal_action_attention_mask=False, +): + """Extracted from ActionModelMixMin._update_joint_attention_mask_2d. + The only self attribute used was self.config.causal_action_attention_mask, now passed as parameter. + """ + if attention_mask.dim() == 3: # bs, seq_len, seq_len + return attention_mask + + bs, seq_len = moe_token_types.shape[0], moe_token_types.shape[1] + # Create a lower-triangular causal mask + causal_mask = torch.tril( + torch.ones( + (seq_len, seq_len), dtype=torch.bfloat16, device=moe_token_types.device + ) + ) + # Expand to the batch dimension + attention_mask = causal_mask.unsqueeze(0).expand(bs, -1, -1) + + if positional_masks is not None and "padding_positions" in positional_masks: + padding_positions = positional_masks["padding_positions"] + # Set padding rows to zero + attention_mask = torch.where( + padding_positions[:, None, :], + torch.zeros_like(attention_mask), + attention_mask, + ) + # Set padding columns to zero + attention_mask = torch.where( + padding_positions[:, :, None], + torch.zeros_like(attention_mask), + attention_mask, + ) + + # Set the moe1 region to 1 and mask it from the fast region + moe1_mask = (moe_token_types[:, :, None]) & (moe_token_types[:, None, :]) + + if ( + not causal_action_attention_mask + ): # If causal action attention mask is disabled, set the whole moe1 region to 1 + attention_mask = torch.where( + moe1_mask, torch.ones_like(attention_mask), attention_mask + ) + + if ( + positional_masks is not None + and "ar_predict_token_positions" in positional_masks + ): + ar_predict_token_positions = positional_masks["ar_predict_token_positions"] + moe1_mask = (moe_token_types[:, :, None]) & ( + ar_predict_token_positions[:, None, :] + ) + attention_mask = torch.where( + moe1_mask, torch.zeros_like(attention_mask), attention_mask + ) + + if ( + positional_masks is not None + and "valid_flow_action_positions" in positional_masks + ): + # true in moe_token_types but false in valid_flow_action_positions + nonvalid_flow_action_positions = ( + moe_token_types & ~positional_masks["valid_flow_action_positions"] + ) + attention_mask = torch.where( + nonvalid_flow_action_positions[:, None, :], + torch.zeros_like(attention_mask), + attention_mask, + ) + attention_mask = torch.where( + nonvalid_flow_action_positions[:, :, None], + torch.zeros_like(attention_mask), + attention_mask, + ) + + # AR and flow are bidirectional + if positional_masks is not None and "ar_action_mask" in positional_masks: + ar_action_mask = positional_masks["ar_action_mask"] != 0 + flow_positions = moe_token_types == 1 + if positional_masks.get("ar_visible", True): + flow_ar_position = ar_action_mask | flow_positions + flow_ar_mask = flow_ar_position[:, :, None] & flow_ar_position[:, None, :] + attention_mask = torch.where( + flow_ar_mask, torch.ones_like(attention_mask), attention_mask + ) + else: + flow_flow_mask = flow_positions[:, :, None] & flow_positions[:, None, :] + ar_ar_mask = ar_action_mask[:, :, None] & ar_action_mask[:, None, :] + flow_ar_mask = flow_flow_mask | ar_ar_mask + + affected = ar_action_mask | flow_positions # (B, N) + affected_pair = affected[:, :, None] & affected[:, None, :] + attention_mask = attention_mask.masked_fill(affected_pair, 0) + + attention_mask = torch.where( + flow_ar_mask, torch.ones_like(attention_mask), attention_mask + ) + + return attention_mask + + +def update_joint_attention_flash_mask( + attention_mask, + moe_token_types, + positional_masks, + causal_action_attention_mask=False, + debug=False, +): + """Extracted from ActionModelMixMin._update_joint_attention_flash_mask. + The only self attribute used was self.config.causal_action_attention_mask, now passed as parameter. + """ + device = moe_token_types.device + B, S = moe_token_types.shape + i32 = torch.int32 + + # ---- Return-vector initialization ---- + LTS = torch.ones((B, S), device=device, dtype=i32) * S + UTE = torch.arange(S, device=device, dtype=i32).unsqueeze(0).expand(B, S).clone() + + # Handle padding positions + if positional_masks is not None and "padding_positions" in positional_masks: + padding_positions = positional_masks["padding_positions"] + LTS[padding_positions] = 0 + UTE[padding_positions] = S + + # Handle AR predict tokens + if ( + positional_masks is not None + and "ar_predict_token_positions" in positional_masks + ): + start_ar_pos, end_ar_pos = find_first_last_ones( + positional_masks["ar_predict_token_positions"] + ) + for bs_i in range(B): + if end_ar_pos[bs_i] != -1: + LTS[bs_i, positional_masks["ar_predict_token_positions"][bs_i]] = ( + end_ar_pos[bs_i].to(i32) + 1 + ) + + flow_mask = moe_token_types == 1 + start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) + if positional_masks is None or "ar_action_mask" not in positional_masks: + # Handle bidirectional flow action masks + if not causal_action_attention_mask: + for bs_i in range(B): + if start_flow_pos[bs_i] != -1: + UTE[bs_i, flow_mask[bs_i]] = start_flow_pos[bs_i].to(i32) + else: + # AR and flow are bidirectional + ar_action_mask = positional_masks["ar_action_mask"] != 0 + flow_mask = moe_token_types == 1 + if positional_masks.get("ar_visible", True): + flow_ar_position = ar_action_mask | flow_mask + for bs_i in range(B): + idx = flow_ar_position[bs_i].nonzero(as_tuple=True)[0] + if idx.numel() == 0: + continue + block_start = idx.min() + block_end = idx.max() + 1 + # Set the visible range of every token in the block to [block_start, block_end) + UTE[bs_i, idx] = block_start.to(i32) + LTS[bs_i, idx] = block_end.to(i32) + else: + for bs_i in range(B): + # 1) Flow sub-block: flow attends only to flow + flow_idx = flow_mask[bs_i].nonzero(as_tuple=True)[0] + if flow_idx.numel() > 0: + flow_start = flow_idx.min() + flow_end = flow_idx.max() + 1 + # Only flow-token rows are set to [flow_start, flow_end) + UTE[bs_i, flow_idx] = flow_start.to(i32) + LTS[bs_i, flow_idx] = flow_end.to(i32) + + # 2) AR sub-block: AR attends only to AR + ar_idx = ar_action_mask[bs_i].nonzero(as_tuple=True)[0] + if ar_idx.numel() > 0: + ar_start = ar_idx.min() + ar_end = ar_idx.max() + 1 + # Only AR-token rows are set to [ar_start, ar_end) + UTE[bs_i, ar_idx] = ar_start.to(i32) + LTS[bs_i, ar_idx] = ar_end.to(i32) + + # Handle validation flow + if ( + positional_masks is not None + and "valid_flow_action_positions" in positional_masks + ): + flow_mask = moe_token_types == 1 + nonvalid_flow_action_positions = ( + flow_mask & ~positional_masks["valid_flow_action_positions"] + ) + if nonvalid_flow_action_positions.any(): + LTS[nonvalid_flow_action_positions] = 0 + UTE[nonvalid_flow_action_positions] = S + + LTS = LTS.unsqueeze(-1) + UTE = UTE.unsqueeze(-1) + + startend_row_indices = torch.cat([LTS, UTE], dim=-1) + + # add num_heads dimension + startend_row_indices = startend_row_indices.unsqueeze(1) + + return startend_row_indices diff --git a/wall_x/model/core/attention/selector.py b/wall_x/model/core/attention/selector.py new file mode 100644 index 00000000..6cda10c4 --- /dev/null +++ b/wall_x/model/core/attention/selector.py @@ -0,0 +1,138 @@ +from typing import Optional, Union, Dict + +from packaging import version + +from transformers.modeling_utils import AttentionInterface +from transformers.utils import logging, is_torch_xla_available + +import torch + +ALL_ATTENTION_FUNCTIONS: AttentionInterface = AttentionInterface() +logger = logging.get_logger(__name__) + +CUSTOM_ATTENTION_FUNCTIONS = [ + "flash_attention_2_ki", + "flash_attention_2_triton", + "flash_mask", + "flash_mask_ki", +] +ATTENTION_TYPES_WITH_2D_MASK = [ + "flash_attention_2_ki", + "flash_attention_2_triton", + "sdpa", +] +ATTENTION_TYPES_WITH_FLASH_MASK = ["flash_mask", "flash_mask_ki"] + + +class AttentionsSelectorMixin: + + @classmethod + def _autoset_attn_implementation( + cls, + config, + use_flash_attention_2: bool = False, + torch_dtype: Optional[torch.dtype] = None, + device_map: Optional[Union[str, Dict[str, int]]] = None, + check_device_map: bool = True, + ): + """ + Automatically checks and dispatches to a default attention implementation. In order of priority: + 1. An implementation specified in `config._attn_implementation` (due for example to the argument attn_implementation="sdpa" in from_pretrained). + 2. DEPRECATED: if use_flash_attention_2 is set to `True` and `flash_attn` is available, flash attention. (`LlamaFlashAttention` for example) + 3. SDPA implementation, if available and supported by the model type. (`LlamaSdpaAttention` for example) + 4. The default model's implementation otherwise (`LlamaAttention` for example) . + """ + # Here we use config._attn_implementation_internal to check whether the attention implementation was explicitly set by the user. + # The property `PretrainedConfig._attn_implementation` is never `None`, for backward compatibility (always fall back on "eager"). + # The `hasattr` here is used as some Transformers tests for some reason do not call PretrainedConfig __init__ (e.g. test_no_super_init_config_and_model) + requested_attn_implementation = None + if ( + hasattr(config, "_attn_implementation_internal") + and config._attn_implementation_internal is not None + ): + if ( + config._attn_implementation != "flash_attention_2" + and use_flash_attention_2 + ): + raise ValueError( + f'Both attn_implementation="{config._attn_implementation}" and `use_flash_attention_2=True` were used when loading the model, which are not compatible.' + ' We recommend to just use `attn_implementation="flash_attention_2"` when loading the model.' + ) + + if ( + not isinstance(config._attn_implementation, dict) + and config._attn_implementation + not in ["eager"] + + ALL_ATTENTION_FUNCTIONS.valid_keys() + + CUSTOM_ATTENTION_FUNCTIONS + ): + message = f'Specified `attn_implementation="{config._attn_implementation}"` is not supported. The only possible arguments are `attn_implementation="eager"` (manual attention implementation)' + if cls._supports_flash_attn_2: + message += ', `"attn_implementation=flash_attention_2"` (implementation using flash attention 2)' + if cls._supports_sdpa: + message += ', `"attn_implementation=sdpa"` (implementation using torch.nn.functional.scaled_dot_product_attention)' + if cls._supports_flex_attn: + message += ', `"attn_implementation=flex_attention"` (implementation using torch\'s flex_attention)' + raise ValueError(message + ".") + + # If a config is passed with a preset attn_implementation, we skip the automatic dispatch and use the user-provided config, with hard checks that the requested attention implementation is available. + requested_attn_implementation = config._attn_implementation_internal + + if use_flash_attention_2: + logger.warning_once( + 'The model was loaded with use_flash_attention_2=True, which is deprecated and may be removed in a future release. Please use `attn_implementation="flash_attention_2"` instead.' + ) + config._attn_implementation = "flash_attention_2" + + if config._attn_implementation == "flash_attention_2": + cls._check_and_enable_flash_attn_2( + config, + torch_dtype=torch_dtype, + device_map=device_map, + hard_check_only=False, + check_device_map=check_device_map, + ) + elif requested_attn_implementation == "flex_attention": + config = cls._check_and_enable_flex_attn(config, hard_check_only=True) + elif ( + requested_attn_implementation in [None, "sdpa"] + and not is_torch_xla_available() + ): + # use_flash_attention_2 takes priority over SDPA, hence SDPA treated in this elif. + config = cls._check_and_enable_sdpa( + config, + hard_check_only=( + False if requested_attn_implementation is None else True + ), + ) + + if ( + torch.version.hip is not None + and config._attn_implementation == "sdpa" + and torch.cuda.device_count() > 1 + and version.parse(torch.__version__) < version.parse("2.4.1") + ): + logger.warning_once( + "Using the `SDPA` attention implementation on multi-gpu setup with ROCM may lead to performance issues due to the FA backend. Disabling it to use alternative backends." + ) + torch.backends.cuda.enable_flash_sdp(False) + elif requested_attn_implementation in ALL_ATTENTION_FUNCTIONS.valid_keys(): + config._attn_implementation = requested_attn_implementation + elif isinstance(requested_attn_implementation, dict): + config._attn_implementation = None + elif config._attn_implementation in CUSTOM_ATTENTION_FUNCTIONS: + pass + else: + config._attn_implementation = "eager" + + config._attn_implementation_autoset = True + return config + + def _check_and_adjust_attn_implementation( + self, attn_implementation: Optional[str], is_init_check: bool = False + ) -> str: + assert ( + attn_implementation + in ["eager", "flash_attention_2", "sdpa"] + CUSTOM_ATTENTION_FUNCTIONS + ) + return attn_implementation diff --git a/wall_x/model/core/ops/__init__.py b/wall_x/model/core/ops/__init__.py new file mode 100644 index 00000000..8fe04a52 --- /dev/null +++ b/wall_x/model/core/ops/__init__.py @@ -0,0 +1,17 @@ +"""Operator proxy layer with runtime fallback backends.""" + +from wall_x.model.core.ops.norm import rmsnorm +from wall_x.model.core.ops.rope import rope, m_rope, rot_pos_emb +from wall_x.model.core.ops.moe import permute, unpermute +from wall_x.model.core.ops.index import get_rope_index, get_window_index + +__all__ = [ + "rmsnorm", + "rope", + "m_rope", + "rot_pos_emb", + "permute", + "unpermute", + "get_rope_index", + "get_window_index", +] diff --git a/wall_x/model/core/ops/base.py b/wall_x/model/core/ops/base.py new file mode 100644 index 00000000..b0a78625 --- /dev/null +++ b/wall_x/model/core/ops/base.py @@ -0,0 +1,99 @@ +"""Base class for operator proxies with lazy backend resolution.""" + +import logging +import threading + +logger = logging.getLogger(__name__) + + +class OpsProxy: + """Callable proxy that lazily resolves to the best available backend. + + Public fallback order: + Level 1: CUDA inline kernel, when available + Level 2: Pure PyTorch fallback + """ + + def __init__(self): + self._resolved_fn = None + self._backend = None + self._resolve_lock = threading.Lock() + + def _resolve(self): + """Resolve the best available backend. Called once on first use.""" + if self._resolved_fn is not None: + return + with self._resolve_lock: + if self._resolved_fn is not None: + return + + cuda_fn = self._get_cuda_kernel() + if cuda_fn is not None: + self._backend = "cuda_inline" + self._resolved_fn = cuda_fn + return + + pytorch_fn = self._get_pytorch_fallback() + if pytorch_fn is None: + raise RuntimeError( + f"{self.__class__.__name__}: no available operator backend. " + f"Available: {self.available_backends()}" + ) + self._backend = "pytorch" + self._resolved_fn = pytorch_fn + logger.info("%s: using PyTorch fallback", self.__class__.__name__) + + def _get_cuda_kernel(self): + """Override in subclasses to provide a CUDA inline kernel.""" + return None + + def _pytorch_fallback(self, *args, **kwargs): + raise NotImplementedError(f"{self.__class__.__name__} has no PyTorch fallback") + + def _get_pytorch_fallback(self): + """Return the PyTorch fallback only when a subclass implements one.""" + if type(self)._pytorch_fallback is OpsProxy._pytorch_fallback: + return None + return self._pytorch_fallback + + def __call__(self, *args, **kwargs): + if self._resolved_fn is None: + self._resolve() + return self._resolved_fn(*args, **kwargs) + + @property + def backend(self) -> str: + if self._resolved_fn is None: + self._resolve() + return self._backend + + def available_backends(self): + """Return available backend names for this operator.""" + backends = [] + if self._get_cuda_kernel() is not None: + backends.append("cuda_inline") + if self._get_pytorch_fallback() is not None: + backends.append("pytorch") + return backends + + def _get_backend_fn(self, backend): + """Return the callable for a specific backend, or None if unavailable.""" + if backend == "pytorch": + return self._get_pytorch_fallback() + if backend == "cuda_inline": + return self._get_cuda_kernel() + raise ValueError(f"Unknown backend: {backend}") + + def call_with_backend(self, backend, *args, **kwargs): + """Call this operator using a specific backend.""" + fn = self._get_backend_fn(backend) + if fn is None: + raise RuntimeError( + f"{self.__class__.__name__}: backend '{backend}' not available. " + f"Available: {self.available_backends()}" + ) + return fn(*args, **kwargs) + + def __repr__(self): + backend = self._backend or "unresolved" + return f"<{self.__class__.__name__} backend={backend}>" diff --git a/wall_x/model/core/ops/index.py b/wall_x/model/core/ops/index.py new file mode 100644 index 00000000..fa5e540a --- /dev/null +++ b/wall_x/model/core/ops/index.py @@ -0,0 +1,339 @@ +"""Index computation operators for attention.""" + +import logging +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from wall_x.model.core.ops.base import OpsProxy + +logger = logging.getLogger(__name__) + +# Sentinel value for padding in window index computation. +# Valid indices are always >= 0 (they are positional offsets within flattened +# grid tensors), so any negative value is safe. -100 is chosen by convention +# (same as HuggingFace's ignore_index for cross-entropy loss). +_PAD_VALUE = -100 + + +def _compute_vision_position_ids( + input_tokens, + token_set, + st, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + image_index, + video_index, + remain_images, + remain_videos, + image_token_id, + video_token_id, + spatial_merge_size, + tokens_per_second, + device, + llm_pos_ids_list, +): + """Compute position IDs for a single vision token (image or video). + + Returns updated (st, image_index, video_index, remain_images, remain_videos). + """ + if image_token_id in token_set and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in token_set and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + + if ed_image < ed_video: + if image_index >= len(image_grid_thw): + raise IndexError( + f"image_index {image_index} out of range (have {len(image_grid_thw)} image grids)" + ) + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + # Images are single-frame: no temporal progression, so second_per_grid_t = 0. + # This makes all image tokens share t_index = 0 (only spatial positions vary). + second_per_grid_t = 0 + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + if video_index >= len(video_grid_thw): + raise IndexError( + f"video_index {video_index} out of range (have {len(video_grid_thw)} video grids)" + ) + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + if second_per_grid_ts is not None: + second_per_grid_t = second_per_grid_ts[video_index] + else: + second_per_grid_t = 1.0 + video_index += 1 + remain_videos -= 1 + ed = ed_video + + llm_grid_t, llm_grid_h, llm_grid_w = ( + int(t), + int(h) // spatial_merge_size, + int(w) // spatial_merge_size, + ) + text_len = ed - st + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append( + torch.arange(text_len, device=device).view(1, -1).expand(3, -1) + st_idx + ) + range_tensor = torch.arange(llm_grid_t, device=device).view(-1, 1) + expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) + time_tensor = expanded_range * second_per_grid_t * tokens_per_second + t_index = time_tensor.long().flatten() + h_index = ( + torch.arange(llm_grid_h, device=device) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w, device=device) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx + ) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + return st, image_index, video_index, remain_images, remain_videos + + +def _compute_text_only_positions(input_ids, attention_mask): + """Compute position IDs for text-only inputs (no vision tokens).""" + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = ( + position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device) + ) + max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[ + 0 + ] + mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] + else: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + return position_ids, mrope_position_deltas + + +class GetRopeIndexOp(OpsProxy): + """Compute 3D RoPE position indices for multimodal inputs. + + Signature: get_rope_index(input_ids, image_grid_thw, video_grid_thw, + second_per_grid_ts, attention_mask, spatial_merge_size, + image_token_id, video_token_id, vision_start_token_id, + tokens_per_second) -> (position_ids, mrope_position_deltas) + """ + + @property + def _external_accel_name(self): + return "get_rope_index" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import GetRopeIndex + + return GetRopeIndex() + except ImportError: + return None + except Exception as e: + logger.warning("GetRopeIndexOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback( + self, + input_ids: torch.LongTensor, + image_grid_thw: Optional[torch.LongTensor], + video_grid_thw: Optional[torch.LongTensor], + second_per_grid_ts: Optional[torch.Tensor], + attention_mask: Optional[torch.Tensor], + spatial_merge_size: int, + image_token_id: int, + video_token_id: int, + vision_start_token_id: int, + tokens_per_second: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if input_ids is None or (image_grid_thw is None and video_grid_thw is None): + return _compute_text_only_positions(input_ids, attention_mask) + + mrope_position_deltas = [] + device = input_ids.device + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + for i, input_ids_i in enumerate(total_input_ids): + input_ids_i = input_ids_i[attention_mask[i] == 1] + vision_start_indices = torch.argwhere( + input_ids_i == vision_start_token_id + ).squeeze(1) + # Boundary check: ensure vision_start + 1 doesn't exceed sequence length + valid_mask = (vision_start_indices + 1) < len(input_ids_i) + vision_start_indices = vision_start_indices[valid_mask] + vision_tokens = input_ids_i[vision_start_indices + 1] + image_nums = int((vision_tokens == image_token_id).sum()) + video_nums = int((vision_tokens == video_token_id).sum()) + input_tokens = input_ids_i.tolist() + token_set = set(input_tokens) + llm_pos_ids_list: list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + st, image_index, video_index, remain_images, remain_videos = ( + _compute_vision_position_ids( + input_tokens, + token_set, + st, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + image_index, + video_index, + remain_images, + remain_videos, + image_token_id, + video_token_id, + spatial_merge_size, + tokens_per_second, + device, + llm_pos_ids_list, + ) + ) + if st < len(input_tokens): + st_idx = ( + llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + ) + text_len = len(input_tokens) - st + llm_pos_ids_list.append( + torch.arange(text_len, device=device).view(1, -1).expand(3, -1) + + st_idx + ) + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( + position_ids.device + ) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]) + ) + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=total_input_ids.device + ).unsqueeze(1) + return position_ids, mrope_position_deltas + + +class GetWindowIndexOp(OpsProxy): + """Compute window attention indices for ViT. + + Signature: get_window_index(grid_thw, window_size, spatial_merge_size, + patch_size, spatial_merge_unit=1) -> (window_index, cu_window_seqlens) + """ + + @property + def _external_accel_name(self): + return "get_window_index" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import get_window_index_cuda + + return get_window_index_cuda + except ImportError: + return None + except Exception as e: + logger.warning("GetWindowIndexOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback( + self, + grid_thw, + window_size, + spatial_merge_size, + patch_size, + spatial_merge_unit=1, + ): + device = grid_thw.device + vit_merger_window_size = window_size // spatial_merge_size // patch_size + window_index_list = [] + cu_window_seqlens = [0] + window_index_id = 0 + for grid_t, grid_h, grid_w in grid_thw: + llm_grid_h = grid_h // spatial_merge_size + llm_grid_w = grid_w // spatial_merge_size + index = torch.arange( + grid_t * llm_grid_h * llm_grid_w, device=device + ).reshape(grid_t, llm_grid_h, llm_grid_w) + pad_h = ( + vit_merger_window_size - llm_grid_h % vit_merger_window_size + ) % vit_merger_window_size + pad_w = ( + vit_merger_window_size - llm_grid_w % vit_merger_window_size + ) % vit_merger_window_size + num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size + num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size + index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", _PAD_VALUE) + index_padded = ( + index_padded.reshape( + grid_t, + num_windows_h, + vit_merger_window_size, + num_windows_w, + vit_merger_window_size, + ) + .permute(0, 1, 3, 2, 4) + .reshape( + grid_t, + num_windows_h * num_windows_w, + vit_merger_window_size, + vit_merger_window_size, + ) + ) + seqlens = (index_padded != _PAD_VALUE).sum([2, 3]).reshape(-1) + index_padded = index_padded.reshape(-1) + index_new = index_padded[index_padded != _PAD_VALUE] + window_index_list.append(index_new + window_index_id) + cu_seqlens_tmp = ( + seqlens.cumsum(0) * spatial_merge_unit + cu_window_seqlens[-1] + ) + cu_window_seqlens.extend(cu_seqlens_tmp.tolist()) + window_index_id += int(grid_t) * int(llm_grid_h) * int(llm_grid_w) + window_index = torch.cat(window_index_list, dim=0) + cu_window_seqlens = torch.tensor( + cu_window_seqlens, dtype=grid_thw.dtype, device=grid_thw.device + ) + return window_index, cu_window_seqlens + + +get_rope_index = GetRopeIndexOp() +get_window_index = GetWindowIndexOp() diff --git a/wall_x/model/core/ops/moe.py b/wall_x/model/core/ops/moe.py new file mode 100644 index 00000000..f10298dd --- /dev/null +++ b/wall_x/model/core/ops/moe.py @@ -0,0 +1,96 @@ +"""MoE routing operators (permute/unpermute). + +These operators reorder tokens by expert assignment for efficient +Mixture-of-Experts processing. +""" + +import logging + +import torch +from wall_x.model.core.ops.base import OpsProxy + +logger = logging.getLogger(__name__) + + +class PermuteOp(OpsProxy): + """Reorder tokens by expert assignment for MoE processing. + + Signature: permute(tokens, indices, num_out_tokens=None, max_token_num=0) -> (permuted_tokens, sorted_indices) + """ + + @property + def _external_accel_name(self): + return "permute" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import permute_kernel + + return permute_kernel + except ImportError: + return None + except Exception as e: + logger.warning("PermuteOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback(self, tokens, indices, num_out_tokens=None, max_token_num=0): + """PyTorch fallback for permute. + + Args: + num_out_tokens: If set, truncate output to this many tokens (matches + external_accel behavior of discarding padding expert assignments). + max_token_num: Unused, kept for external_accel API compatibility. + """ + del max_token_num # unused, external_accel API compat + if indices.dim() == 1: + indices = indices.view(-1, 1) + expand_factor = indices.size(1) + flatten_indices = indices.view(-1) + # Keep int64 throughout to avoid precision loss on large token counts + sorted_indices = torch.argsort(flatten_indices, stable=True) + permuted_tokens = tokens.index_select(0, sorted_indices // expand_factor) + if num_out_tokens is not None: + permuted_tokens = permuted_tokens[:num_out_tokens] + sorted_indices = sorted_indices[:num_out_tokens] + return permuted_tokens, sorted_indices + + +class UnpermuteOp(OpsProxy): + """Restore tokens to original order after MoE processing. + + Signature: unpermute(permuted_tokens, sorted_indices, probs=None) -> restored_tokens + """ + + @property + def _external_accel_name(self): + return "unpermute" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import unpermute_kernel + + return unpermute_kernel + except ImportError: + return None + except Exception as e: + logger.warning("UnpermuteOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback(self, permuted_tokens, sorted_indices, probs=None): + if probs is not None: + merge_factor = probs.size(1) + else: + merge_factor = 1 + unpermuted_tokens = torch.zeros_like(permuted_tokens) + unpermuted_tokens.index_copy_(0, sorted_indices.long(), permuted_tokens) + unpermuted_tokens = unpermuted_tokens.reshape( + -1, merge_factor, permuted_tokens.size(-1) + ) + if probs is not None: + unpermuted_tokens = unpermuted_tokens * probs.unsqueeze(-1) + unpermuted_tokens = unpermuted_tokens.sum(dim=1) + return unpermuted_tokens + + +permute = PermuteOp() +unpermute = UnpermuteOp() diff --git a/wall_x/model/core/ops/norm.py b/wall_x/model/core/ops/norm.py new file mode 100644 index 00000000..dfea59a8 --- /dev/null +++ b/wall_x/model/core/ops/norm.py @@ -0,0 +1,41 @@ +"""RMS normalization operator.""" + +import logging + +import torch +from wall_x.model.core.ops.base import OpsProxy + +logger = logging.getLogger(__name__) + + +class RMSNormOp(OpsProxy): + """RMS normalization: rmsnorm(x, weight, eps). + + Accepts ``rmsnorm(x, weight, eps)``. + """ + + @property + def _external_accel_name(self): + return "rmsnorm" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import Rmsnorm + + return Rmsnorm() + except ImportError: + return None + except Exception as e: + logger.warning("RMSNormOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback(self, hidden_states, weight, eps=1e-6): + """Pure PyTorch RMSNorm. eps default matches external_accel (1e-6, not PyTorch's 1e-5).""" + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + eps) + return (weight * hidden_states).to(input_dtype) + + +rmsnorm = RMSNormOp() diff --git a/wall_x/model/core/ops/rope.py b/wall_x/model/core/ops/rope.py new file mode 100644 index 00000000..a1dc24a8 --- /dev/null +++ b/wall_x/model/core/ops/rope.py @@ -0,0 +1,233 @@ +"""Rotary position embedding operators.""" + +import logging +from typing import List + +import torch +from wall_x.model.core.ops.base import OpsProxy + +logger = logging.getLogger(__name__) + + +def _rotate_half(x): + """Rotate half: split last dim in two halves, negate-swap, concat.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _rotate_interleave(x): + """Interleaved rotation: pairs of (even, odd) elements.""" + x_even = x[..., ::2] + x_odd = x[..., 1::2] + return torch.stack((-x_odd, x_even), dim=-1).flatten(-2) + + +class RoPEOp(OpsProxy): + """Standard rotary position embedding. + + Signature: rope(q, k, cos, sin, interleave=False) -> (q_embed, k_embed) + """ + + @property + def _external_accel_name(self): + return "rope" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import Rope + + return Rope() + except ImportError: + return None + except Exception as e: + logger.warning("RoPEOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback(self, q, k, cos, sin, interleave=False, **kwargs): + cos = cos.float() + sin = sin.float() + rotary_dim = cos.size(-1) * 2 + head_dim = q.size(-1) + if rotary_dim > head_dim: + raise ValueError( + f"rotary_dim ({rotary_dim}) > head_dim ({head_dim}): " + f"cos last dim ({cos.size(-1)}) is too large for q" + ) + partial = rotary_dim < head_dim + if partial: + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + else: + q_rot, k_rot = q, k + cos = cos.unsqueeze(-2) # (..., 1, half_dim) + sin = sin.unsqueeze(-2) + if interleave: + cos = cos.repeat_interleave(2, dim=-1) + sin = sin.repeat_interleave(2, dim=-1) + q_embed = q_rot.float() * cos + _rotate_interleave(q_rot.float()) * sin + k_embed = k_rot.float() * cos + _rotate_interleave(k_rot.float()) * sin + else: + cos = torch.cat((cos, cos), dim=-1) + sin = torch.cat((sin, sin), dim=-1) + q_embed = q_rot.float() * cos + _rotate_half(q_rot.float()) * sin + k_embed = k_rot.float() * cos + _rotate_half(k_rot.float()) * sin + if partial: + q_embed = torch.cat([q_embed.to(q.dtype), q_pass], dim=-1) + k_embed = torch.cat([k_embed.to(k.dtype), k_pass], dim=-1) + else: + q_embed = q_embed.to(q.dtype) + k_embed = k_embed.to(k.dtype) + return q_embed, k_embed + + def pack(self, *args, **kwargs): + """Delegate to resolved backend's pack method if available.""" + if self._resolved_fn is None: + self._resolve() + if hasattr(self._resolved_fn, "pack"): + return self._resolved_fn.pack(*args, **kwargs) + return args + + +class MRoPEOp(OpsProxy): + """Multi-head rotary position embedding (used by Qwen2.5-VL models). + + Signature: m_rope(query_states, key_states, cos, sin, mrope_section, interleaved=False) -> (q_embed, k_embed) + + cos/sin shape: (3, B, S, D//2). mrope_section like [16, 24, 24] specifies + how many head dims each of T/H/W occupies. After split by mrope_section * 2, + uses i % 3 to select from corresponding temporal/height/width rows. + """ + + @property + def _external_accel_name(self): + return "m_rope" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import MRope + + return MRope() + except ImportError: + return None + except Exception as e: + logger.warning("MRoPEOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback( + self, + query_states, + key_states, + cos, + sin, + mrope_section: List[int], + interleaved=False, + **kwargs, + ): + cos = cos.float() + sin = sin.float() + # Double cos/sin along last dim (matching external_accel kernel's internal behavior) + cos = torch.cat((cos, cos), dim=-1) + sin = torch.cat((sin, sin), dim=-1) + # List concatenation: [16,24,24] -> [16,24,24,16,24,24] (NOT element-wise multiply) + mrope_section_doubled = mrope_section + mrope_section + cos_split = torch.cat( + [m[i % 3] for i, m in enumerate(cos.split(mrope_section_doubled, dim=-1))], + dim=-1, + ).unsqueeze( + 2 + ) # (B, S, 1, D) + sin_split = torch.cat( + [m[i % 3] for i, m in enumerate(sin.split(mrope_section_doubled, dim=-1))], + dim=-1, + ).unsqueeze(2) + q_embed = (query_states.float() * cos_split) + ( + _rotate_half(query_states.float()) * sin_split + ) + k_embed = (key_states.float() * cos_split) + ( + _rotate_half(key_states.float()) * sin_split + ) + return q_embed.to(query_states.dtype), k_embed.to(key_states.dtype) + + def pack(self, *args, **kwargs): + """Delegate to resolved backend's pack method if available.""" + if self._resolved_fn is None: + self._resolve() + if hasattr(self._resolved_fn, "pack"): + return self._resolved_fn.pack(*args, **kwargs) + return args + + +class RotPosEmbOp(OpsProxy): + """Rotary position embedding computation for ViT (used by Qwen2.5-VL vision encoder). + + Signature: rot_pos_emb(inv_freq, grid_thw, spatial_merge_size) -> rotary_pos_emb + """ + + @property + def _external_accel_name(self): + return "rot_pos_emb" + + def _get_cuda_kernel(self): + try: + from wall_x.model.core.ops._cuda_wrappers import RotPos + + return RotPos() + except ImportError: + return None + except Exception as e: + logger.warning("RotPosEmbOp: CUDA kernel load failed: %s", e) + return None + + def _pytorch_fallback(self, inv_freq, grid_thw, spatial_merge_size): + if inv_freq.dtype != torch.float32: + inv_freq = inv_freq.to(torch.float32) + pos_ids = [] + for t, h, w in grid_thw: + t, h, w = int(t), int(h), int(w) + if h % spatial_merge_size != 0 or w % spatial_merge_size != 0: + raise ValueError( + f"grid h={h}, w={w} must be divisible by spatial_merge_size={spatial_merge_size}" + ) + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + hpos_ids = ( + hpos_ids.reshape( + h // spatial_merge_size, + spatial_merge_size, + w // spatial_merge_size, + spatial_merge_size, + ) + .permute(0, 2, 1, 3) + .flatten() + ) + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) + wpos_ids = ( + wpos_ids.reshape( + h // spatial_merge_size, + spatial_merge_size, + w // spatial_merge_size, + spatial_merge_size, + ) + .permute(0, 2, 1, 3) + .flatten() + ) + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + pos_ids = torch.cat(pos_ids, dim=0).to(inv_freq.device) + max_grid_size = grid_thw[:, 1:].max() + seq = torch.arange(max_grid_size, device=inv_freq.device, dtype=inv_freq.dtype) + rotary_pos_emb_full = torch.outer(seq, inv_freq) + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) + return rotary_pos_emb.to(torch.float) + + def pack(self, *args, **kwargs): + """Delegate to resolved backend's pack method if available.""" + if self._resolved_fn is None: + self._resolve() + if hasattr(self._resolved_fn, "pack"): + return self._resolved_fn.pack(*args, **kwargs) + return args + + +rope = RoPEOp() +m_rope = MRoPEOp() +rot_pos_emb = RotPosEmbOp() diff --git a/wall_x/model/core/vla_mixin.py b/wall_x/model/core/vla_mixin.py new file mode 100644 index 00000000..d5aff23c --- /dev/null +++ b/wall_x/model/core/vla_mixin.py @@ -0,0 +1,751 @@ +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.utils.checkpoint as cp + +from torch.distributed.fsdp import MixedPrecision as MP +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + + +from peft import LoraConfig, get_peft_model + +from transformers import GenerationMixin +from transformers.utils import logging + +from wall_x.utils.constant import is_action_dataset_name + +from wall_x.model.core.action.processor import ActionProcessor +from wall_x.model.core.action.moe import TokenTypeRouter, SparseMoeBlock +from wall_x.model.core.attention.mask import ( + update_position_ids, + update_joint_attention_mask_2d, + update_joint_attention_flash_mask, +) + +logger = logging.get_logger(__name__) + + +class ActionModelMixMin: + # config: Qwen2_5_VLConfig + action_preprocessor: ActionProcessor + router: "TokenTypeRouter" + moe: "SparseMoeBlock" + + def __init__(self, config, action_preprocessor, router, moe): + self.config = config + self.action_preprocessor = action_preprocessor + self.router = router + self.moe = moe + self._mot_opt_warned = False + + def set_normalizer(self, normalizer_action, normalizer_propri): + if hasattr(self, "action_preprocessor"): + self.action_preprocessor.set_normalizer( + normalizer_action, normalizer_propri + ) + else: + # WARNING: normalizer cannot be set when action_preprocessor is missing + logger.warning( + "ActionModelMixMin.set_normalizer is called but action_preprocessor is not set" + ) + + def _apply_mlp_moe(self, hidden_states, token_types, start_indices, end_indices): + if self.config.mlp_moe: + hidden_states = self.moe( + hidden_states, token_types, start_indices, end_indices + ) + else: + hidden_states = self.mlp(hidden_states) + return hidden_states + + def _apply_norm_moe( + self, + hidden_states, + token_types, + adarms_conds, + norms, # list of norm layers (expert-wise) + norm, # shared norm if not norm_moe + start_indices=None, + end_indices=None, + use_selective_recompute=False, + ): + """ + MoE-aware LayerNorm with optional selective activation recomputation. + + Only activation math is recomputed. No GEMM is recomputed. + Safe for FSDP (use_reentrant=False). + """ + + gate = None + gate_mask = None + + # ------------------------- + # Case 1: norm_moe=True (expert-wise norm) + # ------------------------- + if self.config.norm_moe: + + # --------------------------------------------------------- + # Case 1A: mot_opt=True (segments assigned by start/end) + # --------------------------------------------------------- + if self.config.mot_opt: + new_hidden_states = torch.zeros_like(hidden_states) + + for expert_idx, expert_norm in enumerate(norms): + start = start_indices[expert_idx] + end = end_indices[expert_idx] + if start == end: + continue + + dim_input = self.config.dim_inputs[expert_idx] + selected = hidden_states[start:end] # [K, D] + + # ====== reshape if adarms on flow expert ====== + if self.config.use_adarms and expert_idx == 1: + selected = selected.view( + -1, + self.config.action_horizon_flow, + selected.shape[-1], + ) + input_slice = selected[:, :, :dim_input] + cond = adarms_conds[expert_idx] + else: + input_slice = selected[:, :dim_input] + cond = adarms_conds[expert_idx] + + if use_selective_recompute: + + def norm_chunk(t_x, t_cond, expert_norm=expert_norm): + if t_cond is None or ( + isinstance(t_cond, torch.Tensor) and t_cond.numel() == 0 + ): + out, _ = expert_norm(t_x) + else: + out, _ = expert_norm(t_x, t_cond) + return out + + cond_for_cp = ( + cond + if cond is not None + else torch.empty(0, device=input_slice.device) + ) + processed = cp.checkpoint( + norm_chunk, + input_slice, + cond_for_cp, + use_reentrant=False, + ) + else: + processed, gate = expert_norm(input_slice, cond) + + # reshape back if needed + if self.config.use_adarms and expert_idx == 1: + processed = processed.view(-1, dim_input) + + new_hidden_states[start:end, :dim_input] = processed.to( + hidden_states.dtype + ) + + hidden_states = new_hidden_states + + # --------------------------------------------------------- + # Case 1B: mot_opt=False (token-level mask) + # --------------------------------------------------------- + else: + + new_hidden_states = torch.zeros_like(hidden_states) + B, S, D = hidden_states.shape + + for expert_idx, expert_norm in enumerate(norms): + mask = token_types == expert_idx + if mask.sum() == 0: + continue + + dim_input = self.config.dim_inputs[expert_idx] + selected = hidden_states[mask] # [K, D] + + if self.config.use_adarms and expert_idx == 1: + gate_mask = mask + selected = selected.view( + -1, + self.config.action_horizon_flow, + selected.shape[-1], + ) + input_slice = selected[:, :, :dim_input] + cond = adarms_conds[expert_idx] + else: + input_slice = selected[:, :dim_input] + cond = adarms_conds[expert_idx] + + if use_selective_recompute: + + def norm_chunk(t_x, t_cond, expert_norm=expert_norm): + if t_cond is None or ( + isinstance(t_cond, torch.Tensor) and t_cond.numel() == 0 + ): + out, _ = expert_norm(t_x) + else: + out, _ = expert_norm(t_x, t_cond) + return out + + cond_for_cp = ( + cond + if cond is not None + else torch.empty(0, device=input_slice.device) + ) + + processed = cp.checkpoint( + norm_chunk, + input_slice, + cond_for_cp, + use_reentrant=False, + ) + else: + processed, gate = expert_norm(input_slice, cond) + + if self.config.use_adarms and expert_idx == 1: + processed = processed.view(-1, dim_input) + + # scatter back + b_id, s_id = torch.where(mask) + new_hidden_states[b_id, s_id, :dim_input] = processed.to( + hidden_states.dtype + ) + + hidden_states = new_hidden_states + + # ------------------------- + # Case 2: norm_moe=False (single LN) + # ------------------------- + else: + + def norm_chunk_shared(t_x, dummy, norm_module=norm): + out, _ = norm_module(t_x) + return out + + if use_selective_recompute: + dummy = torch.empty(0, device=hidden_states.device) + hidden_states = cp.checkpoint( + norm_chunk_shared, + hidden_states, + dummy, + use_reentrant=False, + ) + else: + hidden_states, gate = norm(hidden_states) + + return hidden_states, gate, gate_mask + + def _gated_residual(self, x, y, gate, start_indices=None, end_indices=None): + """ + Applies gated residual connection with optional gate parameter. + + Args: + x: Input tensor (residual) + y: Output tensor to be added + gate: Optional gate tensor to modulate the addition + + Returns: + x + y if gate is None, otherwise x + y * gate + """ + if x is None and y is None: + return None + if x is None or y is None: + return x if x is not None else y + if gate is None: + return x + y + + new_y = y.clone() + selected_y = y[start_indices[1] : end_indices[1]] + selected_y = selected_y.view( + -1, self.config.action_horizon_flow, selected_y.shape[-1] + )[:, :, : self.config.dim_inputs[1]] + selected_y = selected_y.to(torch.float32) * gate + new_y[start_indices[1] : end_indices[1], : self.config.dim_inputs[1]] = ( + selected_y.view(-1, self.config.dim_inputs[1]).to(new_y.dtype) + ) + + return x + new_y + + def scatter_proprioception_embeddings( + self, input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask + ): + use_state_string_representation = getattr( + self.config, "use_state_string_representation", False + ) + if proprioception is not None and not use_state_string_representation: + proprioception = proprioception.to(inputs_embeds.device).to( + inputs_embeds.dtype + ) + agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to( + inputs_embeds.dtype + ) + proprioception = self.action_preprocessor.proprioception_proj( + proprioception, + dataset_names, + agent_pos_mask, + use_history=proprioception.shape[1] > 1, + ) + mask = input_ids == self.action_token_id_set["propri_token_id"] + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + proprioception_mask = mask_expanded.to(inputs_embeds.device) + + proprioception = proprioception.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter( + proprioception_mask, proprioception + ) + + return inputs_embeds + + def scatter_flow_action_embeddings( + self, + input_ids, + inputs_embeds, + action_chunk, + dataset_names, + sample_time, + dof_mask, + ): + if not self.config.use_flow_action_expert: + return inputs_embeds, None, None + adarms_cond, flow = None, None + if action_chunk is not None: + action_chunk = action_chunk.to(inputs_embeds.device) + dof_mask = dof_mask.to(inputs_embeds.device) + noisy_action_emb, flow, adarms_cond = self.action_preprocessor( + action_chunk, dataset_names, sample_time, dof_mask + ) + mask = input_ids == self.action_token_id_set["action_token_id"] + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + action_mask = mask_expanded.to(inputs_embeds.device) + + noisy_action_emb = noisy_action_emb.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(action_mask, noisy_action_emb) + + return inputs_embeds, flow, adarms_cond + + @staticmethod + def _update_position_ids(position_ids, moe_token_types, positional_masks): + return update_position_ids(position_ids, moe_token_types, positional_masks) + + def _update_joint_attention_mask_2d( + self, attention_mask, moe_token_types, positional_masks + ): + return update_joint_attention_mask_2d( + attention_mask, + moe_token_types, + positional_masks, + causal_action_attention_mask=self.config.causal_action_attention_mask, + ) + + def _update_joint_attention_flash_mask( + self, attention_mask, moe_token_types, positional_masks, debug=False + ): + return update_joint_attention_flash_mask( + attention_mask, + moe_token_types, + positional_masks, + causal_action_attention_mask=self.config.causal_action_attention_mask, + debug=debug, + ) + + +class ActionGenerationMixin(GenerationMixin): + action_preprocessor: ActionProcessor + + def to_bfloat16_for_selected_params(self, fsdp_plugin=None, accelerator=None): + """ + Keep selected model parameters in float32 and cast the rest to bfloat16. + - If fsdp_plugin exists, use FSDP v1 mixed_precision wrapping. + - Otherwise, modify parameter dtypes directly. + """ + + def _assign_child(root_module, dotted_name: str, new_child): + parts = dotted_name.split(".") + parent = root_module + for p in parts[:-1]: + parent = getattr(parent, p) + setattr(parent, parts[-1], new_child) + + if fsdp_plugin: + fsdp_version = getattr(fsdp_plugin, "fsdp_version", None) + if fsdp_version != 1: + raise RuntimeError("Only FSDP v1 (fsdp_version=1) is supported.") + + device = getattr( + accelerator, "device", torch.device("cuda", torch.cuda.current_device()) + ) + if isinstance(device, torch.device) and device.type == "cuda": + if device.index is not None: + torch.cuda.set_device(device.index) + device_id = device.index + + # Move the model to the target device first + self = self.to(device) + + # Define the mixed precision policy + bf16_policy = MP( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.bfloat16, + cast_forward_inputs=False, + cast_root_forward_inputs=False, + ) + + fp32_policy = MP( + param_dtype=torch.float32, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, + cast_forward_inputs=False, + cast_root_forward_inputs=False, + ) + + from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl_act import ( + Qwen2_5_VLDecoderLayer_with_MoE, + ) + from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl import ( + Qwen2_5_VLVisionBlock, + ) + + target_classes = ( + Qwen2_5_VLDecoderLayer_with_MoE, + Qwen2_5_VLVisionBlock, + ) + + # Step 1️⃣: Find top-level ActionProcessor modules and wrap them separately with FSDP (FP32) + for name, module in list(self.named_modules()): + if isinstance(module, nn.Module) and any( + k in name.lower() + for k in [ + "input_layernorm", + "post_attention_layernorm", + "model.norm", + "action_preprocessor", + ] + ): + if any( + True for _ in module.children() + ): # Wrap only leaves to avoid parent-child duplication + continue + if getattr(module, "_fsdp_wrapped", False): + continue + + logger.info("[FSDP v1] wrapping module in FP32: %s", name) + wrapped = FSDP( + module, + mixed_precision=fp32_policy, + sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, + backward_prefetch="BACKWARD_PRE", + device_id=device_id, + use_orig_params=True, + ) + _assign_child(self, name, wrapped) + setattr(wrapped, "_fsdp_wrapped", True) + + for name, module in list(self.named_modules()): + if isinstance(module, target_classes): + + if getattr(module, "_fsdp_wrapped", False): + continue + + logger.info("[FSDP v1] wrapping module in BF16: %s", name) + wrapped = FSDP( + module, + mixed_precision=bf16_policy, + sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, + backward_prefetch="BACKWARD_PRE", + device_id=device_id, + use_orig_params=True, + ) + _assign_child(self, name, wrapped) + setattr(wrapped, "_fsdp_wrapped", True) + + # Step 2️⃣: Wrap the outer module with FSDP using the BF16 policy + logger.info("[FSDP v1] wrapping root model with bf16 mixed precision...") + self = FSDP( + self, + mixed_precision=bf16_policy, + sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, + backward_prefetch="BACKWARD_PRE", + device_id=device_id, + use_orig_params=True, + ) + + return self + + # ----------------- Non-FSDP path ----------------- + else: + logger.info("Running manual dtype conversion (no FSDP).") + params_to_keep_float32 = [] + for name, _ in self.named_parameters(): + if any( + k in name + for k in [ + "input_layernorm", + "post_attention_layernorm", + "model.norm", + "action_preprocessor", + "action_processor", + ] + ): + params_to_keep_float32.append(name) + + for name, param in self.named_parameters(): + if name not in params_to_keep_float32: + param.data = param.data.to(torch.bfloat16) + if name in params_to_keep_float32: + param.data = param.data.to(torch.float32) + + return self + + def define_action_token_id(self): + # Get the action token list through tokenizer_mixin; compatible with pure flow mode + if self.processor is not None: + action_token_list = [] + if self.tokenizer_mixin is not None: + action_token_list = self.tokenizer_mixin.get_action_token_list( + self.processor + ) + + action_token_id = self.processor.tokenizer.convert_tokens_to_ids( + "<|action|>" + ) + propri_token_id = self.processor.tokenizer.convert_tokens_to_ids( + "<|propri|>" + ) + self.action_token_id_set = { + "action_token_list": action_token_list, + "propri_token_id": propri_token_id, + "action_token_id": action_token_id, + } + + def add_lora( + self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1 + ): + """Add LoRA adapters""" + config = LoraConfig( + r=r, + lora_alpha=lora_alpha, + target_modules=target_modules, + lora_dropout=lora_dropout, + bias="none", + task_type="CAUSAL_LM", + ) + self.model = get_peft_model(self.model, config) + # Log trainable parameter information + self.model.print_trainable_parameters() + + def compute_loss( + self, + hidden_states, + logits, + input_ids=None, + dataset_names=None, + labels=None, + action_chunk=None, + dof_mask=None, + flow=None, + flow_loss_mask=None, + _lm_loss_mask=None, + action_hidden_states=None, + **kwargs, + ): + if input_ids is not None: + batch_size, seq_length = input_ids.shape + + loss = 0 + cross_entropy_loss, flow_loss = None, None + + if dataset_names is not None: + unique_datasets_name = list(set(dataset_names)) + _device = hidden_states.device + _flow_channel_names = [ + f"{name}_flow" + for name in unique_datasets_name + if is_action_dataset_name(name) + ] + channel_loss_dict = { + name: torch.tensor(0.0, device=_device) + for name in unique_datasets_name + _flow_channel_names + } + channel_loss_count_dict = { + name: torch.tensor(0, device=_device) + for name in unique_datasets_name + _flow_channel_names + } + else: + unique_datasets_name, channel_loss_dict, channel_loss_count_dict = ( + None, + None, + None, + ) + + if labels is not None: + if _lm_loss_mask is not None: + # ===== Optimized path: logits already gathered for loss tokens only ===== + # logits: [N_loss, V] or None, _lm_loss_mask: [B, S-1] + if logits is not None: + shift_logits = logits.to(torch.float32) # [N_loss, V] + shift_labels = labels[..., 1:].contiguous() + shift_labels_flat = shift_labels[_lm_loss_mask] # [N_loss] + shift_labels_flat = shift_labels_flat.to(shift_logits.device) + + _cross_entropy_loss = self.loss_fct(shift_logits, shift_labels_flat) + cross_entropy_loss = _cross_entropy_loss.mean() + + # compute channel loss + if unique_datasets_name is not None: + batch_idx = ( + torch.arange(batch_size, device=_lm_loss_mask.device) + .unsqueeze(1) + .expand_as(_lm_loss_mask) + ) + loss_batch_idx = batch_idx[_lm_loss_mask] # [N_loss] + for dataset_name_i in unique_datasets_name: + ds_mask = torch.tensor( + [name == dataset_name_i for name in dataset_names], + device=_lm_loss_mask.device, + dtype=torch.bool, + ) + tok_ds_mask = ds_mask[loss_batch_idx] + channel_loss_dict[dataset_name_i] = ( + _cross_entropy_loss[tok_ds_mask].sum() + if tok_ds_mask.any() + else torch.tensor(0.0, device=shift_logits.device) + ) + channel_loss_count_dict[dataset_name_i] += tok_ds_mask.sum() + else: + cross_entropy_loss = torch.tensor(0.0, device=hidden_states.device) + else: + # ===== Original path (inference / no labels optimization) ===== + shift_logits = logits[..., :-1, :].contiguous().to(torch.float32) + shift_labels = labels[..., 1:].contiguous() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + non_ignored_mask = shift_labels != -100 + _cross_entropy_loss = self.loss_fct(shift_logits, shift_labels) + cross_entropy_loss = ( + _cross_entropy_loss[non_ignored_mask].mean() + if non_ignored_mask.any() + else torch.tensor(0.0, device=shift_logits.device) + ) + + # compute channel loss + _cross_entropy_loss = _cross_entropy_loss.view( + batch_size, seq_length - 1 + ) + non_ignored_mask = non_ignored_mask.view(batch_size, seq_length - 1) + if unique_datasets_name is not None: + for dataset_name_i in unique_datasets_name: + dataset_mask = torch.tensor( + [name == dataset_name_i for name in dataset_names], + device=logits.device, + ) + combined_mask = dataset_mask.unsqueeze(1) & non_ignored_mask + channel_loss_dict[dataset_name_i] = ( + _cross_entropy_loss[combined_mask].sum() + if combined_mask.any() + else torch.tensor(0.0, device=shift_logits.device) + ) + channel_loss_count_dict[dataset_name_i] += combined_mask.sum() + + if not torch.isnan(cross_entropy_loss): + loss += cross_entropy_loss * self.config.ar_loss_weight + else: + with torch.no_grad(): + cross_entropy_loss.detach() + + # compute action token accuracy（computed uniformly through tokenizer_mixin） + # if self.tokenizer_mixin is not None and self.action_mapper is not None: + # accuracy_dict = self.tokenizer_mixin.compute_accuracy( + # logits, labels, self.action_mapper, self.action_token_id_set + # ) + # channel_loss_dict.update(accuracy_dict) + + if action_chunk is not None: + action_mask = input_ids == self.action_token_id_set["action_token_id"] + if action_mask.any(): + if action_hidden_states is None: + action_hidden_states = hidden_states[action_mask].to(torch.float32) + else: + action_hidden_states = action_hidden_states.reshape( + -1, action_hidden_states.shape[-1] + ).to(torch.float32) + flow = flow.reshape(-1, flow.shape[-1]) + _flow_loss = self.action_preprocessor.flow_loss( + action_hidden_states, flow, action_chunk, dof_mask, flow_loss_mask + ) + if isinstance(_flow_loss, torch.Tensor): + # Compute the valid-element mask as the intersection of dof_mask and flow_loss_mask + valid_mask = ( + dof_mask.reshape(-1, dof_mask.shape[-1]) + if dof_mask is not None + else None + ) + if flow_loss_mask is not None: + flow_mask_expanded = ( + flow_loss_mask.unsqueeze(-1) + .reshape(-1, 1) + .expand(-1, _flow_loss.shape[-1]) + ) + valid_mask = ( + valid_mask * flow_mask_expanded + if valid_mask is not None + else flow_mask_expanded + ) + if valid_mask is not None and not valid_mask.all(): + flow_loss = _flow_loss.sum() / valid_mask.sum() + else: + flow_loss = _flow_loss.mean() + loss += flow_loss + _flow_loss = _flow_loss.view( + dof_mask.shape[0], dof_mask.shape[1], dof_mask.shape[2] + ) + + # compute flow channel loss + if unique_datasets_name is not None: + B, T, D = _flow_loss.shape + action_ds_names = [ + name + for name in dataset_names + if is_action_dataset_name(name) + ] + if valid_mask is not None: + valid_mask_3d = valid_mask.view(B, T, D) + else: + valid_mask_3d = torch.ones_like( + _flow_loss, dtype=torch.bool, device=_flow_loss.device + ) + for dataset_name_i in unique_datasets_name: + ds_mask = torch.tensor( + [name == dataset_name_i for name in action_ds_names], + device=_flow_loss.device, + dtype=torch.bool, + ) + if not ds_mask.any(): + continue + flow_key = f"{dataset_name_i}_flow" + ds_mask_3d = ds_mask.view(-1, 1, 1).expand(B, T, D) + flow_loss_sum = (_flow_loss * ds_mask_3d).sum() + flow_count = (valid_mask_3d * ds_mask_3d).sum() + channel_loss_dict[flow_key] = ( + channel_loss_dict[flow_key] + flow_loss_sum + ) + channel_loss_count_dict[flow_key] = ( + channel_loss_count_dict[flow_key] + flow_count + ) + + return ( + loss, + cross_entropy_loss, + flow_loss, + channel_loss_dict, + channel_loss_count_dict, + ) diff --git a/wall_x/model/model_utils.py b/wall_x/model/model_utils.py deleted file mode 100644 index 5aaeecc6..00000000 --- a/wall_x/model/model_utils.py +++ /dev/null @@ -1,319 +0,0 @@ -import torch -import os -import numpy as np -from transformers import AutoProcessor -from wall_x.model.action_head import Normalizer - - -def update_model_config(train_config, model_config): - model_config.use_state_string_representation = train_config["data"].get( - "use_state_string_representation", False - ) - model_config.flow_loss_weight = train_config.get("flow_loss_weight", 1.0) - - model_config.dof_config = train_config["dof_config"] - model_config.agent_pos_config = train_config["agent_pos_config"] - - model_config.action_horizon_flow = train_config["data"].get( - "action_horizon_flow", 32 - ) - - if train_config.get("_attn_implementation", None) is not None: - model_config._attn_implementation = train_config["_attn_implementation"] - - return model_config - - -def load_wallx_processors(config): - processor = AutoProcessor.from_pretrained(config["processor_path"], use_fast=True) - # pad side = left - processor.tokenizer.padding_side = "left" - - new_tokens = ["<|propri|>", "<|action|>"] - # special_tokens = [] - action_tokenizer_type = config.get("action_tokenizer_type", None) - if action_tokenizer_type == "fast": - train_action_tokenizer = AutoProcessor.from_pretrained( - config["action_tokenizer_path"], trust_remote_code=True - ) - val_action_tokenizer = AutoProcessor.from_pretrained( - config["action_tokenizer_path"], trust_remote_code=True - ) - new_tokens += [ - f"<|action_token_{i}|>" for i in range(train_action_tokenizer.vocab_size) - ] - elif action_tokenizer_type == "spatialvla": - raise NotImplementedError("SpatialActionTokenizer is not implemented") - else: - train_action_tokenizer = None - val_action_tokenizer = None - - num_added_tokens = processor.tokenizer.add_tokens(new_tokens) - - if action_tokenizer_type and train_action_tokenizer.vocab_size > 0: - action_mapper = {} - for i in range(train_action_tokenizer.vocab_size): - token = f"<|action_token_{i}|>" - token_id = processor.tokenizer.convert_tokens_to_ids(token) - action_mapper[token_id] = i - else: - action_mapper = None - - return { - "processor": processor, - "train_action_tokenizer": train_action_tokenizer, - "val_action_tokenizer": val_action_tokenizer, - "action_mapper": action_mapper, - "num_added_tokens": num_added_tokens, - } - - -def register_normalizers(config, model_path): - # if config.get("customized_action_statistic_dof", None): - # action_statistic_dof = json.load(open(config["customized_action_statistic_dof"], "r")) - # else: - # action_statistic_dof = default_action_statistic_dof - - action_statistic_dof = None - - if os.path.exists(model_path + "/normalizer_action.pth"): - print( - "Loading normalizer_action from checkpoint", - model_path + "/normalizer_action.pth", - flush=True, - ) - normalizer_action = Normalizer.from_ckpt(model_path + "/normalizer_action.pth") - else: - normalizer_action = Normalizer( - action_statistic_dof, - config["dof_config"], - min_key=config.get("min_key", "min"), - delta_key=config.get("delta_key", "delta"), - ) - - # print("action_statistic_dof",action_statistic_dof) - - if os.path.exists(model_path + "/normalizer_propri.pth"): - print( - "Loading normalizer_propri from checkpoint", - model_path + "/normalizer_propri.pth", - flush=True, - ) - normalizer_propri = Normalizer.from_ckpt(model_path + "/normalizer_propri.pth") - else: - normalizer_propri = Normalizer( - action_statistic_dof, - config["agent_pos_config"], - min_key=config.get("min_key", "min"), - delta_key=config.get("delta_key", "delta"), - ) - - return normalizer_action, normalizer_propri - - -def find_first_last_ones(tensor): - """ - Input: a tensor of shape (bs, seq_len) containing 0s and 1s - Output: (first_indices, last_indices), each of shape (bs,) - where first_indices[i] is the index of the first 1 in the i-th batch, or -1 if none exists. - last_indices[i] is the index of the last 1 in the i-th batch, or -1 if none exists. - """ - bs, seq_len = tensor.shape - masks = tensor == 1 - has_ones = masks.any(dim=1) - - first = torch.full((bs,), -1, dtype=torch.long, device=tensor.device) - last = first.clone() - - first[has_ones] = torch.argmax(masks[has_ones].float(), dim=1) - - flipped_masks = masks.flip(dims=[1]) - last_argmax = torch.argmax(flipped_masks[has_ones].float(), dim=1) - last[has_ones] = seq_len - 1 - last_argmax - - return first, last - - -def flashmask_to_densemask(startend_row_indices, dtype, causal=True): - if startend_row_indices is None: - return None - bz, num_head, seq_len, bound_num = startend_row_indices.shape - m = np.ones((bz, num_head, seq_len, seq_len), dtype=dtype) - has_end = (causal and bound_num == 2) or ((not causal) and bound_num == 4) - for bi in range(bz): - for hi in range(num_head): - for j in range(seq_len): - downstart = startend_row_indices[bi, hi, j, 0] - if has_end: - downend = startend_row_indices[bi, hi, j, 1] - m[bi, hi, downstart:downend, j] = 0 - else: - m[bi, hi, downstart:, j] = 0 - if causal: - m[bi, hi, :j, j] = 0 - else: - if has_end: - upstart = startend_row_indices[bi, hi, j, 2] - upend = startend_row_indices[bi, hi, j, 3] - m[bi, hi, upstart:upend, j] = 0 - else: - upend = startend_row_indices[bi, hi, j, 1] - m[bi, hi, :upend, j] = 0 - return m - - -def num_floating_point_operations( - args, - batch_size: int, - num_lang_tokens: int, - num_action_tokens: int, - vision_seq_length: int = 756, -): - """ - Accurately estimate the training FLOPs of Transformer + MoE + MoT + Vision. - - Supported: - - expert0 = language tokens - - expert1 = action tokens - - MoE MLP (2 experts) - - MoT Attention (2 experts) - - GQA - - Vision Transformer (full+window attention) - """ - assert args.num_experts == 2, "The current model only supports 2 experts." - - dim_lang, dim_act = args.dim_inputs - - # Number of tokens per layer (flattened across batch) - N_lang = batch_size * num_lang_tokens - N_action = batch_size * num_action_tokens - N_total = N_lang + N_action # Used for non-MoT attention - - # ================================================================ - # Text MLP FLOPs - # ================================================================ - hidden_size = args.hidden_size - ffn_hidden_size = args.intermediate_size - num_layers = args.num_hidden_layers - - use_moe_mlp = getattr(args, "mlp_moe", False) - - # ---------- Forward-only MLP FLOPs ---------- - def forward_mlp_flops(N, d_in, d_ff): - """ - SwiGLU forward: - gate = x @ W1 (2*N*d_in*d_ff) - up = x @ W2 (2*N*d_in*d_ff) - act = silu + mul (~2*N*d_ff) - down = h @ W3 (2*N*d_ff*d_in) - - Forward ≈ 4*N*d_in*d_ff + 2*N*d_ff*d_in = 6*N*d_in*d_ff + 2*N*d_ff - """ - return 6 * N * d_in * d_ff + 2 * N * d_ff - - if not use_moe_mlp: - # Dense MLP - F_fwd = forward_mlp_flops(N_total, hidden_size, ffn_hidden_size) - total_mlp_flops_text = 3 * num_layers * F_fwd # <-- training FLOPs - else: - # MoE: expert0(language) + expert1(action) - hid_lang = args.experts[0]["intermediate_size"] - hid_act = args.experts[1]["intermediate_size"] - - F_lang_fwd = forward_mlp_flops(N_lang, dim_lang, hid_lang) - F_act_fwd = forward_mlp_flops(N_action, dim_act, hid_act) - - total_mlp_flops_text = 3 * num_layers * (F_lang_fwd + F_act_fwd) - - # ================================================================ - # Text Attention FLOPs - # ================================================================ - num_heads = args.num_attention_heads - num_kv = args.num_key_value_heads - H = hidden_size - B = batch_size - S = num_lang_tokens + num_action_tokens - N = B * S - - use_mot = getattr(args, "attention_moe", False) - - # ---------- attention matmul ---------- - F_matmul_fwd = 4 * B * (S**2) * H - - if not use_mot: - # -------- GQA + QKV / O -------- - F_q_fwd = 2 * N * H * H - F_kv_fwd = 4 * N * H * H * (num_kv / num_heads) - F_o_fwd = 2 * N * H * H - - F_attn_fwd = F_q_fwd + F_kv_fwd + F_o_fwd + F_matmul_fwd - - else: - # -------- MoT: expert0 + expert1 QKV -------- - F_lang_qkv = N_lang * dim_lang * H * (2 + 4 * num_kv / num_heads) - F_act_qkv = N_action * dim_act * H * (2 + 4 * num_kv / num_heads) - F_attn_fwd = F_lang_qkv + F_act_qkv + F_matmul_fwd - - # Training FLOPs - total_attn_flops_text = 3 * num_layers * F_attn_fwd - - # ================================================================ - # Logits projection FLOPs - # ================================================================ - vocab_size = getattr(args, "padded_vocab_size", args.vocab_size) - - F_logits_fwd = 2 * N * H * vocab_size - total_logits_flops = 3 * F_logits_fwd - - total_text_flops = total_mlp_flops_text + total_attn_flops_text + total_logits_flops - - # ================================================================ - # Vision Transformer FLOPs - # ================================================================ - total_vision_flops = 0 - - if hasattr(args, "vision_config") and vision_seq_length is not None: - vcfg = args.vision_config - - Bv = batch_size - Sv = vision_seq_length - Nv = Bv * Sv - - Hv = vcfg.hidden_size - Iv = vcfg.intermediate_size - num_heads_v = vcfg.num_heads - window_size = vcfg.window_size - out_hidden = vcfg.out_hidden_size - - depth_v = vcfg.depth - fullatt = set(vcfg.fullatt_block_indexes) - num_full = len(fullatt) - num_local = depth_v - num_full - - # ---------- forward FLOPs ---------- - def forward_vit_mlp(N, H, Inner): - return 6 * N * H * Inner + 2 * N * Inner - - F_mlp_v = forward_vit_mlp(Nv, Hv, Iv) - F_qkv_v = 6 * Nv * Hv * Hv - F_o_v = 2 * Nv * Hv * Hv - - F_full = 4 * Bv * (Sv**2) * Hv - num_windows = Sv / window_size - F_win = 4 * Bv * num_windows * (window_size**2) * (Hv / num_heads_v) - - F_block_full_fwd = F_mlp_v + F_qkv_v + F_o_v + F_full - F_block_local_fwd = F_mlp_v + F_qkv_v + F_o_v + F_win - - # ---------- train FLOPs ---------- - total_vision_flops = 3 * ( - num_full * F_block_full_fwd + num_local * F_block_local_fwd - ) - - # merger - total_vision_flops += 3 * (2 * Nv * Hv * out_hidden) - - # ================================================================ - # TOTAL TRAIN FLOPs - # ================================================================ - return total_text_flops + total_vision_flops diff --git a/wall_x/model/qact/__init__.py b/wall_x/model/qact/__init__.py new file mode 100644 index 00000000..860ac446 --- /dev/null +++ b/wall_x/model/qact/__init__.py @@ -0,0 +1 @@ +"""QAct (Qwen-VLA) model family.""" diff --git a/wall_x/model/qwen2_5_based/__init__.py b/wall_x/model/qact/qwen2_5/__init__.py similarity index 58% rename from wall_x/model/qwen2_5_based/__init__.py rename to wall_x/model/qact/qwen2_5/__init__.py index d9ed878f..1daa3c4d 100644 --- a/wall_x/model/qwen2_5_based/__init__.py +++ b/wall_x/model/qact/qwen2_5/__init__.py @@ -1,8 +1,2 @@ from .modeling_qwen2_5_vl_act import Qwen2_5_VLMoEModel, Qwen2_5_VLMoEForAction from .configuration_qwen2_5_vl import Qwen2_5_VLConfig - -__all__ = [ - "Qwen2_5_VLMoEModel", - "Qwen2_5_VLMoEForAction", - "Qwen2_5_VLConfig", -] diff --git a/wall_x/model/qact/qwen2_5/adapter.py b/wall_x/model/qact/qwen2_5/adapter.py new file mode 100644 index 00000000..ace8a451 --- /dev/null +++ b/wall_x/model/qact/qwen2_5/adapter.py @@ -0,0 +1,56 @@ +"""Qwen2.5 VLA adapter — variant-specific overrides on top of VLAdapter.""" + +from wall_x.model.registry import register_model +from wall_x.trainer.adapters.vla_model_adapter import VLAdapter + + +@register_model("qwen2_5") +class Qwen2_5Adapter(VLAdapter): + MODEL_TYPE = "qwen2_5" + + @classmethod + def model_class(cls): + from wall_x.model.qact.qwen2_5 import Qwen2_5_VLMoEForAction + + return Qwen2_5_VLMoEForAction + + @classmethod + def config_class(cls): + from wall_x.model.qact.qwen2_5 import Qwen2_5_VLConfig + + return Qwen2_5_VLConfig + + @classmethod + def inference_model_class(cls): + return cls.model_class() + + def get_transformer_layer_cls(self): + layer_classes = set() + try: + from transformers.models.qwen2_vl.modeling_qwen2_vl import ( + Qwen2VLDecoderLayer, + ) + + layer_classes.add(Qwen2VLDecoderLayer) + except ImportError: + pass + try: + from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl import ( + Qwen2_5_VLDecoderLayer, + ) + + layer_classes.add(Qwen2_5_VLDecoderLayer) + except ImportError: + pass + return layer_classes if layer_classes else None + + @staticmethod + def log_attention_implementation(logger, model): + logger.info( + f"*** model attention implementation: " + f"{model.model._attn_implementation} ***" + ) + logger.info( + f"*** model.visual attention implementation: " + f"{model.visual.config._attn_implementation} ***" + ) diff --git a/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py b/wall_x/model/qact/qwen2_5/configuration_qwen2_5_vl.py similarity index 77% rename from wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py rename to wall_x/model/qact/qwen2_5/configuration_qwen2_5_vl.py index a2d51a1b..44a6bff2 100644 --- a/wall_x/model/qwen2_5_based/configuration_qwen2_5_vl.py +++ b/wall_x/model/qact/qwen2_5/configuration_qwen2_5_vl.py @@ -1,6 +1,35 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_qwen2_5_vl.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# coding=utf-8 +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import logging + from transformers.configuration_utils import PretrainedConfig from transformers.modeling_rope_utils import rope_config_validation +logger = logging.getLogger(__name__) + class Qwen2_5_VLVisionConfig(PretrainedConfig): model_type = "qwen2_5_vl" @@ -203,18 +232,24 @@ def __init__( mlp_moe=False, norm_moe=False, mot_opt=False, - flow_loss_weight=10, + ar_loss_weight=1.0, use_state_string_representation=False, use_adarms=False, proj_with_mask=True, - use_flow_action_expert=True, adarms_cond_dim=None, action_horizon_flow=32, causal_action_attention_mask=False, + use_flow_action_expert=True, use_x_pred=False, attn_deterministic=False, + use_x_loss=False, **kwargs, ): + # Compatibility with newer transformers versions (5.x): + # - Older versions: super() sets self.pad_token_id; override it with the saved value so kwargs are preserved + # - Newer versions: super() does not set self.pad_token_id; assign it afterward + _pad_token_id = kwargs.pop("pad_token_id", None) + self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size @@ -249,7 +284,7 @@ def __init__( self.mlp_moe = mlp_moe self.norm_moe = norm_moe self.mot_opt = mot_opt - self.flow_loss_weight = flow_loss_weight + self.ar_loss_weight = ar_loss_weight self.use_state_string_representation = use_state_string_representation self.use_adarms = use_adarms @@ -258,9 +293,10 @@ def __init__( self.use_flow_action_expert = use_flow_action_expert self.action_horizon_flow = action_horizon_flow self.causal_action_attention_mask = causal_action_attention_mask + self.use_flow_action_expert = use_flow_action_expert self.use_x_pred = use_x_pred self.attn_deterministic = attn_deterministic - + self.use_x_loss = use_x_loss # Validate the correctness of rotary position embeddings parameters # BC: if there is a 'type' field, move it to 'rope_type'. # and change type from 'mrope' to 'default' because `mrope` does defeault RoPE calculations @@ -274,6 +310,9 @@ def __init__( super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + # Assign pad_token_id as described above, overriding older super() values or filling newer missing attributes + self.pad_token_id = _pad_token_id + # move vision config initialization after super init to avoid recursively set in latest transformers version # TODO: make it better if isinstance(vision_config, dict): @@ -281,5 +320,38 @@ def __init__( elif vision_config is None: self.vision_config = self.sub_configs["vision_config"]() + def update_model_config(self, train_config): + """Update model configuration from training config. + + This method updates the model configuration with training-specific + settings such as action horizon, DOF config, attention implementation, etc. + + Args: + train_config: dict containing training configuration parameters. + """ + self.use_state_string_representation = train_config["data"].get( + "use_state_string_representation", False + ) + self.ar_loss_weight = train_config.get("ar_loss_weight", 1.0) + + self.dof_config = train_config["dof_config"] + self.agent_pos_config = train_config["agent_pos_config"] + + self.action_horizon_flow = train_config["data"].get("action_horizon_flow", 32) + + if train_config.get("_attn_implementation", None) is not None: + self._attn_implementation = train_config["_attn_implementation"] + + if train_config.get("attn_deterministic", None) is not None: + self.attn_deterministic = train_config["attn_deterministic"] + self.vision_config.attn_deterministic = train_config["attn_deterministic"] + logger.debug("Attention is using deterministic kernel for this run") + else: + self.attn_deterministic = True + self.vision_config.attn_deterministic = True + + if train_config.get("noise_scheduler", None) is not None: + self.noise_scheduler = train_config["noise_scheduler"] + __all__ = ["Qwen2_5_VLConfig"] diff --git a/wall_x/model/qact/qwen2_5/inference_mixin.py b/wall_x/model/qact/qwen2_5/inference_mixin.py new file mode 100644 index 00000000..b1de47a3 --- /dev/null +++ b/wall_x/model/qact/qwen2_5/inference_mixin.py @@ -0,0 +1,1039 @@ +import logging + +import numpy as np +import torch +from typing import Optional, List, Tuple +from torchdiffeq import odeint + +logger = logging.getLogger(__name__) + + +def topk_right_tie_break_1d(x, k): + L = x.size(0) + x_rev = torch.flip(x, [0]) + idx_in_rev = torch.argsort(x_rev, dim=0, descending=True, stable=True) + orig_idx = (L - 1) - idx_in_rev + topk_idx = orig_idx[:k] + topk_vals = x[topk_idx] + return topk_vals, topk_idx + + +def add_gumbel_noise(logits, temperature): + """ + The Gumbel max is a method for sampling categorical distributions. + According to arXiv:2409.02908, for MDM, low-precision Gumbel Max improves perplexity score but reduces generation quality. + Thus, we use float64. + """ + if temperature == 0: + return logits + logits = logits.to(torch.float64) + noise = torch.rand_like(logits, dtype=torch.float64) + gumbel_noise = (-torch.log(noise)) ** temperature + return logits.exp() / gumbel_noise + + +class VLAInferenceMixin: + + # TODO: Integrate with the optimized implementation. + def prepare_inputs_embeds( + self, + input_ids: torch.LongTensor, + inputs_embeds: Optional[torch.FloatTensor] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, + dataset_names: Optional[str] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ) -> Tuple[torch.FloatTensor, Optional[torch.Tensor]]: + """ + Prepare model input embeddings, including text, image, video, and proprioception embeddings. + + Args: + input_ids: Input token IDs. + inputs_embeds: Precomputed input embeddings, if provided. + pixel_values: Image pixel values. + pixel_values_videos: Video pixel values. + image_grid_thw: Image grid time, height, and width. + video_grid_thw: Video grid time, height, and width. + proprioception: Proprioception data. + dataset_names: Dataset names. + agent_pos_mask: Agent position mask. + attention_mask: Attention mask. + + Returns: + inputs_embeds: Complete input embeddings. + attention_mask: Processed attention mask. + """ + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + + # Process image embeddings. + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_features = image_embeds.shape[0] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + + mask = input_ids == self.config.image_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + image_mask = mask_expanded.to(inputs_embeds.device) + + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + # Process video embeddings. + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + + mask = input_ids == self.config.video_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + video_mask = mask_expanded.to(inputs_embeds.device) + + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + # Process proprioception embeddings. + if proprioception is not None and not getattr( + self.config, "use_state_string_representation", False + ): + proprioception = proprioception.to(inputs_embeds.device) + agent_pos_mask = agent_pos_mask.to(inputs_embeds.device) + proprio_embed = self.action_preprocessor.proprioception_proj( + proprioception, + dataset_names, + agent_pos_mask, + use_history=proprioception.shape[1] > 1, + ) + proprioception_mask = ( + input_ids == self.action_token_id_set["propri_token_id"] + ) + inputs_embeds[proprioception_mask] = proprio_embed.reshape( + -1, inputs_embeds.shape[-1] + ).to(inputs_embeds.dtype) + + # Process the attention mask. + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + return inputs_embeds, attention_mask + + def prepare_position_ids( + self, + input_ids: torch.LongTensor, + inputs_embeds: torch.FloatTensor, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + ) -> torch.LongTensor: + """ + Prepare position IDs, including RoPE delta calculation and caching. + + Args: + input_ids: Input token IDs. + inputs_embeds: Input embeddings. + image_grid_thw: Image grid time, height, and width. + video_grid_thw: Video grid time, height, and width. + second_per_grid_ts: Time step for each grid. + attention_mask: Attention mask. + position_ids: Precomputed position IDs, if provided. + cache_position: Cache position. + past_key_values: Previous key/value cache. + + Returns: + position_ids: Calculated position IDs. + """ + # RoPE deltas cannot be calculated once the attention mask is 4D. + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2 + ): + # Calculate RoPE indices once per generation in the prefill stage. + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + attention_mask, + ) + self.rope_deltas = rope_deltas + # Reuse the previous RoPE deltas to obtain the correct position IDs. + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + return position_ids + + @torch.no_grad() + def generate_dllm_action( + self, + input_ids, + action_horizon, + action_dim, + ar_action_dim, + total_ar_step, + use_ar_action: bool = False, + num_inference_timesteps: int = 10, + prefix_length: Optional[int] = None, # Prefix length. + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, + positional_masks: Optional[torch.LongTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + action_chunk: Optional[torch.FloatTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, + unnorm_proprioception: Optional[torch.FloatTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + dataset_names: Optional[str] = None, + dof_mask: Optional[torch.FloatTensor] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + robot_type_id: Optional[int] = None, # Used for v3.1 delta decoding. + **kwargs, + ): + batch_size = ( + input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] + ) + inputs_embeds, attention_mask = self.prepare_inputs_embeds( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + proprioception=proprioception, + dataset_names=dataset_names, + agent_pos_mask=agent_pos_mask, + attention_mask=attention_mask, + ) + position_ids = self.prepare_position_ids( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + attention_mask=attention_mask, + position_ids=position_ids, + cache_position=cache_position, + past_key_values=past_key_values, + ) + + noisy_action = torch.randn( + size=(batch_size, action_horizon, action_dim), + dtype=torch.float32, + device=inputs_embeds.device, + ) + + times = torch.linspace( + 0.0, + 1.0, + num_inference_timesteps + 1, + device=inputs_embeds.device, + dtype=torch.float32, + ) + + dt = times[1] - times[0] + time_0 = times[0].unsqueeze(0).repeat(noisy_action.shape[0]) + action_embed, adarms_cond = self.action_preprocessor.step( + timestep=time_0, noisy_action=noisy_action, dof_mask=dof_mask + ) + action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]).to( + inputs_embeds.dtype + ) + flow_action_mask = input_ids == self.action_token_id_set["action_token_id"] + inputs_embeds[flow_action_mask] = action_embed + + # Compute the token span for each expert group after permutation. + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + + prefetch_output = self.model( + input_ids=None, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=None, + inputs_embeds=inputs_embeds, + moe_token_types=moe_token_types, + positional_masks=positional_masks, + use_cache=True, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + adarms_conds=[None, adarms_cond], + start_indices=start_indices, + end_indices=end_indices, + ) + + hidden_states = prefetch_output.last_hidden_state + prefix_kv_cache = prefetch_output.past_key_values + + action_hidden_states = hidden_states[flow_action_mask].to(torch.float32) + v_0 = self.action_preprocessor.action_proj_back( + action_hidden_states[:, : self.action_preprocessor.action_hidden_size] + ) + noisy_action = noisy_action + dt * v_0.reshape( + batch_size, action_horizon, action_dim + ) + + ar_prefix_length = positional_masks["ar_action_mask"].nonzero(as_tuple=True)[ + -1 + ][ + 0 + ] # Attention, not tested yet with batch + start_indices[1] -= ar_prefix_length.item() + end_indices -= ar_prefix_length.item() + if hasattr(prefix_kv_cache, "key_cache"): + for layer_i in range(len(prefix_kv_cache.key_cache)): + prefix_kv_cache.key_cache[layer_i] = prefix_kv_cache.key_cache[layer_i][ + :, :, :ar_prefix_length, : + ] + prefix_kv_cache.value_cache[layer_i] = prefix_kv_cache.value_cache[ + layer_i + ][:, :, :ar_prefix_length, :] + else: + for layer_i in range(len(prefix_kv_cache.layers)): + prefix_kv_cache.layers[layer_i].keys = prefix_kv_cache.layers[ + layer_i + ].keys[:, :, :ar_prefix_length, :] + prefix_kv_cache.layers[layer_i].values = prefix_kv_cache.layers[ + layer_i + ].values[:, :, :ar_prefix_length, :] + + ar_postfix_position_ids = position_ids[:, :, ar_prefix_length:] + ar_postfix_inputs_embeds = inputs_embeds[:, ar_prefix_length:, :] + ar_postfix_attention_mask = attention_mask[:, ar_prefix_length:] + ar_postfix_moe_token_types = moe_token_types[:, ar_prefix_length:] + ar_postfix_input_ids = input_ids[:, ar_prefix_length:] + + noisy_steps_per_sample = ( + torch.ceil((1.0 - times[0] - dt) * (total_ar_step + 1)).long() - 1 + ) + noisy_steps_per_sample = noisy_steps_per_sample.clamp(min=0, max=total_ar_step) + ar_step = total_ar_step - noisy_steps_per_sample + ready_step = [False] * total_ar_step + # Always initialize ar_tokens to record generated discrete codes. + # Use a tensor to support multi-token steps in v3.1 delta mode. + ar_tokens = torch.full( + (ar_postfix_inputs_embeds.shape[1],), + -1, + dtype=torch.long, + device=inputs_embeds.device, + ) + ar_postfix_inputs_embeds, ready_step, ar_tokens = self.update_ar_action( + ar_postfix_inputs_embeds, + hidden_states[:, ar_prefix_length:], + positional_masks["ar_action_mask"][:, ar_prefix_length:], + ar_step=ar_step, + ready_step=ready_step, + ar_tokens=ar_tokens, + ) + + pad_token_id = self.processor.tokenizer.pad_token_id + padding_mask = input_ids == pad_token_id + + ar_postfix_length = input_ids.shape[-1] - ar_prefix_length + _ar_postfix_attention_mask = torch.ones( + (batch_size, ar_postfix_length, ar_prefix_length + ar_postfix_length), + dtype=torch.bool, + device=ar_postfix_attention_mask.device, + ) + + ar_postfix_padding_mask = padding_mask[ + :, ar_prefix_length: + ] # [batch_size, postfix_length] + full_padding_mask = padding_mask # [batch_size, prefix_length + postfix_length] + + for batch_idx in range(padding_mask.shape[0]): + # Disable rows for padded query positions. + _ar_postfix_attention_mask[ + batch_idx, ar_postfix_padding_mask[batch_idx], : + ] = False + # Disable columns for padded key positions. + _ar_postfix_attention_mask[batch_idx, :, full_padding_mask[batch_idx]] = ( + False + ) + + # The 3D attention mask already filters padding, but kv-cache generation + # differs from training, so rebuild the mask before passing it to the model. + if not positional_masks.get("ar_visible", True): + ar_action_mask = positional_masks["ar_action_mask"] != 0 + flow_positions = moe_token_types == 1 + flow_flow_mask = flow_positions[:, :, None] & flow_positions[:, None, :] + ar_ar_mask = ar_action_mask[:, :, None] & ar_action_mask[:, None, :] + flow_ar_mask = flow_flow_mask | ar_ar_mask + kvcache_flow_ar_mask = flow_ar_mask[:, ar_prefix_length:] + + affected = ar_action_mask | flow_positions # (B, N) + affected_pair = affected[:, :, None] & affected[:, None, :] + affected_pair = affected_pair[:, ar_prefix_length:] + + _ar_postfix_attention_mask = torch.where( + affected_pair, kvcache_flow_ar_mask, _ar_postfix_attention_mask + ) + + def step_with_kvcache( + timestep, + noisy_action, + ar_postfix_inputs_embeds, + _ar_postfix_attention_mask, + ar_prefix_length, + ready_step, + dt, + ar_tokens, + ): + action_mask = ( + ar_postfix_input_ids == self.action_token_id_set["action_token_id"] + ) + assert action_mask.any(), "No action token found in input_ids" + timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0]) + action_embed, adarms_cond = self.action_preprocessor.step( + timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask + ) + action_embed = action_embed.reshape(-1, ar_postfix_inputs_embeds.shape[-1]) + + # Clone inputs_embeds locally to keep this step isolated. + temp_inputs_embeds = ar_postfix_inputs_embeds.clone() + temp_inputs_embeds[action_mask] = action_embed.to(temp_inputs_embeds.dtype) + + transformer_outputs = self.model( + input_ids=None, + attention_mask=_ar_postfix_attention_mask, + position_ids=ar_postfix_position_ids, + past_key_values=prefix_kv_cache, + inputs_embeds=temp_inputs_embeds, + moe_token_types=ar_postfix_moe_token_types, + use_cache=False, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + adarms_conds=[None, adarms_cond], + start_indices=start_indices, + end_indices=end_indices, + ) + + hidden_states = transformer_outputs.last_hidden_state + + noisy_steps_per_sample = ( + torch.ceil((1.0 - timestep - dt) * (total_ar_step + 1)).long() - 1 + ) + noisy_steps_per_sample = noisy_steps_per_sample.clamp( + min=0, max=total_ar_step + ) + ar_step = total_ar_step - noisy_steps_per_sample + ar_postfix_inputs_embeds, ready_step, ar_tokens = self.update_ar_action( + ar_postfix_inputs_embeds, + hidden_states, + positional_masks["ar_action_mask"][:, ar_prefix_length:], + ar_step=ar_step, + ready_step=ready_step, + ar_tokens=ar_tokens, + ) + + action_hidden_states = hidden_states[action_mask].to(torch.float32) + v_t = self.action_preprocessor.action_proj_back( + action_hidden_states[:, : self.action_preprocessor.action_hidden_size] + ) + + return v_t.reshape(batch_size, action_horizon, action_dim) + + action_trajectory = odeint( + lambda timestep, noisy_action: step_with_kvcache( + timestep, + noisy_action, + ar_postfix_inputs_embeds, + _ar_postfix_attention_mask, + ar_prefix_length, + ready_step, + dt, + ar_tokens, + ), + noisy_action, + times[1:], + method="euler", + ) + output = {} + if use_ar_action: + # Decode through tokenizer_mixin.decode_action() for both fast and v3.1 delta modes. + # ar_tokens is already a tensor. + ar_tokens_tensor = ar_tokens.unsqueeze(0) # [1, seq_len] + + predict_action, decode_success = self.tokenizer_mixin.decode_action( + output_ids=ar_tokens_tensor, + action_mapper=self.action_mapper, + action_horizon=action_horizon, + action_dim=ar_action_dim, + device=inputs_embeds.device, + proprioception=proprioception, + dof_mask=dof_mask, + robot_type_id=robot_type_id, + ) + + if not decode_success: + logger.warning("Error in DLLM decoding action, predict_action is None") + output["predict_action"] = None + else: + # unnormalize + if isinstance(predict_action, np.ndarray): + predict_action = torch.tensor( + predict_action, device=inputs_embeds.device + ) + elif predict_action.device != inputs_embeds.device: + predict_action = predict_action.to(inputs_embeds.device) + + # Add the batch dimension when needed. + if predict_action.dim() == 2: + predict_action = predict_action.unsqueeze(0) + + # Decide whether dof_mask is needed from the mixin setting. + uses_dof_mask = ( + self.tokenizer_mixin.uses_dof_mask_for_unnorm + if self.tokenizer_mixin is not None + else True + ) + + if uses_dof_mask: + predict_action = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names, dof_mask + ) + ) + else: + predict_action = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names, None + ) + ) + output["predict_action"] = predict_action + + else: + predict_action = action_trajectory[-1] + predict_action = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names + ) + ) + output["predict_action"] = predict_action + + if action_chunk is not None: + output["gt_action"] = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + action_chunk, dataset_names + ) + ) + return output + + def update_ar_action2( + self, + ar_postfix_inputs_embeds, + ar_hidden_states, + ar_action_mask, + ar_step, + ready_step, + remask=False, + ar_tokens=None, + ): + ready_num = sum(ready_step) + if ready_num >= ar_step and not remask: + return ar_postfix_inputs_embeds, ready_step, ar_tokens + valid_steps = torch.unique( + ar_action_mask[(ar_action_mask != 0) & (ar_action_mask != -1)] + ) + placeholder_seq_embed = self.model.embed_tokens( + torch.Tensor(self.processor.placeholder_seq) + .to(self.model.device) + .to(torch.int) + ) + step_confidences = [] + step_token_ids = [] + for step in valid_steps: + if ready_step[step - 1] and not remask: + step_confidences.append(1.0) + step_token_ids.append(None) + continue + step_mask = ar_action_mask == step + step_hidden_states = ar_hidden_states[step_mask] + logits = self.lm_head(step_hidden_states) + pred_ids = logits.argmax(dim=-1) + step_confidence = torch.softmax(logits, dim=-1) + total_step_confidence = step_confidence.max(dim=-1).values.mean() + step_confidences.append(total_step_confidence) + step_token_ids.append(pred_ids) + step_confidences = torch.tensor(step_confidences).to(self.model.device) + _, top_k_indices = topk_right_tie_break_1d(step_confidences, ar_step) + for i, step in enumerate(valid_steps): + if ready_step[step - 1] and not remask: + continue + step_mask = ar_action_mask == step + step_indices = step_mask.nonzero(as_tuple=True)[1] + if i in top_k_indices: + # update + ar_postfix_inputs_embeds[0, step_indices, :] = self.model.embed_tokens( + step_token_ids[i] + ) + if ar_tokens is not None: + ar_tokens[step_indices] = step_token_ids[i] + ready_step[step - 1] = True + elif remask and ready_step[step - 1]: + # back to placeholder + ar_postfix_inputs_embeds[0, step_indices, :] = placeholder_seq_embed[ + : len(step_indices) + ] + ready_step[step - 1] = False + + return ar_postfix_inputs_embeds, ready_step, ar_tokens + + def update_ar_action( + self, + ar_postfix_inputs_embeds, + ar_hidden_states, + ar_action_mask, + ar_step, + ready_step, + remask=False, + ar_tokens=None, + ): + ready_num = sum(ready_step) + if ready_num >= ar_step and not remask: + return ar_postfix_inputs_embeds, ready_step, ar_tokens + update_step = ar_step - ready_num + valid_steps = torch.unique( + ar_action_mask[(ar_action_mask != 0) & (ar_action_mask != -1)] + ) + placeholder_seq_embed = self.model.embed_tokens( + torch.Tensor(self.processor.placeholder_seq) + .to(self.model.device) + .to(torch.int) + ) + step_confidences = [] + step_token_ids = [] + + logits = self.lm_head(ar_hidden_states) + confidence = torch.softmax(logits, dim=-1) + pred_ids = logits.argmax(dim=-1) + pad_mask = pred_ids == 151668 # pad token id + step_is_ready = [] + step_is_pad = [] + for step in valid_steps: + step_mask = ar_action_mask == step + step_pred_ids = pred_ids[step_mask] + step_token_ids.append(step_pred_ids) + + is_pad = pad_mask[step_mask].all() + step_is_pad.append(is_pad) + + is_ready = ready_step[step - 1] + step_is_ready.append(is_ready) + + step_confidence = confidence[step_mask] + total_step_confidence = step_confidence.max(dim=-1).values.mean() + step_confidences.append(total_step_confidence) + + step_confidences = torch.tensor(step_confidences).to(self.model.device) + step_is_ready = torch.tensor(step_is_ready, device=self.model.device) + step_is_pad = torch.tensor(step_is_pad, device=self.model.device) + + non_ready_mask = ~step_is_ready + nr_steps = valid_steps[non_ready_mask] # 1-based step ID + nr_conf = step_confidences[non_ready_mask] + nr_pad = step_is_pad[non_ready_mask] + + if nr_steps.numel() == 0: + return ar_postfix_inputs_embeds, ready_step, ar_tokens + + chosen_steps = [] + used_mask = torch.zeros_like(nr_steps, dtype=torch.bool) + + # Step 1: choose one non-PAD step by confidence when available. + non_pad_mask = ~nr_pad + if non_pad_mask.any(): + non_pad_conf = nr_conf[non_pad_mask] + non_pad_steps = nr_steps[non_pad_mask] + + best_idx = torch.argmax(non_pad_conf) # choose highest confidence + chosen_step = non_pad_steps[best_idx] + chosen_steps.append(chosen_step.item()) + + # Mark this step as used. + used_mask[nr_steps == chosen_step] = True + + # If enough steps were selected, reuse chosen_steps below. + if len(chosen_steps) >= update_step: + chosen_steps = torch.tensor(chosen_steps, device=self.model.device) + # chosen_steps will be used below. + else: + # ============================================================ + # Step 2: choose PAD steps by descending index, ignoring confidence. + # ============================================================ + remaining = update_step - len(chosen_steps) + + pad_mask_only = nr_pad & (~used_mask) + pad_steps_only = nr_steps[pad_mask_only] + + if pad_steps_only.numel() > 0: + # Sort from right to left. + pad_sorted = torch.argsort(pad_steps_only, descending=True) + pick = pad_steps_only[pad_sorted[:remaining]] + + chosen_steps.extend(pick.tolist()) + used_mask[(nr_steps.unsqueeze(1) == pick).any(dim=-1)] = True + + # Continue to Step 3 if more steps are still needed. + if len(chosen_steps) < update_step: + remaining = update_step - len(chosen_steps) + + # ============================================================ + # Step 3: choose remaining non-PAD steps by confidence. + # ============================================================ + non_pad_left_mask = (~nr_pad) & (~used_mask) + + if non_pad_left_mask.any(): + left_conf = nr_conf[non_pad_left_mask] + left_steps = nr_steps[non_pad_left_mask] + + sorted_idx = torch.argsort(left_conf, descending=True) + pick = left_steps[sorted_idx[:remaining]] + + chosen_steps.extend(pick.tolist()) + used_mask[(nr_steps.unsqueeze(1) == pick).any(dim=-1)] = True + + chosen_steps = torch.tensor(chosen_steps, device=self.model.device) + for i, step in enumerate(valid_steps): + if ready_step[step - 1] and not remask: + continue + step_mask = ar_action_mask == step + step_indices = step_mask.nonzero(as_tuple=True)[1] + if step in chosen_steps: + # update + ar_postfix_inputs_embeds[0, step_indices, :] = self.model.embed_tokens( + step_token_ids[i] + ) + if ar_tokens is not None: + ar_tokens[step_indices] = step_token_ids[i] + ready_step[step - 1] = True + elif remask and ready_step[step - 1]: + # back to placeholder + ar_postfix_inputs_embeds[0, step_indices, :] = placeholder_seq_embed[ + : len(step_indices) + ] + ready_step[step - 1] = False + + return ar_postfix_inputs_embeds, ready_step, ar_tokens + + def update_infer_dllm_position_mask(self, model_input): + positional_masks = self.tokenizer_mixin.update_placeholder_mask( + self.processor, + model_input["prefix_length"], + model_input["input_ids"], + ) + model_input["positional_masks"] = positional_masks + model_input["positional_masks"]["ar_visible"] = True + return model_input + + @torch.no_grad() + def generate_flow_action_no_cache( + self, + input_ids, + action_horizon, + action_dim, + num_inference_timesteps: int = 10, + prefix_length: Optional[int] = None, # Prefix length. + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + positional_masks: Optional[torch.LongTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + action_chunk: Optional[torch.FloatTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, + unnorm_proprioception: Optional[torch.FloatTensor] = None, + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + dataset_names: Optional[str] = None, + dof_mask: Optional[torch.FloatTensor] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + **kwargs, + ): + # assert self.config._attn_implementation == "sdpa", "generate_flow_action only support sdpa attn implementation" + batch_size = ( + input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] + ) + assert ( + batch_size == 1 + ), "generate_flow_action_no_cache only support batch size 1" + + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_features = image_embeds.shape[0] + if n_image_tokens != n_image_features: + raise ValueError( + f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" + ) + + mask = input_ids == self.config.image_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + image_mask = mask_expanded.to(inputs_embeds.device) + + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + + mask = input_ids == self.config.video_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + video_mask = mask_expanded.to(inputs_embeds.device) + + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + if proprioception is not None and not getattr( + self.config, "use_state_string_representation", False + ): + proprioception = proprioception.to(inputs_embeds.device) + agent_pos_mask = agent_pos_mask.to(inputs_embeds.device) + proprio_embed = self.action_preprocessor.proprioception_proj( + proprioception, + dataset_names, + agent_pos_mask, + use_history=proprioception.shape[1] > 1, + ) + proprioception_mask = ( + input_ids == self.action_token_id_set["propri_token_id"] + ) + inputs_embeds[proprioception_mask] = proprio_embed.reshape( + -1, inputs_embeds.shape[-1] + ).to(inputs_embeds.dtype) + + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2 + ): + # calculate RoPE index once per generation in the pre-fill stage only + if ( + (cache_position is not None and cache_position[0] == 0) + or self.rope_deltas is None + or (past_key_values is None or past_key_values.get_seq_length() == 0) + ): + position_ids, rope_deltas = self.get_rope_index( + input_ids, + image_grid_thw, + video_grid_thw, + second_per_grid_ts, + attention_mask, + ) + self.rope_deltas = rope_deltas + # then use the prev pre-calculated rope-deltas to get the correct position ids + else: + batch_size, seq_length, _ = inputs_embeds.shape + delta = ( + (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = position_ids.view(1, -1).expand(batch_size, -1) + if cache_position is not None: # otherwise `deltas` is an int `0` + delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + position_ids = position_ids.add(delta) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + if start_indices is None or end_indices is None: + # Compute the token span for each expert group after permutation. + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + + if action_chunk is not None: + action_chunk = action_chunk.to(inputs_embeds.device).to(torch.float32) + + output = {} + + noise = torch.randn( + size=(batch_size, action_horizon, action_dim), + dtype=torch.float32, + device=inputs_embeds.device, + ) + noisy_action = noise.clone() + dof_mask = dof_mask.to(inputs_embeds.device).to(torch.float32) + + times = self.action_preprocessor.get_inference_times( + num_inference_timesteps, inputs_embeds.device, torch.float32 + ) + flow_action_mask = input_ids == self.action_token_id_set["action_token_id"] + + if not dof_mask.all(): + padding_action = ( + torch.zeros((1, dof_mask.shape[-1])) + .to(dof_mask.device) + .to(torch.float32) + ) + padding_action = self.action_preprocessor.normalizer_action.normalize_data( + padding_action, dataset_names + ) + v_padding = padding_action - noisy_action + + def step(timestep, noisy_action): + timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0]) + action_embed, adarms_cond = self.action_preprocessor.step( + timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask + ) + action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]).to( + inputs_embeds.dtype + ) + + inputs_embeds[flow_action_mask] = action_embed + model_output = self.model( + input_ids=None, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=None, + inputs_embeds=inputs_embeds, + moe_token_types=moe_token_types, + start_indices=start_indices, + end_indices=end_indices, + positional_masks=positional_masks, + use_cache=True, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + adarms_conds=[None, adarms_cond], + ) + + hidden_states = model_output.last_hidden_state + action_hidden_states = hidden_states[flow_action_mask].to(torch.float32) + action_pred = self.action_preprocessor.action_proj_back( + action_hidden_states[:, : self.action_preprocessor.action_hidden_size] + ) + if getattr(self.config, "use_x_pred", False): + v_t = (action_pred - noisy_action) / torch.clamp(1 - timestep, min=0.05) + else: + v_t = action_pred + + if not dof_mask.all(): + v_t = (v_padding) * (1 - dof_mask) + v_t * dof_mask + + return v_t.reshape(batch_size, action_horizon, action_dim) + + action_trajectory = odeint(step, noisy_action, times, method="euler") + + predict_action = action_trajectory[-1] + predict_action = self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names + ) + output["predict_action"] = predict_action + # normalize action chunk to get gt_action + if action_chunk is not None: + output["gt_action"] = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + action_chunk, dataset_names + ) + ) + + return output diff --git a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py similarity index 82% rename from wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py rename to wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py index 951bed9c..7e6bcb68 100644 --- a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py +++ b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl.py @@ -1,18 +1,51 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_qwen2_5_vl.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# coding=utf-8 +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import math +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + import torch import torch.nn as nn import torch.nn.functional as F -from dataclasses import dataclass +import torch.utils.checkpoint as cp from torch.nn import CrossEntropyLoss -from typing import Any, Dict, List, Optional, Tuple, Union from transformers.activations import ACT2FN from transformers.cache_utils import ( Cache, DynamicCache, - SlidingWindowCache, StaticCache, ) + +try: + from transformers.cache_utils import SlidingWindowCache +except ImportError: + # Compatibility with newer transformers versions + SlidingWindowCache = StaticCache from transformers.generation import GenerationMixin from transformers.modeling_attn_mask_utils import AttentionMaskConverter from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput @@ -27,17 +60,20 @@ logging, replace_return_docstrings, ) + + from .configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig -from wall_x.fusions import ops +from wall_x.model.core.attention.selector import AttentionsSelectorMixin +from wall_x.model.core.ops import rot_pos_emb, get_window_index, m_rope if is_flash_attn_2_available(): + from flash_attn import flash_attn_func from flash_attn import flash_attn_varlen_func from flash_attn.layers.rotary import apply_rotary_emb - from flash_attn import flash_attn_func + else: flash_attn_varlen_func = None apply_rotary_emb = None - flash_attn_func = None if is_flash_attn_2_available(): @@ -46,26 +82,142 @@ flash_attn_varlen_func = None +try: + from flash_mask.flash_mask_interface import flash_mask_attn_func +except ImportError: + flash_mask_attn_func = None + + logger = logging.get_logger(__name__) _CONFIG_FOR_DOC = "Qwen2_5_VLConfig" class Qwen2_5_VLMLP(nn.Module): - def __init__(self, config, bias: bool = False): + # Align intermediate_size to this value for cuBLAS GEMM tile efficiency. + # E.g. 3420 -> 3456 (= 128 * 27), improving MFU from 34.8% to ~60%+. + # Zero-padded weights don't affect forward output or gradient flow. + _GEMM_ALIGN = 128 + + def __init__(self, config, bias=False, use_selective_recompute=False): super().__init__() self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias) + + # Pad intermediate_size to next multiple of _GEMM_ALIGN + align = self._GEMM_ALIGN + self.padded_intermediate_size = ( + (self.intermediate_size + align - 1) // align * align + ) + + self.gate_up_proj = nn.Linear( + self.hidden_size, 2 * self.padded_intermediate_size, bias=bias + ) + self.down_proj = nn.Linear( + self.padded_intermediate_size, self.hidden_size, bias=bias + ) + self.act_fn = ACT2FN[config.hidden_act] + self.use_selective_recompute = use_selective_recompute + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + """Pad checkpoint weights from original intermediate_size to padded size.""" + pad = self.padded_intermediate_size - self.intermediate_size + if pad > 0: + orig = self.intermediate_size + # gate_up_proj.weight: [2*orig, hidden] -> [2*padded, hidden] + key_w = prefix + "gate_up_proj.weight" + if key_w in state_dict and state_dict[key_w].shape[0] == 2 * orig: + w = state_dict[key_w] + gate_w, up_w = w.split([orig, orig], dim=0) + gate_w = F.pad(gate_w, (0, 0, 0, pad)) + up_w = F.pad(up_w, (0, 0, 0, pad)) + state_dict[key_w] = torch.cat([gate_w, up_w], dim=0) + # gate_up_proj.bias: [2*orig] -> [2*padded] + key_b = prefix + "gate_up_proj.bias" + if key_b in state_dict and state_dict[key_b].shape[0] == 2 * orig: + b = state_dict[key_b] + gate_b, up_b = b.split([orig, orig]) + state_dict[key_b] = torch.cat( + [F.pad(gate_b, (0, pad)), F.pad(up_b, (0, pad))] + ) + # down_proj.weight: [hidden, orig] -> [hidden, padded] + key_dw = prefix + "down_proj.weight" + if key_dw in state_dict and state_dict[key_dw].shape[1] == orig: + state_dict[key_dw] = F.pad(state_dict[key_dw], (0, pad)) + # down_proj.bias: [hidden] — no change needed + + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def _save_to_state_dict(self, destination, prefix, keep_vars): + """Strip padding so saved checkpoints use the original intermediate_size.""" + super()._save_to_state_dict(destination, prefix, keep_vars) + pad = self.padded_intermediate_size - self.intermediate_size + if pad == 0: + return + orig = self.intermediate_size + padded = self.padded_intermediate_size + + # gate_up_proj.weight: [2*padded, hidden] → [2*orig, hidden] + key_w = prefix + "gate_up_proj.weight" + if key_w in destination: + w = destination[key_w] + gate_w, up_w = w.split([padded, padded], dim=0) + destination[key_w] = torch.cat([gate_w[:orig], up_w[:orig]], dim=0) + + # gate_up_proj.bias: [2*padded] → [2*orig] + key_b = prefix + "gate_up_proj.bias" + if key_b in destination: + b = destination[key_b] + gate_b, up_b = b.split([padded, padded]) + destination[key_b] = torch.cat([gate_b[:orig], up_b[:orig]]) + + # down_proj.weight: [hidden, padded] → [hidden, orig] + key_dw = prefix + "down_proj.weight" + if key_dw in destination: + destination[key_dw] = destination[key_dw][:, :orig] + + def _activation_chunk(self, gate_out, up_out): + # gate_out / up_out are precomputed GEMM results + act = self.act_fn(gate_out) + return act * up_out def forward(self, hidden_state): - return self.down_proj( - self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state) + # ----------------------------- + # 1. Precompute GEMM without checkpointing + # ----------------------------- + gate_up_out = self.gate_up_proj(hidden_state) + gate_out, up_out = gate_up_out.split( + [self.padded_intermediate_size, self.padded_intermediate_size], dim=-1 ) + if self.use_selective_recompute: + # checkpoint only activation, not GEMM + out = cp.checkpoint( + self._activation_chunk, gate_out, up_out, use_reentrant=False + ) + else: + out = self._activation_chunk(gate_out, up_out) + + return self.down_proj(out) + class Qwen2_5_VisionPatchEmbed(nn.Module): def __init__( @@ -99,9 +251,18 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: self.patch_size, self.patch_size, ) - hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view( - -1, self.embed_dim + # hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view( + # -1, self.embed_dim + # ) + weight = self.proj.weight.view(self.embed_dim, -1) + out = hidden_states.view( + -1, + self.in_channels + * self.temporal_patch_size + * self.patch_size + * self.patch_size, ) + hidden_states = F.linear(out.to(target_dtype), weight) return hidden_states @@ -196,7 +357,8 @@ def __init__(self, dim: int, context_dim: int, spatial_merge_size: int = 2) -> N ) def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.mlp(self.ln_q(x)[0].view(-1, self.hidden_size)) + normed_x, _ = self.ln_q(x) + x = self.mlp(normed_x.view(-1, self.hidden_size)) return x @@ -211,11 +373,12 @@ def apply_rotary_pos_emb_flashatt( class Qwen2_5_VLVisionFlashAttention2(nn.Module): - def __init__(self, dim: int, num_heads: int = 16) -> None: + def __init__(self, config: Qwen2_5_VLConfig, dim: int, num_heads: int = 16) -> None: super().__init__() self.num_heads = num_heads self.qkv = nn.Linear(dim, dim * 3, bias=True) self.proj = nn.Linear(dim, dim) + self.deterministic = config.attn_deterministic def forward( self, @@ -251,7 +414,14 @@ def forward( if max_seqlen is None: max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() attn_output = flash_attn_varlen_func( - q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen + q, + k, + v, + cu_seqlens, + cu_seqlens, + max_seqlen, + max_seqlen, + deterministic=self.deterministic, ).reshape(seq_length, -1) attn_output = self.proj(attn_output) return attn_output @@ -279,7 +449,7 @@ def apply_rotary_pos_emb_vision( class Qwen2_5_VLVisionAttention(nn.Module): - def __init__(self, dim: int, num_heads: int = 16) -> None: + def __init__(self, config: Qwen2_5_VLConfig, dim: int, num_heads: int = 16) -> None: super().__init__() self.num_heads = num_heads self.head_dim = dim // num_heads @@ -290,9 +460,9 @@ def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, - max_seqlen: Optional[int] = None, rotary_pos_emb: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, ) -> torch.Tensor: seq_length = hidden_states.shape[0] q, k, v = ( @@ -344,7 +514,7 @@ def forward( class Qwen2_5_VLVisionSdpaAttention(nn.Module): - def __init__(self, dim: int, num_heads: int = 16) -> None: + def __init__(self, config: Qwen2_5_VLConfig, dim: int, num_heads: int = 16) -> None: super().__init__() self.num_heads = num_heads self.qkv = nn.Linear(dim, dim * 3, bias=True) @@ -354,7 +524,6 @@ def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, - max_seqlen: Optional[int] = None, rotary_pos_emb: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: @@ -408,14 +577,33 @@ def forward( class Qwen2_5_VLVisionBlock(nn.Module): - def __init__(self, config, attn_implementation: str = "sdpa") -> None: + def __init__( + self, + config, + attn_implementation: str = "sdpa", + use_selective_recompute: bool = False, + ): super().__init__() + self.use_selective_recompute = use_selective_recompute + self.norm1 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) self.norm2 = Qwen2RMSNorm(config.hidden_size, eps=1e-6) + self.attn = QWEN2_5_VL_VISION_ATTENTION_CLASSES[attn_implementation]( - config.hidden_size, num_heads=config.num_heads + config, config.hidden_size, num_heads=config.num_heads + ) + self.mlp = Qwen2_5_VLMLP( + config, bias=True, use_selective_recompute=use_selective_recompute ) - self.mlp = Qwen2_5_VLMLP(config, bias=True) + + # ----------------------------- + # selective checkpoint for norm + # ----------------------------- + def _norm1_chunk(self, hidden_states): + return self.norm1(hidden_states)[0] # return only normed_hidden_states + + def _norm2_chunk(self, hidden_states): + return self.norm2(hidden_states)[0] def forward( self, @@ -425,14 +613,46 @@ def forward( rotary_pos_emb: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: + + # ============================ + # 1) Norm1 with checkpoint + # ============================ + if self.use_selective_recompute: + normed_hidden_states = cp.checkpoint( + self._norm1_chunk, + hidden_states, + use_reentrant=False, + ) + # residual RMS is not used by subsequent computation and is usually unused in large-model training + _ = None + else: + normed_hidden_states, _ = self.norm1(hidden_states) + + # Attention does not recompute GEMM hidden_states = hidden_states + self.attn( - self.norm1(hidden_states)[0], + normed_hidden_states, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, rotary_pos_emb=rotary_pos_emb, position_embeddings=position_embeddings, ) - hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)[0]) + + # ============================ + # 2) Norm2 with checkpoint + # ============================ + if self.use_selective_recompute: + normed_hidden_states = cp.checkpoint( + self._norm2_chunk, + hidden_states, + use_reentrant=False, + ) + _ = None + else: + normed_hidden_states, _ = self.norm2(hidden_states) + + # MLP only recomputes the hooked activation and does not recompute GEMM + hidden_states = hidden_states + self.mlp(normed_hidden_states) + return hidden_states @@ -457,7 +677,7 @@ def forward( "The bare Qwen2_5_VL Model outputting raw hidden-states without any specific head on top.", Qwen2_5_VL_START_DOCSTRING, ) -class Qwen2_5_VLPreTrainedModel(PreTrainedModel): +class Qwen2_5_VLPreTrainedModel(AttentionsSelectorMixin, PreTrainedModel): config_class = Qwen2_5_VLConfig base_model_prefix = "model" supports_gradient_checkpointing = True @@ -484,8 +704,11 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): config_class = Qwen2_5_VLVisionConfig _no_split_modules = ["Qwen2_5_VLVisionBlock"] - def __init__(self, config, *inputs, **kwargs) -> None: + def __init__( + self, config, use_selective_recompute=False, *inputs, **kwargs + ) -> None: super().__init__(config, *inputs, **kwargs) + self.use_selective_recompute = use_selective_recompute self.spatial_merge_size = config.spatial_merge_size self.patch_size = config.patch_size self.fullatt_block_indexes = config.fullatt_block_indexes @@ -504,7 +727,9 @@ def __init__(self, config, *inputs, **kwargs) -> None: self.blocks = nn.ModuleList( [ - Qwen2_5_VLVisionBlock(config, config._attn_implementation) + Qwen2_5_VLVisionBlock( + config, config._attn_implementation, use_selective_recompute + ) for _ in range(config.depth) ] ) @@ -605,12 +830,13 @@ def forward( `torch.Tensor`: hidden_states. """ hidden_states = self.patch_embed(hidden_states) - rotary_pos_emb = ops.rot_pos_emb( - self.rotary_pos_emb.inv_freq, grid_thw, self.spatial_merge_size + rotary_pos_emb = rot_pos_emb( + self.rotary_pos_emb.inv_freq.to(torch.float32), + grid_thw, + self.spatial_merge_size, ) - - window_index, cu_window_seqlens = ops.get_window_index( - grid_thw=grid_thw, + window_index, cu_window_seqlens = get_window_index( + grid_thw=grid_thw.to(torch.int32), window_size=self.window_size, spatial_merge_size=self.spatial_merge_size, patch_size=self.patch_size, @@ -658,9 +884,9 @@ def forward( hidden_states = self._gradient_checkpointing_func( blk.__call__, hidden_states, - cu_seqlens_now, - None, - position_embeddings, + cu_seqlens=cu_seqlens_now, + max_seqlen=max_seqlen_now, + position_embeddings=position_embeddings, ) else: hidden_states = blk( @@ -683,7 +909,7 @@ def __init__(self, config: Qwen2_5_VLConfig, device=None): # BC: "rope_type" was originally "type" if hasattr(config, "rope_scaling") and config.rope_scaling is not None: self.rope_type = config.rope_scaling.get( - "rope_type", config.rope_scaling.get("type") + "rope_type", config.rope_scaling.get("type", "default") ) else: self.rope_type = "default" @@ -691,12 +917,43 @@ def __init__(self, config: Qwen2_5_VLConfig, device=None): self.original_max_seq_len = config.max_position_embeddings self.config = config - self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + # Compatibility with newer transformers versions (5.x): + # Newer versions remove the "default" key from ROPE_INIT_FUNCTIONS and compute default RoPE inline. + # Older versions (4.x): ROPE_INIT_FUNCTIONS["default"] exists and can be used normally. + # Newer versions (5.x): when rope_type=="default", use compute_default_rope_parameters instead. + if self.rope_type == "default": + self.rope_init_fn = self.compute_default_rope_parameters + else: + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = self.inv_freq + @staticmethod + def compute_default_rope_parameters( + config: Qwen2_5_VLConfig, + device=None, + seq_len=None, + ): + """Compute default RoPE parameters, matching older ROPE_INIT_FUNCTIONS["default"]. + Transformers 5.x moved this logic out of ROPE_INIT_FUNCTIONS and into the model class. + """ + base = config.rope_theta + dim = config.hidden_size // config.num_attention_heads + attention_factor = 1.0 + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, dim, 2, dtype=torch.int64).to( + device=device, dtype=torch.float + ) + / dim + ) + ) + return inv_freq, attention_factor + def _dynamic_frequency_update(self, position_ids, device): """ dynamic RoPE layers should recompute `inv_freq` in the following situations: @@ -763,13 +1020,18 @@ def __init__(self, config): self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.gate_up_proj = nn.Linear( + self.hidden_size, 2 * self.intermediate_size, bias=False + ) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + gate_up_out = self.gate_up_proj(x) + gate_out, up_out = gate_up_out.split( + [self.intermediate_size, self.intermediate_size], dim=-1 + ) + down_proj = self.down_proj(self.act_fn(gate_out) * up_out) return down_proj @@ -863,15 +1125,13 @@ def __init__(self, config: Qwen2_5_VLConfig, layer_idx: Optional[int] = None): f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" f" and `num_heads`: {self.num_heads})." ) - self.q_proj = nn.Linear( - self.hidden_size, self.num_heads * self.head_dim, bias=True - ) - self.k_proj = nn.Linear( - self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True - ) - self.v_proj = nn.Linear( - self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True + + qkv_dim = ( + self.num_heads * self.head_dim + + 2 * self.num_key_value_heads * self.head_dim ) + self.qkv_proj = nn.Linear(self.hidden_size, qkv_dim, bias=True) + self.o_proj = nn.Linear( self.num_heads * self.head_dim, self.hidden_size, bias=False ) @@ -890,16 +1150,22 @@ def forward( position_embeddings: Optional[ Tuple[torch.Tensor, torch.Tensor] ] = None, # necessary, but kept here for BC + **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + qkv = self.qkv_proj(hidden_states) # [bsz, q_len, q_dim + 2 * kv_dim] + kv_dim = self.num_key_value_heads * self.head_dim + q, k, v = torch.split( + qkv, [self.num_heads * self.head_dim, kv_dim, kv_dim], dim=-1 + ) + query_states = q.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = k.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = v.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_multimodal_rotary_pos_emb( @@ -970,13 +1236,14 @@ class Qwen2_5_VLFlashAttention2(Qwen2_5_VLAttention): config.max_window_layers layers. """ - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, config: Qwen2_5_VLConfig, layer_idx: Optional[int] = None): + super().__init__(config, layer_idx) # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + self.deterministic = config.attn_deterministic def forward( self, @@ -993,19 +1260,25 @@ def forward( ): bsz, q_len, _ = hidden_states.size() - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) + qkv = self.qkv_proj(hidden_states) # [bsz, q_len, q_dim + 2 * kv_dim] + kv_dim = self.num_key_value_heads * self.head_dim + q, k, v = torch.split( + qkv, [self.num_heads * self.head_dim, kv_dim, kv_dim], dim=-1 + ) - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + query_states = q.view(bsz, q_len, self.num_heads, self.head_dim) + key_states = k.view(bsz, q_len, self.num_key_value_heads, self.head_dim) + value_states = v.view(bsz, q_len, self.num_key_value_heads, self.head_dim) # Because the input can be padded, the absolute sequence length depends on the max position id. cos, sin = position_embeddings - query_states, key_states = ops.multimodal_rope( - query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] + query_states, key_states = m_rope( + query_states.contiguous(), + key_states.contiguous(), + cos[..., : (cos.size(-1) // 2)].contiguous().float(), + sin[..., : (sin.size(-1) // 2)].contiguous().float(), + self.rope_scaling["mrope_section"], ) if past_key_value is not None: cache_kwargs = { @@ -1014,8 +1287,14 @@ def forward( "cache_position": cache_position, } # Specific to RoPE models key_states, value_states = past_key_value.update( - key_states, value_states, self.layer_idx, cache_kwargs + key_states.transpose(1, 2), + value_states.transpose(1, 2), + self.layer_idx, + cache_kwargs, ) + key_states, value_states = key_states.transpose( + 1, 2 + ), value_states.transpose(1, 2) # repeat k/v heads if n_kv_heads < n_heads # key_states = repeat_kv(key_states, self.num_key_value_groups) @@ -1033,7 +1312,7 @@ def forward( elif hasattr(self.config, "_pre_quantization_dtype"): target_dtype = self.config._pre_quantization_dtype else: - target_dtype = self.q_proj.weight.dtype + target_dtype = self.qkv_proj.weight.dtype logger.warning_once( f"The input hidden states seems to be silently casted in float32, this might be related to" @@ -1045,11 +1324,6 @@ def forward( key_states = key_states.to(target_dtype) value_states = value_states.to(target_dtype) - # Reashape to the expected shape for Flash Attention - query_states = query_states.transpose(1, 2) - key_states = key_states.transpose(1, 2) - value_states = value_states.transpose(1, 2) - attn_output = flash_attn_func( query_states, key_states, @@ -1057,6 +1331,7 @@ def forward( dropout_rate, softmax_scale=None, causal=self.is_causal, + deterministic=self.deterministic, ) attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() @@ -1064,6 +1339,39 @@ def forward( if not output_attentions: attn_weights = None + if output_attentions: + with torch.no_grad(): + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + key_states = repeat_kv(key_states, self.num_key_value_groups) + scale = 1.0 / torch.sqrt( + torch.tensor( + self.head_dim, device=hidden_states.device, dtype=torch.float32 + ) + ) + attention_score = ( + torch.matmul(query_states, key_states.transpose(-2, -1)) * scale + ) + causal_mask_local = torch.tril( + torch.ones( + attention_score.size(-2), + attention_score.size(-1), + device=attention_score.device, + dtype=torch.bool, + ) + ) + mask = ( + causal_mask_local.unsqueeze(0) + .unsqueeze(0) + .expand(attention_score.size(0), attention_score.size(1), -1, -1) + ) + if mask.dtype != attention_score.dtype: + mask = mask.to(dtype=attention_score.dtype) + attention_score = attention_score.masked_fill( + ~(mask.bool()), float("-inf") + ) + attention_score = torch.softmax(attention_score, dim=-1) + attn_weights = attention_score[0].mean(0) return attn_output, attn_weights, past_key_value @@ -1108,13 +1416,18 @@ def forward( bsz, q_len, _ = hidden_states.size() - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + qkv = self.qkv_proj(hidden_states) # [bsz, q_len, q_dim + 2 * kv_dim] + kv_dim = self.num_key_value_heads * self.head_dim + q, k, v = torch.split( + qkv, [self.num_heads * self.head_dim, kv_dim, kv_dim], dim=-1 + ) + query_states = q.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = k.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = v.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_multimodal_rotary_pos_emb( @@ -1132,7 +1445,18 @@ def forward( key_states, value_states, self.layer_idx, cache_kwargs ) else: - past_key_states, past_value_states = past_key_value[self.layer_idx] + # Compatible across transformers versions: + # v5.x: DynamicCache uses .layers[idx].keys/.values + # v4.x: DynamicCache uses .key_cache[idx]/.value_cache[idx] + # old: Cache object is subscriptable, returns (key, value) tuple + if hasattr(past_key_value, "layers"): + past_key_states = past_key_value.layers[self.layer_idx].keys + past_value_states = past_key_value.layers[self.layer_idx].values + elif hasattr(past_key_value, "key_cache"): + past_key_states = past_key_value.key_cache[self.layer_idx] + past_value_states = past_key_value.value_cache[self.layer_idx] + else: + past_key_states, past_value_states = past_key_value[self.layer_idx] key_states = torch.cat([past_key_states, key_states], dim=-2) value_states = torch.cat([past_value_states, value_states], dim=-2) @@ -1141,15 +1465,15 @@ def forward( causal_mask = attention_mask if attention_mask is not None: # no matter the length, we just slice it - # Ensure the attention_mask correctly matches the head dimension + # Ensure attention_mask matches the head dimension if len(attention_mask.shape) == 2: # [batch_size, seq_len] - # Expand to [batch_size, 1, seq_len, seq_len] causal mask format + # Expand to causal mask format [batch_size, 1, seq_len, seq_len] bsz, seq_len = attention_mask.shape causal_mask = attention_mask.view(bsz, 1, 1, seq_len).expand( bsz, 1, seq_len, seq_len ) elif len(attention_mask.shape) == 3: # [batch_size, seq_len, seq_len] - # Add head dimension: [batch_size, 1, seq_len, seq_len] + # Add the head dimension: [batch_size, 1, seq_len, seq_len] causal_mask = attention_mask.unsqueeze(1) elif ( len(attention_mask.shape) == 4 @@ -1157,10 +1481,10 @@ def forward( causal_mask = attention_mask else: raise ValueError( - f"Unsupported attention_mask dim: {attention_mask.shape}" + f"Unsupported attention_mask shape: {attention_mask.shape}" ) - # Convert the attention mask to boolean type + # Convert attention mask to bool causal_mask = causal_mask.to(torch.bool) # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, @@ -1204,10 +1528,132 @@ def forward( return attn_output, None, past_key_value +class Qwen2_5_VLFlashMaskAttention(Qwen2_5_VLAttention): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[ + Tuple[torch.Tensor, torch.Tensor] + ] = None, # necessary, but kept here for BC + ): + bsz, q_len, _ = hidden_states.size() + + qkv = self.qkv_proj(hidden_states) # [bsz, q_len, q_dim + 2 * kv_dim] + kv_dim = self.num_key_value_heads * self.head_dim + q, k, v = torch.split( + qkv, [self.num_heads * self.head_dim, kv_dim, kv_dim], dim=-1 + ) + query_states = q.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = k.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = v.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + query_states, key_states = apply_multimodal_rotary_pos_emb( + query_states, key_states, cos, sin, self.rope_scaling["mrope_section"] + ) + + if past_key_value is not None: + cache_kwargs = { + "sin": sin, + "cos": cos, + "cache_position": cache_position, + } # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + # dropout_rate = 0.0 if not self.training else self.attention_dropout + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in float16 just to be sure everything works as expected. + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.qkv_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + # Reashape to the expected shape for Flash Attention + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + # if ( + # self.config.use_sliding_window + # and getattr(self.config, "sliding_window", None) is not None + # and self.layer_idx >= self.config.max_window_layers + # ): + # sliding_window = self.config.sliding_window + # else: + # sliding_window = None + + # Expand the attention_mask head dimension from 1 to num_heads + if attention_mask is not None and attention_mask.shape[1] == 1: + # attention_mask: (bsz, 1, q_len, 2) or (bsz, 1, q_len, 4) + attention_mask = attention_mask.expand( + -1, self.num_heads, -1, -1 + ).contiguous() + + if attention_mask is not None: + attention_mask = attention_mask.contiguous() + + query_states_contiguous = query_states.contiguous() + key_states_contiguous = key_states.contiguous() + value_states_contiguous = value_states.contiguous() + + attn_output = flash_mask_attn_func( + query_states_contiguous, + key_states_contiguous, + value_states_contiguous, + startend_row_indices=attention_mask, + causal=False, + ) + + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + QWEN2_5_VL_ATTENTION_CLASSES = { "eager": Qwen2_5_VLAttention, "flash_attention_2": Qwen2_5_VLFlashAttention2, "sdpa": Qwen2_5_VLSdpaAttention, + "flash_mask": Qwen2_5_VLFlashMaskAttention, } @@ -1765,14 +2211,16 @@ class Qwen2_5_VLCausalLMOutputWithPast(ModelOutput): class Qwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel, GenerationMixin): - _tied_weights_keys = ["lm_head.weight"] + # Compatibility with newer transformers versions (5.x):_tied_weights_keys changed from list[str] to dict[str, str] (target -> source mapping). + # Older versions (4.x) iterate over dict keys, so the behavior is equivalent and compatible. + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} config_class = Qwen2_5_VLConfig _no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"] - def __init__(self, config): + def __init__(self, config, use_selective_recompute): super().__init__(config) self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config( - config.vision_config + config.vision_config, use_selective_recompute ) self.model = Qwen2_5_VLModel(config) self.vocab_size = config.vocab_size diff --git a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py similarity index 60% rename from wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py rename to wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py index 996f8799..b9de74f7 100644 --- a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py +++ b/wall_x/model/qact/qwen2_5/modeling_qwen2_5_vl_act.py @@ -1,61 +1,80 @@ -import os +import re import torch -import yaml -import numpy as np -import glob import torch.nn as nn -from torchdiffeq import odeint from dataclasses import dataclass from torch.nn import CrossEntropyLoss -from safetensors.torch import load_file -from peft import LoraConfig, get_peft_model -from typing import Optional, List, Tuple, Any, Dict, Union -import time -from transformers import AutoConfig, AutoProcessor -from transformers.utils import logging, is_torchdynamo_compiling +from typing import Optional, List, Tuple +from torchdiffeq import odeint +import numpy as np +from typing import Any, Dict, Union +from wall_x.model.qact.qwen2_5.configuration_qwen2_5_vl import Qwen2_5_VLConfig +from wall_x.model.qact.qwen2_5.modeling_qwen2_5_vl import ( + QWEN2_5_VL_ATTENTION_CLASSES, + Qwen2_5_VLMLP, + Qwen2_5_VLRotaryEmbedding, + Qwen2_5_VLPreTrainedModel, + Qwen2RMSNorm, + Qwen2_5_VLPatchMerger, + Qwen2_5_VLVisionBlock, + Qwen2_5_VisionTransformerPretrainedModel, +) + +from transformers import AutoConfig + from transformers.cache_utils import ( Cache, DynamicCache, - SlidingWindowCache, StaticCache, ) -from transformers.modeling_attn_mask_utils import AttentionMaskConverter -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import ( - Qwen2_5_VLMLP, - Qwen2_5_VLRotaryEmbedding, - Qwen2_5_VLPreTrainedModel, - Qwen2RMSNorm, - Qwen2_5_VLForConditionalGeneration, -) +try: + from transformers.cache_utils import SlidingWindowCache +except ImportError: + # Compatibility with newer transformers versions + SlidingWindowCache = StaticCache + from transformers.modeling_outputs import ( - ModelOutput, BaseModelOutputWithPast, + ModelOutput, ) - -from wall_x.fusions import ops -from wall_x.model.action_head import ActionProcessor -from wall_x.model.qwen2_5_based.configuration_qwen2_5_vl import Qwen2_5_VLConfig -from wall_x.model.model_utils import load_wallx_processors, update_model_config -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import ( - Qwen2_5_VisionTransformerPretrainedModel, - Qwen2_5_VLAttention, - Qwen2_5_VLFlashAttention2, - Qwen2_5_VLSdpaAttention, +from transformers.utils import ( + logging, + is_torchdynamo_compiling, +) +from transformers.modeling_attn_mask_utils import ( + AttentionMaskConverter, +) +from wall_x.model.core.vla_mixin import ActionGenerationMixin, ActionModelMixMin +from wall_x.model.core.action.normalizer import ( + normalize_data_with_virtual_tail, + unnormalize_data_with_virtual_tail, ) -from wall_x.model.vla_mixin import ActionGenerationMixin, ActionModelMixMin -from wall_x.model.vla_mixin import TokenTypeRouter, SparseMoeBlock -from wall_x.model.vla_mixin import ( +from wall_x.model.core.action.moe import TokenTypeRouter, SparseMoeBlock +from wall_x.model.core.attention.selector import ( ATTENTION_TYPES_WITH_2D_MASK, + ATTENTION_TYPES_WITH_FLASH_MASK, ) -from wall_x.model.joint_attention import JOINT_QWEN_ATTENTION_CLASSES +from wall_x.model.core.attention.joint import JOINT_QWEN_ATTENTION_CLASSES +from wall_x.model.qact.qwen2_5.inference_mixin import VLAInferenceMixin -from wall_x.data.utils import update_action_statistics -from wall_x.utils.constant import action_statistic_dof -from pprint import pprint +from wall_x.model.core.ops import unpermute, permute, get_rope_index logger = logging.get_logger(__name__) +_QWEN25_DMUON_BLOCKED_NAME_PARTS = ( + "embed_tokens", + "pos_embed", + "lm_head", + "norm", +) + + +def _is_qwen25_dmuon_target_param(name: str, param: nn.Parameter) -> bool: + """Select trainable Qwen2.5 VLA matrices for DMuon.""" + if not param.requires_grad or param.ndim != 2 or not name.endswith(".weight"): + return False + return not any(part in name for part in _QWEN25_DMUON_BLOCKED_NAME_PARTS) + @dataclass class Qwen2_5_VLACausalLMOutputWithPast(ModelOutput): @@ -72,13 +91,6 @@ class Qwen2_5_VLACausalLMOutputWithPast(ModelOutput): channel_loss_count_dict: Optional[dict[torch.FloatTensor]] = None -QWEN2_5_VL_ATTENTION_CLASSES = { - "eager": Qwen2_5_VLAttention, - "flash_attention_2": Qwen2_5_VLFlashAttention2, - "sdpa": Qwen2_5_VLSdpaAttention, -} - - class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module, ActionModelMixMin): def __init__( self, @@ -203,9 +215,9 @@ def forward( Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code into the model """ - residual = hidden_states - hidden_states, gate, _ = self._apply_norm_moe( + residual = hidden_states + hidden_states, gate, gate_mask = self._apply_norm_moe( hidden_states, token_types, adarms_conds, @@ -286,7 +298,7 @@ class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin): def from_pretrained( cls, pretrained_model_name_or_path, num_experts=None, *args, **kwargs ): - # If `num_experts` is provided, ensure it is added to the config. + # If num_experts is provided, make sure it is added to config config = kwargs.get("config", None) if config is None: config = AutoConfig.from_pretrained(pretrained_model_name_or_path) @@ -360,10 +372,12 @@ def forward( position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - moe_token_types: Optional[torch.LongTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, # new parameter start_indices: Optional[torch.Tensor] = None, end_indices: Optional[torch.Tensor] = None, - positional_masks: Optional[dict] = None, + positional_masks: Optional[ + dict + ] = None, # stores token position masks needed by each category use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, @@ -393,7 +407,7 @@ def forward( "You must specify exactly one of input_ids or inputs_embeds" ) if moe_token_types is None: - raise ValueError("moe_token_types must be provided for MoE routing.") + raise ValueError("moe_token_types must be provided for MoE routing") if start_indices is None or end_indices is None: raise ValueError( "start_indices and end_indices must be provided for MoE routing" @@ -456,18 +470,16 @@ def forward( # create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) - # If `mot_opt` is enabled, the tokens from different experts will be permuted first, resulting in a dimension of [Tokens, HiddenSize]. + # When mot_opt is enabled, permute tokens from different experts first; shape becomes [Tokens, HiddenSize] orig_shape = hidden_states.shape + moe_token_types = moe_token_types.contiguous() if self.config.mot_opt: hidden_states = hidden_states.view(-1, hidden_states.size(-1)) - hidden_states, row_id_map = ops.permute( - hidden_states, moe_token_types.view(-1) - ) + hidden_states, row_id_map = permute(hidden_states, moe_token_types.view(-1)) else: row_id_map = None - probs = torch.ones_like(moe_token_types.view(-1), dtype=torch.float32).view( - -1, 1 - ) + # Use reshape instead of view to support non-contiguous tensors + probs = torch.ones_like(moe_token_types, dtype=torch.float32).reshape(-1, 1) # decoder layers all_hidden_states = () if output_hidden_states else None @@ -486,6 +498,16 @@ def forward( positional_masks=positional_masks, ) + if ( + self.config._attn_implementation in ATTENTION_TYPES_WITH_FLASH_MASK + and self.config.attention_moe is True + ): + causal_mask = self._update_joint_attention_flash_mask( + attention_mask=causal_mask, + moe_token_types=moe_token_types, + positional_masks=positional_masks, + ) + for decoder_layer in self.layers: if output_hidden_states: assert ( @@ -538,9 +560,6 @@ def forward( next_decoder_cache = layer_outputs[2 if output_attentions else 1] if output_attentions: - assert ( - self.config.mot_opt is False - ), "When using mot_opt, output_hidden_states is not supported yet." all_self_attns += (layer_outputs[1],) hidden_states, _, _ = self._apply_norm_moe( @@ -564,7 +583,7 @@ def forward( next_cache = next_decoder_cache if use_cache else None if self.config.mot_opt: - hidden_states = ops.unpermute(hidden_states, row_id_map, probs) + hidden_states = unpermute(hidden_states, row_id_map, probs) hidden_states = hidden_states.view(orig_shape) if not return_dict: @@ -658,22 +677,22 @@ def _update_causal_mask( config=self.config, past_key_values=past_key_values, ) - # Modify the mask to support bidirectional attention. + # Modify the mask to support bidirectional attention if moe_token_types is not None: - # Find the positions of all tokens of type 1. + # Find all token positions with type 1 type1_tokens = ( (moe_token_types == 1).unsqueeze(1).unsqueeze(2) ) # [B, 1, 1, S] - # Create a square mask for the type1 region. + # Create a square mask for the type-1 region type1_mask = torch.zeros_like(causal_mask) # [B, num_heads, S, S] type1_region = type1_tokens & type1_tokens.transpose(-1, -2) # [B, 1, S, S] type1_mask = type1_mask.masked_fill(type1_region, 1.0).to(torch.bool) - # Set the original causal_mask to zero in the type1 region, and then add the type1_mask. + # Zero the original causal_mask in the type-1 region, then add type1_mask causal_mask = torch.where( - type1_mask, - torch.zeros_like(causal_mask), - causal_mask, + type1_mask, # Expand dimensions to match causal_mask + torch.zeros_like(causal_mask), # zero the type-1 region + causal_mask, # keep other regions unchanged ) if ( self.config._attn_implementation == "sdpa" @@ -771,210 +790,59 @@ def _prepare_4d_causal_attention_mask_with_cache_position( class Qwen2_5_VLMoEForAction( - Qwen2_5_VLForConditionalGeneration, ActionGenerationMixin, ActionModelMixMin + Qwen2_5_VLPreTrainedModel, + ActionGenerationMixin, + ActionModelMixMin, + VLAInferenceMixin, ): - """ - Qwen2.5 Vision-Language Mixture of Experts model for action processing. - - This model extends the base Qwen2.5 VL model with action token processing capabilities - and optional LoRA fine-tuning support. - """ - - _tied_weights_keys = ["lm_head.weight"] + # Compatibility with newer transformers versions (5.x):_tied_weights_keys changed from list[str] to dict[str, str] (target -> source mapping). + # Older versions (4.x) iterate over dict keys, so the behavior is equivalent and compatible. + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} config_class = Qwen2_5_VLConfig _no_split_modules = ["Qwen2_5_VLDecoderLayer_with_MoE", "Qwen2_5_VLVisionBlock"] - @classmethod - def _set_customized_config(cls, config): - """ - Processing norm_stats.json and reconstruct the DoF mapping - """ - dataload_config = config["data"] - if not dataload_config.get("use_lerobot", False): - raise NotImplementedError( - "Not implemented for non-lerobot dataset currently" - ) - - enable_customized_robot_config = config.get( - "enable_customized_robot_config", False - ) - assert ( - enable_customized_robot_config - ), "enable_customized_robot_config must be true when use lerobot dataset" - - customized_dof_config = config["customized_robot_config"][ - "customized_dof_config" - ] - customized_agent_pos_config = config["customized_robot_config"][ - "customized_agent_pos_config" - ] - norm_stats_path = config["norm_stats_path"] - - # Use the compute_action_statistics function from utils - - name = config["customized_robot_config"]["name"] - - update_action_statistics( - action_statistic_dof=action_statistic_dof, # Assuming this is a global variable - norm_stats_path=norm_stats_path, - repo_id=config["data"]["lerobot_config"]["repo_id"], - robot_name=name, - customized_dof_config=customized_dof_config, - customized_agent_pos_config=customized_agent_pos_config, - ) - - print("Customized robot config added") - pprint(action_statistic_dof) - - @classmethod - def from_pretrained( - cls, - pretrained_model_path, - train_config=None, - config_path=None, - processor_path=None, - action_tokenizer_path=None, - is_train=False, - **kwargs, - ): - """ - Load model from pretrained model path. - - Args: - pretrained_model_path (str): Model directory path containing model.safetensors file - config_path (str, optional): Configuration file path, if None will look for qwen25_config.json in pretrained_model_path - processor_path (str, optional): Processor path, if None will load from default config - action_tokenizer_path (str, optional): Action tokenizer path, if None will load from default config - **kwargs: Additional arguments - - Returns: - Qwen2_5_VLMoEForAction: Loaded model instance - """ - # Load model components from pretrained path - - if train_config is None: - try: - with open(os.path.join(pretrained_model_path, "config.yml"), "r") as f: - train_config = yaml.load(f, Loader=yaml.FullLoader) - except Exception as e: - print(f"load train_config.yml fail: {e}") - train_config = None - - model_config_path = os.path.join(pretrained_model_path, "config.json") - model_config = cls.config_class.from_pretrained(model_config_path) - - if train_config is not None: - model_config = update_model_config(train_config, model_config) - processors_dict = load_wallx_processors(train_config) - processor = processors_dict["processor"] - else: - processor = AutoProcessor.from_pretrained( - pretrained_model_path, use_fast=True - ) - - if not is_train: - model_config._attn_implementation = "sdpa" - - if action_tokenizer_path is not None: - processor.action_processor = AutoProcessor.from_pretrained( - action_tokenizer_path, trust_remote_code=True - ) - - # Set the customized robot configuration to ensure consistency between cross-embodiment - # representations and the Wall-X action dimensionality. - # if not train_config: - # cls._set_customized_config(train_config) - # customized_dof_config = train_config["customized_robot_config"][ - # "customized_dof_config" - # ] - # customized_agent_pos_config = train_config["customized_robot_config"][ - # "customized_agent_pos_config" - # ] - # setattr(model_config, "customized_dof_config", customized_dof_config) - # setattr(model_config, "customized_agent_pos_config", customized_agent_pos_config) - - # Initialize model with configuration and processor - model = cls(model_config, processor=processor, **kwargs) - - # Resize token embeddings to match processor tokenizer vocabulary size - model.resize_token_embeddings(len(processor.tokenizer)) - - # Load model state dict from safetensors file - safetensor_files = glob.glob( - os.path.join(pretrained_model_path, "*.safetensors") - ) - state_dict = {} - embed_tokens_size = len(processor.tokenizer) - for file in safetensor_files: - sd = load_file(file, device="cpu") - # filter normalizer statistic params - del_keys = [] - for key in sd.keys(): - if "action_preprocessor.normalizer" in key: - print(f"filter load model weight {key}") - del_keys.append(key) - if "embed_tokens.weight" in key: - embed_tokens_size = sd[key].shape[0] - # if train_config is not None: - for key in del_keys: - del sd[key] - state_dict.update(sd) - if embed_tokens_size != len(processor.tokenizer): - model.resize_token_embeddings(embed_tokens_size) - model.load_state_dict(state_dict, strict=False) - - return model - def __init__( self, - config, - use_fast_tokenizer=False, + config: Qwen2_5_VLConfig, processor=None, - action_tokenizer=None, - action_mapper=None, - flow_loss_weight=1.0, + tokenizer_mixin=None, use_selective_recompute=False, ): - """ - Initialize the Qwen2.5 VLMoE model for action processing. - - Args: - config: Model configuration - use_fast_tokenizer (bool): Whether to use fast tokenizer - processor: Text and image processor - action_tokenizer: Action-specific tokenizer - action_mapper: Action mapping utility - flow_loss_weight (float): Weight for flow loss computation - """ super().__init__(config) - - # Initialize vision transformer and language model components - self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config( - config.vision_config - ) + self.visual = self._build_visual(config, use_selective_recompute) self.model = Qwen2_5_VLMoEModel( config, use_selective_recompute=use_selective_recompute ) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) - # Initialize loss function without reduction for channel-wise loss computation - self.loss_fct = CrossEntropyLoss(reduction="none") - self.flow_loss_weight = flow_loss_weight - self.use_fast_tokenizer = use_fast_tokenizer - self.processor = processor + self.loss_fct = CrossEntropyLoss( + reduction="none" + ) # do not do reduction to compute channel loss - # Define action token IDs + # Manage through tokenizer_mixin; in pure flow mode the mixin is None + self.processor = processor + self.tokenizer_mixin = tokenizer_mixin + if self.tokenizer_mixin is not None: + self.action_tokenizer_type = self.tokenizer_mixin.tokenizer_type + self.action_tokenizer = self.tokenizer_mixin.tokenizer + self.action_mapper = self.tokenizer_mixin.build_action_mapper(processor) + else: + self.action_tokenizer_type = None + self.action_tokenizer = None + self.action_mapper = None self.define_action_token_id() + self.times_cache = {} # cache times linspace for each num_inference_timesteps + self._infer_stable_cache: dict = {} # cache stable-state inference structures - # Cache for rope deltas - self.rope_deltas = None + self.rope_deltas = None # cache rope_deltas here - # Initialize action preprocessor - self.action_preprocessor = ActionProcessor(config) + from wall_x.model.core.action.processor import ActionProcessor - # Apply LoRA if specified in configuration + self.action_preprocessor = ActionProcessor(config) # action processing + + # Apply LoRA if the configuration contains LoRA settings if hasattr(config, "use_lora") and config.use_lora: self.add_lora( r=config.lora_r, @@ -982,85 +850,37 @@ def __init__( target_modules=config.lora_target_modules, lora_dropout=config.lora_dropout, ) - # Initialize weights and apply final processing self.post_init() + self._post_init_engine(config) - def define_action_token_id(self): - """ - Define action token IDs based on tokenizer configuration. - - Creates mappings for fast action tokens, proprioception tokens, and general action tokens. - """ - # Create list of fast action token IDs - fast_action_token_list = [] - if self.use_fast_tokenizer: - for i in range( - self.processor.tokenizer.init_kwargs["action_token_vocab_size"] - ): - action_token_id = self.processor.tokenizer.convert_tokens_to_ids( - f"<|action_token_{i}|>" - ) - fast_action_token_list.append(action_token_id) - - # Get special action token IDs - action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") - propri_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|propri|>") - - # Store action token ID mappings - self.action_token_id_set = { - "fast_action_token_list": fast_action_token_list, - "propri_token_id": propri_token_id, - "action_token_id": action_token_id, - } - - def add_lora( - self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1 - ): - """ - Add LoRA (Low-Rank Adaptation) adapters to the model. - - Args: - r (int): Rank of adaptation - lora_alpha (int): LoRA scaling parameter - target_modules (list): List of module names to apply LoRA to - lora_dropout (float): Dropout probability for LoRA layers - """ - config = LoraConfig( - r=r, - lora_alpha=lora_alpha, - target_modules=target_modules, - lora_dropout=lora_dropout, - bias="none", - task_type="CAUSAL_LM", + def _build_visual(self, config, use_selective_recompute): + """Factory method for visual encoder. Subclasses may override to swap implementation.""" + return Qwen2_5_VisionTransformerPretrainedModel( + config=config.vision_config, + use_selective_recompute=use_selective_recompute, ) - self.model = get_peft_model(self.model, config) - # Print information about trainable parameters - self.model.print_trainable_parameters() + def _post_init_engine(self, config): + """Hook called at the end of __init__. Subclasses may override to add engine-specific state.""" + pass def get_input_embeddings(self): - """Get input embeddings layer.""" return self.model.embed_tokens def set_input_embeddings(self, value): - """Set input embeddings layer.""" self.model.embed_tokens = value def get_output_embeddings(self): - """Get output embeddings layer.""" return self.lm_head def set_output_embeddings(self, new_embeddings): - """Set output embeddings layer.""" self.lm_head = new_embeddings def set_decoder(self, decoder): - """Set the decoder model.""" self.model = decoder def get_decoder(self): - """Get the decoder model.""" return self.model def get_rope_index( @@ -1284,46 +1104,44 @@ def get_rope_index( return position_ids, mrope_position_deltas - def train_step_forward( + def get_attention_maps( self, - input_ids: torch.LongTensor = None, + input_ids, + action_horizon, + action_dim, + num_inference_timesteps: int = 10, + prefix_length: Optional[int] = None, # prefix length attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + positional_masks: Optional[torch.LongTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, pixel_values: Optional[torch.Tensor] = None, pixel_values_videos: Optional[torch.FloatTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, + action_chunk: Optional[torch.FloatTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, + unnorm_proprioception: Optional[torch.FloatTensor] = None, rope_deltas: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, - # for vla - moe_token_types: Optional[torch.LongTensor] = None, - start_indices: Optional[torch.Tensor] = None, - end_indices: Optional[torch.Tensor] = None, - positional_masks: Optional[dict] = None, - action_chunk: Optional[torch.FloatTensor] = None, - proprioception: Optional[torch.FloatTensor] = None, dataset_names: Optional[str] = None, dof_mask: Optional[torch.FloatTensor] = None, agent_pos_mask: Optional[torch.FloatTensor] = None, - flow_loss_mask: Optional[torch.FloatTensor] = None, **kwargs, - ) -> Union[Tuple, Qwen2_5_VLACausalLMOutputWithPast]: + ): + if input_ids is not None: batch_size, seq_length = input_ids.shape - output_attentions = ( - output_attentions - if output_attentions is not None - else self.config.output_attentions - ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None @@ -1334,7 +1152,7 @@ def train_step_forward( ) if start_indices is None or end_indices is None: - # Calculate the start and end positions of each expert group's tokens after permutation + # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here group_size = torch.zeros( self.config.num_experts, dtype=torch.long, device="cpu" ) @@ -1355,12 +1173,17 @@ def train_step_forward( or self.rope_deltas is None or (past_key_values is None or past_key_values.get_seq_length() == 0) ): - position_ids, rope_deltas = self.get_rope_index( + position_ids, rope_deltas = get_rope_index( input_ids=input_ids, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, second_per_grid_ts=second_per_grid_ts, attention_mask=attention_mask, + spatial_merge_size=self.config.vision_config.spatial_merge_size, + image_token_id=self.config.image_token_id, + video_token_id=self.config.video_token_id, + vision_start_token_id=self.config.vision_start_token_id, + tokens_per_second=self.config.vision_config.tokens_per_second, ) self.rope_deltas = rope_deltas # then use the prev pre-calculated rope-deltas to get the correct position ids @@ -1383,8 +1206,8 @@ def train_step_forward( if pixel_values is not None: pixel_values = pixel_values.type(self.visual.dtype) image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) - mask = input_ids == self.config.image_token_id - mask_unsqueezed = mask.unsqueeze(-1) + img_mask = input_ids == self.config.image_token_id + mask_unsqueezed = img_mask.unsqueeze(-1) mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) image_mask = mask_expanded.to(inputs_embeds.device) @@ -1416,7 +1239,6 @@ def train_step_forward( inputs_embeds = self.scatter_proprioception_embeddings( input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask ) - inputs_embeds, flow, adarms_cond = self.scatter_flow_action_embeddings( input_ids, inputs_embeds, action_chunk, dataset_names, dof_mask ) @@ -1430,87 +1252,32 @@ def train_step_forward( attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, - moe_token_types=moe_token_types, + moe_token_types=moe_token_types, # pass token_types start_indices=start_indices, end_indices=end_indices, - positional_masks=positional_masks, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, + positional_masks=positional_masks, # pass masks for each category + use_cache=False, + output_attentions=True, + output_hidden_states=False, return_dict=return_dict, adarms_conds=[None, adarms_cond], # cache_position=cache_position, ) + attention_maps = torch.stack(outputs.attentions, dim=0)[ + None, :, -action_horizon: + ] # [flow steps, layer depths, all tokens, all tokens] + attention_maps = attention_maps.mean(2) + image_mask_indices = img_mask[0].nonzero(as_tuple=True)[0] + attention_maps = attention_maps[:, :, image_mask_indices] + return attention_maps - hidden_states = outputs[0] - logits = self.lm_head(hidden_states) - - ( - loss, - cross_entropy_loss, - flow_loss, - channel_loss_dict, - channel_loss_count_dict, - ) = self.compute_loss( - hidden_states=hidden_states, - logits=logits, - input_ids=input_ids, - dataset_names=dataset_names, - labels=labels, - action_chunk=action_chunk, - dof_mask=dof_mask, - flow=flow, - flow_loss_mask=flow_loss_mask, - ) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - - return Qwen2_5_VLACausalLMOutputWithPast( - loss=loss, - cross_entropy_loss=( - cross_entropy_loss.clone() if cross_entropy_loss is not None else None - ), - flow_loss=flow_loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - rope_deltas=self.rope_deltas, - channel_loss_dict=channel_loss_dict, - channel_loss_count_dict=channel_loss_count_dict, - ) - - def predict_action(self, predict_mode: str, **kwargs): - """ - Predict actions using specified prediction mode. - - Args: - predict_mode (str): Prediction mode, either "fast" or "diffusion" - **kwargs: Additional arguments passed to the predict method - - Returns: - tuple: (predicted_action, ground_truth_action) where ground_truth_action may be None - """ - assert predict_mode in ["fast", "diffusion"] - - output = self.predict(predict_mode=predict_mode, **kwargs) - - return output["predict_action"], output.get("gt_action", None) - - @torch.no_grad() - def predict( + def train_step_forward( self, - predict_mode: str, - pred_horizon: Optional[int] = None, - action_dim: Optional[int] = None, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, - moe_token_types: Optional[torch.LongTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, @@ -1520,76 +1287,26 @@ def predict( pixel_values_videos: Optional[torch.FloatTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, - action_chunk: Optional[torch.FloatTensor] = None, - proprioception: Optional[torch.FloatTensor] = None, rope_deltas: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, - num_inference_timesteps: Optional[int] = 10, + # for vla + sample_time: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + positional_masks: Optional[dict] = None, + action_chunk: Optional[torch.FloatTensor] = None, + proprioception: Optional[torch.FloatTensor] = None, dataset_names: Optional[str] = None, dof_mask: Optional[torch.FloatTensor] = None, agent_pos_mask: Optional[torch.FloatTensor] = None, - re_generate: bool = False, + flow_loss_mask: Optional[torch.FloatTensor] = None, **kwargs, - ): - """ - Multi-modal prediction method supporting text generation, fast action prediction, and diffusion-based action prediction. - - This method handles three prediction modes: - 1. "text": Pure text generation using autoregressive decoding - 2. "fast": Fast action prediction using discrete action tokens - 3. "diffusion": Continuous action prediction using diffusion/flow matching - - Args: - predict_mode (str): Prediction mode ("text", "fast", or "diffusion") - pred_horizon (int, optional): Prediction horizon for action sequences - action_dim (int, optional): Dimensionality of action space - input_ids (torch.LongTensor, optional): Input token IDs - attention_mask (torch.Tensor, optional): Attention mask for input tokens - position_ids (torch.LongTensor, optional): Position IDs for tokens - past_key_values (List[torch.FloatTensor], optional): Cached key-value pairs - inputs_embeds (torch.FloatTensor, optional): Pre-computed input embeddings - moe_token_types (torch.LongTensor, optional): Token type assignments for MoE routing - labels (torch.LongTensor, optional): Target labels for evaluation - use_cache (bool, optional): Whether to use key-value caching - output_attentions (bool, optional): Whether to return attention weights - output_hidden_states (bool, optional): Whether to return hidden states - return_dict (bool, optional): Whether to return structured output - pixel_values (torch.Tensor, optional): Image pixel values - pixel_values_videos (torch.FloatTensor, optional): Video pixel values - image_grid_thw (torch.LongTensor, optional): Image grid dimensions - video_grid_thw (torch.LongTensor, optional): Video grid dimensions - action_chunk (torch.FloatTensor, optional): Ground truth action sequences - proprioception (torch.FloatTensor, optional): Proprioceptive sensor data - rope_deltas (torch.LongTensor, optional): RoPE position deltas - cache_position (torch.LongTensor, optional): Cache position indices - second_per_grid_ts (torch.Tensor, optional): Time interval per temporal grid - num_inference_timesteps (int, optional): Number of diffusion inference steps - dataset_names (str, optional): Dataset names for normalization - dof_mask (torch.FloatTensor, optional): Degrees of freedom mask - agent_pos_mask (torch.FloatTensor, optional): Agent position mask - re_generate (bool, optional): Whether to use sampling for regeneration - **kwargs: Additional keyword arguments - - Returns: - dict: Dictionary containing prediction results with keys like: - - 'predict_action': Predicted action sequences - - 'gt_action': Ground truth actions (if available) - - 'input_text': Input text (for text/fast modes) - - 'predict_output_text': Generated text (for text/fast modes) - - 'gt_output_text': Ground truth text (for text/fast modes) - """ - batch_size = ( - input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] - ) - - # Text and fast modes require batch size 1 for autoregressive generation - if predict_mode in ["text", "fast"]: - assert ( - batch_size == 1 - ), "predict only support batch size 1 for ar generation" + ) -> Union[Tuple, Qwen2_5_VLACausalLMOutputWithPast]: + if input_ids is not None: + batch_size, seq_length = input_ids.shape - # Set output configuration from model config if not specified output_attentions = ( output_attentions if output_attentions is not None @@ -1604,92 +1321,29 @@ def predict( return_dict if return_dict is not None else self.config.use_return_dict ) - # Process input embeddings with multi-modal data - if inputs_embeds is None: - inputs_embeds = self.model.embed_tokens(input_ids) - - # Process image embeddings - if pixel_values is not None: - pixel_values = pixel_values.type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) - n_image_tokens = (input_ids == self.config.image_token_id).sum().item() - n_image_features = image_embeds.shape[0] + if start_indices is None or end_indices is None: + # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() - # Validate image token and feature count match - if n_image_tokens != n_image_features: - raise ValueError( - f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" - ) + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) - mask = input_ids == self.config.image_token_id - mask_unsqueezed = mask.unsqueeze(-1) - mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) - image_mask = mask_expanded.to(inputs_embeds.device) - - image_embeds = image_embeds.to( - inputs_embeds.device, inputs_embeds.dtype - ) - inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) - - # Process video embeddings - if pixel_values_videos is not None: - pixel_values_videos = pixel_values_videos.type(self.visual.dtype) - video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) - n_video_tokens = (input_ids == self.config.video_token_id).sum().item() - n_video_features = video_embeds.shape[0] - - # Validate video token and feature count match - if n_video_tokens != n_video_features: - raise ValueError( - f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" - ) - - mask = input_ids == self.config.video_token_id - mask_unsqueezed = mask.unsqueeze(-1) - mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) - video_mask = mask_expanded.to(inputs_embeds.device) - - video_embeds = video_embeds.to( - inputs_embeds.device, inputs_embeds.dtype - ) - inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) - - # Process proprioceptive data - if proprioception is not None: - proprioception = proprioception.to(inputs_embeds.device).to( - inputs_embeds.dtype - ) - agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to( - inputs_embeds.dtype - ) - proprio_embed = self.action_preprocessor.proprioception_proj( - proprioception, - dataset_names, - agent_pos_mask, - use_history=proprioception.shape[1] > 1, - ) - proprioception_mask = ( - input_ids == self.action_token_id_set["propri_token_id"] - ) - inputs_embeds[proprioception_mask] = proprio_embed.reshape( - -1, inputs_embeds.shape[-1] - ).to(inputs_embeds.dtype) - - if attention_mask is not None: - attention_mask = attention_mask.to(inputs_embeds.device) - - # Calculate RoPE position IDs if not provided - # Note: Cannot calculate rope deltas with 4D attention mask. TODO: Fix this limitation + # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme if position_ids is None and ( attention_mask is None or attention_mask.ndim == 2 ): - # Calculate RoPE index once per generation in the pre-fill stage only + # calculate RoPE index once per generation in the pre-fill stage only if ( (cache_position is not None and cache_position[0] == 0) or self.rope_deltas is None or (past_key_values is None or past_key_values.get_seq_length() == 0) ): - position_ids, rope_deltas = ops.get_rope_index( + position_ids, rope_deltas = get_rope_index( input_ids=input_ids, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, @@ -1702,161 +1356,220 @@ def predict( tokens_per_second=self.config.vision_config.tokens_per_second, ) self.rope_deltas = rope_deltas - # Use previously calculated rope deltas to get correct position IDs + # then use the prev pre-calculated rope-deltas to get the correct position ids else: - batch_size, seq_length, _ = inputs_embeds.shape + # batch_size, seq_length, _ = inputs_embeds.shape delta = ( - (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) + (cache_position[0] + self.rope_deltas).to(cache_position.device) if cache_position is not None else 0 ) - position_ids = torch.arange(seq_length, device=inputs_embeds.device) + position_ids = torch.arange(seq_length, device=cache_position.device) position_ids = position_ids.view(1, -1).expand(batch_size, -1) if cache_position is not None: # otherwise `deltas` is an int `0` delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) position_ids = position_ids.add(delta) position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) - # Prepare action chunk data if provided - if action_chunk is not None: - action_chunk = action_chunk.to(inputs_embeds.device).to(inputs_embeds.dtype) + if inputs_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + if pixel_values is not None: + pixel_values = pixel_values.type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + mask = input_ids == self.config.image_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + image_mask = mask_expanded.to(inputs_embeds.device) - output = {} + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) - # Split input sequence for text and fast modes (not needed for diffusion) - if predict_mode == "text" or predict_mode == "fast": - # Look for generation prompt tokens: <|im_start|>assistant - generation_prompt_ids = torch.tensor( - [151644, 77091], device=input_ids.device, dtype=input_ids.dtype + if pixel_values_videos is not None: + pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + n_video_features = video_embeds.shape[0] + if n_video_tokens != n_video_features: + raise ValueError( + f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}" + ) + + mask = input_ids == self.config.video_token_id + mask_unsqueezed = mask.unsqueeze(-1) + mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) + video_mask = mask_expanded.to(inputs_embeds.device) + + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + inputs_embeds = self.scatter_proprioception_embeddings( + input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask ) - matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & ( - input_ids[0, 1:] == generation_prompt_ids[1] + inputs_embeds, flow, adarms_cond = self.scatter_flow_action_embeddings( + input_ids, + inputs_embeds, + action_chunk, + dataset_names, + sample_time, + dof_mask, ) - if matches.any(): - split_pos = torch.nonzero(matches, as_tuple=True)[0][0].item() - # Extract ground truth output tokens (including newline) - gt_output_ids = input_ids[:, split_pos + 3 :] - # Remove output part from input, keeping prompt - input_ids = input_ids[:, : split_pos + 3] - inputs_embeds = inputs_embeds[:, : split_pos + 3, :] - if attention_mask is not None: - attention_mask = attention_mask[:, : split_pos + 3] - if labels is not None: - labels = labels[:, split_pos + 3 :] + if attention_mask is not None: + attention_mask = attention_mask.to(inputs_embeds.device) + + outputs = self.model( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + moe_token_types=moe_token_types, # pass token_types + start_indices=start_indices, + end_indices=end_indices, + positional_masks=positional_masks, # pass masks for each category + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + adarms_conds=[None, adarms_cond], + # cache_position=cache_position, + ) + + hidden_states = outputs[0] + + # --- optimization: loss-token-only LM Head --- + # During training, only compute lm_head for tokens that contribute to CE loss, + # reducing GEMM size from [B*S, V] to [N_loss, V]. + # FSDP fp32 layernorms emit fp32 hidden_states; lm_head runs in bf16. + _lm_dtype = self.lm_head.weight.dtype + _lm_loss_mask = None + if labels is not None and self.training: + shift_labels = labels[..., 1:].contiguous() + _lm_loss_mask = shift_labels != -100 # [B, S-1] + if _lm_loss_mask.any(): + lm_hidden = hidden_states[:, :-1, :][_lm_loss_mask] # [N_loss, H] + logits = self.lm_head(lm_hidden.to(_lm_dtype)) # [N_loss, V] else: - raise Warning( - "input_ids does not contain the generation prompt tokens <|im_start|>assistant" - ) + logits = None + else: + logits = self.lm_head(hidden_states.to(_lm_dtype)) - # Decode input text for output - input_text = self.processor.batch_decode( - input_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True - ) - output["input_text"] = input_text - - # Handle text and fast prediction modes using autoregressive generation - if predict_mode == "text" or predict_mode == "fast": - # Initialize MoE token types for generation - moe_token_types = torch.zeros_like(input_ids) - batch = { - "input_ids": input_ids, - "attention_mask": attention_mask, - "pixel_values": pixel_values, - "moe_token_types": moe_token_types, - "image_grid_thw": image_grid_thw, - "dof_mask": dof_mask, - "agent_pos_mask": agent_pos_mask, - "proprioception": proprioception, - "dataset_names": dataset_names, - } + ( + loss, + cross_entropy_loss, + flow_loss, + channel_loss_dict, + channel_loss_count_dict, + ) = self.compute_loss( + hidden_states=hidden_states, + logits=logits, + input_ids=input_ids, + dataset_names=dataset_names, + labels=labels, + action_chunk=action_chunk, + dof_mask=dof_mask, + flow=flow, + flow_loss_mask=flow_loss_mask, + _lm_loss_mask=_lm_loss_mask, + ) - # Generate output tokens - predict_output_ids = self.generate( - **batch, - max_new_tokens=100, - eos_token_id=[self.processor.tokenizer.eos_token_id], - use_cache=True, - pad_token_id=self.processor.tokenizer.pad_token_id, - temperature=( - 1.0 if not re_generate else 0.7 - ), # Higher temperature for regeneration - do_sample=( - False if not re_generate else True - ), # Enable sampling for regeneration - ) + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output - # Decode generated and ground truth text - gt_output_text = self.processor.batch_decode( - gt_output_ids, - skip_special_tokens=False, - clean_up_tokenization_spaces=True, + return Qwen2_5_VLACausalLMOutputWithPast( + loss=loss, + cross_entropy_loss=( + cross_entropy_loss.clone() if cross_entropy_loss is not None else None + ), + flow_loss=flow_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + rope_deltas=self.rope_deltas, + channel_loss_dict=channel_loss_dict, + channel_loss_count_dict=channel_loss_count_dict, + ) + + def predict_action(self, predict_mode: str, **kwargs): + assert predict_mode in ["fast", "spatial_token", "flow", "diffusion", "dllm"] + + if predict_mode in ["fast", "spatial_token"]: + output = self.generate_ar_action(**kwargs) + elif predict_mode in ["diffusion", "flow"]: + output = self.generate_flow_action(**kwargs) + elif predict_mode == "dllm": + output = self.generate_dllm_action(**kwargs) + + return output["predict_action"], output.get("gt_action", None) + + @torch.no_grad() + def generate_text( + self, + input_ids: torch.LongTensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + moe_token_types: Optional[torch.LongTensor] = None, # new parameter + start_indices: Optional[torch.Tensor] = None, + end_indices: Optional[torch.Tensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_side_image: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + side_image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + action_chunk: Optional[torch.FloatTensor] = None, # action chunk + proprioception: Optional[torch.FloatTensor] = None, # joint positions + rope_deltas: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + second_per_grid_ts: Optional[torch.Tensor] = None, + dataset_names: Optional[str] = None, + dof_mask: Optional[torch.FloatTensor] = None, + agent_pos_mask: Optional[torch.FloatTensor] = None, + gt_output_ids: Optional[torch.LongTensor] = None, + prefix_length: Optional[int] = None, # prefix length + positional_masks: Optional[dict] = None, # position masks for each modality + re_generate: bool = False, + **kwargs, + ): + # assert self.config._attn_implementation == "flash_attention_2", "generate_text only support flash_attention_2 attn implementation" + + # generate text only need prefix part + if prefix_length is not None: + input_ids, _gt_output_ids = ( + input_ids[:, :prefix_length], + input_ids[:, prefix_length:], ) - predict_output_text = self.processor.batch_decode( - predict_output_ids, - skip_special_tokens=False, - clean_up_tokenization_spaces=True, + gt_output_ids = ( + gt_output_ids if gt_output_ids is not None else _gt_output_ids ) - output["gt_output_text"] = gt_output_text - output["predict_output_text"] = predict_output_text - - # Convert tokens to actions for fast prediction mode - if predict_mode == "fast": - action_id = [] - # Extract action tokens from generated sequence - for token_id_i in predict_output_ids[0]: - if ( - token_id_i.item() - >= self.processor.tokenizer.init_kwargs["action_token_start_index"] - ): - action_id.append( - token_id_i.item() - - self.processor.tokenizer.init_kwargs[ - "action_token_start_index" - ] - ) - - predict_action = self.processor.action_processor.decode( - [action_id], time_horizon=pred_horizon, action_dim=action_dim + attention_mask = ( + attention_mask[:, :prefix_length] + if attention_mask is not None + else None ) - # Handle action decoding errors - if np.sum(predict_action) == 0: - print("Error in decoding action, predict_action is None") - output["predict_action"] = None - else: - # Convert discrete tokens to continuous actions - predict_action = torch.tensor(predict_action, device=self.device) - dof_mask = dof_mask.to(self.device).to(pixel_values.dtype) - predict_action = ( - self.action_preprocessor.normalizer_action.unnormalize_data( - predict_action, dataset_names, dof_mask - ) - ) - output["predict_action"] = predict_action - - # Process ground truth actions if available - if action_chunk is not None: - # Apply DOF mask and unnormalize action chunk to get ground truth actions - action_chunk = action_chunk[:, :, dof_mask[0, 0, :].bool()] - output["gt_action"] = ( - self.action_preprocessor.normalizer_action.unnormalize_data( - action_chunk, dataset_names, dof_mask - ) - ) - else: - output["gt_action"] = None - - # Handle diffusion-based action prediction - if predict_mode == "diffusion": - # Initialize with random noise - noisy_action = torch.randn( - size=(batch_size, pred_horizon, action_dim), - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, + moe_token_types = ( + moe_token_types[:, :prefix_length] + if moe_token_types is not None + else None ) - dof_mask = dof_mask.to(inputs_embeds.device).to(inputs_embeds.dtype) + else: + prefix_length = input_ids.shape[1] - # Calculate token distribution across MoE expert groups + if start_indices is None or end_indices is None: + # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here group_size = torch.zeros( self.config.num_experts, dtype=torch.long, device="cpu" ) @@ -1867,92 +1580,189 @@ def predict( start_indices = torch.cumsum(group_size, dim=0) - group_size end_indices = torch.cumsum(group_size, dim=0) - def step(timestep, noisy_action): - """ - Single denoising step for diffusion process. - - Args: - timestep: Current diffusion timestep - noisy_action: Current noisy action estimate - - Returns: - torch.Tensor: Predicted clean action - """ - action_mask = input_ids == self.action_token_id_set["action_token_id"] - assert action_mask.any(), "No action token found in input_ids" - - # Prepare timestep for batch processing - timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0]) - action_embed, _ = self.action_preprocessor.step( - timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask - ) - action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]) + batch = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "moe_token_types": moe_token_types, + "start_indices": start_indices, + "end_indices": end_indices, + "image_grid_thw": image_grid_thw, + "dof_mask": dof_mask, + "agent_pos_mask": agent_pos_mask, + "proprioception": proprioception, + "dataset_names": dataset_names, + # "prefix_length": prefix_length, + } + predict_output_ids = self.generate( + **batch, + max_new_tokens=100, + eos_token_id=[ + self.processor.tokenizer.eos_token_id + ], # set multiple end markers + use_cache=True, + pad_token_id=self.processor.tokenizer.pad_token_id, # explicitly set pad_token_id + temperature=( + 1.0 if not re_generate else 0.7 + ), # use a higher temperature when regenerating + do_sample=( + False if not re_generate else True + ), # use sampling when regenerating + ) + input_text = self.processor.batch_decode( + input_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True + ) + if gt_output_ids is not None and gt_output_ids.shape[-1] != 0: + gt_output_text = self.processor.batch_decode( + gt_output_ids, + skip_special_tokens=False, + clean_up_tokenization_spaces=True, + ) + else: + gt_output_text = None + predict_output_text = self.processor.batch_decode( + predict_output_ids[:, prefix_length:], + skip_special_tokens=False, + clean_up_tokenization_spaces=True, + ) + output = { + "input_text": input_text, + "gt_output_text": gt_output_text, + "predict_output_text": predict_output_text, + "predict_output_ids": predict_output_ids, + } + return output - # Create temporary copy of embeddings for thread safety - temp_inputs_embeds = inputs_embeds.clone() - temp_inputs_embeds[action_mask] = action_embed.to( - temp_inputs_embeds.dtype - ) + @torch.no_grad() + def generate_ar_action( + self, + input_ids, + action_horizon, + action_dim, + state=None, + dataset_names=None, + action_chunk=None, + dof_mask=None, + unnorm=True, + max_retries=3, # maximum retry count + proprioception=None, # proprioception for v3.1_delta decoding + robot_type_id=None, # robot type ID for v3.1_delta decoding + **kwargs, + ): + # Regeneration mechanism + predict_action = None + output = None + decode_success = False + if dataset_names is None and (unnorm or action_chunk is not None): + raise KeyError( + "generate_ar_action requires dataset_names when normalizing or " + "unnormalizing actions." + ) - # Forward pass through transformer - transformer_outputs = self.model( - input_ids=None, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=temp_inputs_embeds, - moe_token_types=moe_token_types, - start_indices=start_indices, - end_indices=end_indices, - use_cache=True, - output_attentions=False, - output_hidden_states=False, - return_dict=True, - ) + for retry_attempt in range(max_retries): + use_re_generate = retry_attempt > 0 - # Extract action predictions from hidden states - hidden_states = transformer_outputs.last_hidden_state - action_mask = input_ids == self.action_token_id_set["action_token_id"] - action_hidden_states = hidden_states[action_mask] - pred = self.action_preprocessor.action_proj_back( - action_hidden_states[ - :, : self.action_preprocessor.action_hidden_size - ] + if retry_attempt > 0: + logger.info( + "Retrying action generation (attempt %s/%s)...", + retry_attempt + 1, + max_retries, ) - return pred.reshape(batch_size, pred_horizon, action_dim) - # Perform ODE integration for diffusion sampling - times = torch.linspace( - 0, - 1, - num_inference_timesteps + 1, - device=inputs_embeds.device, - dtype=inputs_embeds.dtype, + output = self.generate_text( + input_ids=input_ids, re_generate=use_re_generate, **kwargs ) - action_trajectory = odeint( - step, - noisy_action.to(torch.float32), - times.to(torch.float32), - method="euler", + predict_output_ids = output["predict_output_ids"] + + # Decode uniformly through tokenizer_mixin + predict_action, decode_success = self.tokenizer_mixin.decode_action( + output_ids=predict_output_ids, + action_mapper=self.action_mapper, + action_horizon=action_horizon, + action_dim=action_dim, + device=input_ids.device, + proprioception=proprioception, + dof_mask=dof_mask, + robot_type_id=robot_type_id, + state=state, ) - # Extract final predicted action and unnormalize - predict_action = action_trajectory[-1] - predict_action = ( - self.action_preprocessor.normalizer_action.unnormalize_data( - predict_action, dataset_names - ) - ) - output["predict_action"] = predict_action + if decode_success: + if retry_attempt > 0: + logger.info( + "Action generation succeeded on attempt %s", + retry_attempt + 1, + ) + break + else: + if retry_attempt < max_retries - 1: + logger.warning( + "Error in decoding action (attempt %s/%s), retrying with re_generate mode...", + retry_attempt + 1, + max_retries, + ) + else: + logger.warning( + "Error in decoding action after %s attempts, returning None", + max_retries, + ) + + # Check whether dof_mask is required through mixin attributes + uses_dof_mask = ( + self.tokenizer_mixin.uses_dof_mask_for_unnorm + if self.tokenizer_mixin is not None + else True + ) - # Process ground truth actions if available - if action_chunk is not None: + if action_chunk is not None: + action_chunk = action_chunk.to(input_ids.device).to(torch.bfloat16) + if uses_dof_mask: + action_chunk = action_chunk[:, :, dof_mask[0, 0, :].bool()] + output["gt_action"] = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + action_chunk, dataset_names, dof_mask + ) + ) + else: + # Output full dimensions without mask filtering output["gt_action"] = ( self.action_preprocessor.normalizer_action.unnormalize_data( - action_chunk, dataset_names + action_chunk, dataset_names, None ) ) + if not decode_success or predict_action is None: + logger.warning("Error in decoding action, predict_action is None") + output["predict_action"] = None + else: + if unnorm: + if isinstance(predict_action, np.ndarray): + predict_action = torch.tensor( + predict_action, device=input_ids.device + ) + elif predict_action.device != input_ids.device: + predict_action = predict_action.to(input_ids.device) + + # Add the batch dimension + if predict_action.dim() == 2: + predict_action = predict_action.unsqueeze(0) + + if uses_dof_mask: + predict_action = unnormalize_data_with_virtual_tail( + self.action_preprocessor.normalizer_action, + predict_action, + dataset_names, + dof_mask, + ) + else: + # Output full dimensions without mask filtering + predict_action = ( + self.action_preprocessor.normalizer_action.unnormalize_data( + predict_action, dataset_names, None + ) + ) + output["predict_action"] = predict_action + return output @torch.no_grad() @@ -1962,8 +1772,7 @@ def generate_flow_action( action_horizon, action_dim, num_inference_timesteps: int = 10, - padding_action: Optional[torch.Tensor] = None, - prefix_length: Optional[int] = None, + prefix_length: Optional[int] = None, # prefix length attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, @@ -1990,13 +1799,94 @@ def generate_flow_action( dataset_names: Optional[str] = None, dof_mask: Optional[torch.FloatTensor] = None, agent_pos_mask: Optional[torch.FloatTensor] = None, - unnorm: Optional[bool] = True, **kwargs, ): - total_start_time = time.time() - timing_results = {} + # assert self.config._attn_implementation == "sdpa", "generate_flow_action only support sdpa attn implementation" + ctx = self._prepare_flow_action_inputs( + input_ids=input_ids, + action_horizon=action_horizon, + action_dim=action_dim, + num_inference_timesteps=num_inference_timesteps, + prefix_length=prefix_length, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + moe_token_types=moe_token_types, + start_indices=start_indices, + end_indices=end_indices, + positional_masks=positional_masks, + labels=labels, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + action_chunk=action_chunk, + proprioception=proprioception, + unnorm_proprioception=unnorm_proprioception, + rope_deltas=rope_deltas, + cache_position=cache_position, + second_per_grid_ts=second_per_grid_ts, + dataset_names=dataset_names, + dof_mask=dof_mask, + agent_pos_mask=agent_pos_mask, + ) + + action_trajectory, attention_maps = self._execute_flow_action(ctx) + + output = self._finalize_flow_action_output(action_trajectory, ctx) + if output_attentions: + output["attention_maps"] = attention_maps + return output + + def _execute_flow_action(self, ctx): + """Strategy hook: selects which flow-action backend to use. + + Returns: + action_trajectory: list of action tensors + attention_maps: attention maps or None + """ + return self._generate_flow_action_vanilla(ctx) + def _prepare_flow_action_inputs( + self, + input_ids, + action_horizon, + action_dim, + num_inference_timesteps, + prefix_length, + attention_mask, + position_ids, + past_key_values, + inputs_embeds, + moe_token_types, + start_indices, + end_indices, + positional_masks, + labels, + use_cache, + output_attentions, + output_hidden_states, + return_dict, + pixel_values, + pixel_values_videos, + image_grid_thw, + video_grid_thw, + action_chunk, + proprioception, + unnorm_proprioception, + rope_deltas, + cache_position, + second_per_grid_ts, + dataset_names, + dof_mask, + agent_pos_mask, + ): batch_size = ( input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] ) @@ -2014,11 +1904,14 @@ def generate_flow_action( return_dict if return_dict is not None else self.config.use_return_dict ) - embed_start_time = time.time() + # Timing: input embedding processing + img_mask = None if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) if pixel_values is not None: pixel_values = pixel_values.type(self.visual.dtype) + # from wall_x.utils.timers import ScopeTimer + # with ScopeTimer("pixel_values.visual"): image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) n_image_tokens = (input_ids == self.config.image_token_id).sum().item() n_image_features = image_embeds.shape[0] @@ -2027,8 +1920,8 @@ def generate_flow_action( f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}" ) - mask = input_ids == self.config.image_token_id - mask_unsqueezed = mask.unsqueeze(-1) + img_mask = input_ids == self.config.image_token_id + mask_unsqueezed = img_mask.unsqueeze(-1) mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) image_mask = mask_expanded.to(inputs_embeds.device) @@ -2079,9 +1972,6 @@ def generate_flow_action( if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) - timing_results["embed_processing"] = time.time() - embed_start_time - - position_start_time = time.time() # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme if position_ids is None and ( attention_mask is None or attention_mask.ndim == 2 @@ -2092,12 +1982,17 @@ def generate_flow_action( or self.rope_deltas is None or (past_key_values is None or past_key_values.get_seq_length() == 0) ): - position_ids, rope_deltas = self.get_rope_index( + position_ids, rope_deltas = get_rope_index( input_ids, image_grid_thw, video_grid_thw, second_per_grid_ts, attention_mask, + self.config.vision_config.spatial_merge_size, + self.config.image_token_id, + self.config.video_token_id, + self.config.vision_start_token_id, + self.config.vision_config.tokens_per_second, ) self.rope_deltas = rope_deltas # then use the prev pre-calculated rope-deltas to get the correct position ids @@ -2115,25 +2010,46 @@ def generate_flow_action( position_ids = position_ids.add(delta) position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) - if start_indices is None or end_indices is None: - # Calculate the start and end positions of each expert group's tokens after permutation (the dataset does not contain `num_expert` information, so this calculation must be done here). - group_size = torch.zeros( - self.config.num_experts, dtype=torch.long, device="cpu" - ) - for i in range(self.config.num_experts): - group_size[i] = (moe_token_types == i).sum() + flow_action_mask = input_ids == self.action_token_id_set["action_token_id"] - # Calculate start and end indices for each expert group - start_indices = torch.cumsum(group_size, dim=0) - group_size - end_indices = torch.cumsum(group_size, dim=0) + if prefix_length is None: + has_true = flow_action_mask.any(dim=1) + prefix_length = torch.argmax(flow_action_mask.float(), dim=1, keepdim=True) + prefix_length[~has_true] = flow_action_mask.shape[1] + # check if prefix_length is the same for all batch + if not torch.all(prefix_length == prefix_length[0]): + raise ValueError( + "prefix_length differs across batch; batch prompts must align" + ) + prefix_length = int(prefix_length[0].item()) - timing_results["position_encoding"] = time.time() - position_start_time + if start_indices is None or end_indices is None: + # Compute start and end positions for each expert token group after permutation; the dataset has no num_expert metadata, so this is computed here + # Cache key: shape + sum is sufficient for stable-state (moe_token_types structure doesn't change) + _moe_key = (moe_token_types.shape, int(moe_token_types.sum().item())) + _moe_cache = self._infer_stable_cache.get("moe_indices") + if _moe_cache is None or _moe_cache[0] != _moe_key: + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" + ) + for i in range(self.config.num_experts): + group_size[i] = (moe_token_types == i).sum() + + # Calculate start and end indices for each expert group + start_indices = torch.cumsum(group_size, dim=0) - group_size + end_indices = torch.cumsum(group_size, dim=0) + self._infer_stable_cache["moe_indices"] = ( + _moe_key, + start_indices, + end_indices, + ) + else: + _, start_indices, end_indices = _moe_cache - action_init_start_time = time.time() + # Timing: action initialization if action_chunk is not None: action_chunk = action_chunk.to(inputs_embeds.device).to(torch.float32) - output = {} # Reproduce # torch.manual_seed(0) noise = torch.randn( @@ -2142,18 +2058,16 @@ def generate_flow_action( device=inputs_embeds.device, ) noisy_action = noise.clone() + dof_mask = dof_mask.to(inputs_embeds.device).to(torch.float32) if num_inference_timesteps not in self.times_cache: - self.times_cache[num_inference_timesteps] = torch.linspace( - 0.0, - 1.0, - num_inference_timesteps + 1, - device=inputs_embeds.device, - dtype=torch.float32, + self.times_cache[num_inference_timesteps] = ( + self.action_preprocessor.get_inference_times( + num_inference_timesteps, inputs_embeds.device, torch.float32 + ) ) times = self.times_cache[num_inference_timesteps] - dt = times[1] - times[0] time_0 = times[0].unsqueeze(0).repeat(noisy_action.shape[0]) action_embed, adarms_cond = self.action_preprocessor.step( timestep=time_0, noisy_action=noisy_action, dof_mask=dof_mask @@ -2161,134 +2075,101 @@ def generate_flow_action( action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]).to( inputs_embeds.dtype ) - flow_action_mask = input_ids == self.action_token_id_set["action_token_id"] - inputs_embeds[flow_action_mask] = action_embed - timing_results["action_initialization"] = time.time() - action_init_start_time + # Automatically generate padding actions from dof_mask + if not dof_mask.all(): + padding_action = ( + torch.zeros((1, dof_mask.shape[-1])) + .to(dof_mask.device) + .to(torch.float32) + ) + padding_action = normalize_data_with_virtual_tail( + self.action_preprocessor.normalizer_action, + padding_action, + dataset_names, + dof_mask, + ) + else: + padding_action = None + + return { + "batch_size": batch_size, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "position_ids": position_ids, + "prefix_length": prefix_length, + "flow_action_mask": flow_action_mask, + "start_indices": start_indices, + "end_indices": end_indices, + "dof_mask": dof_mask, + "noise": noise, + "noisy_action": noisy_action, + "times": times, + "adarms_cond": adarms_cond, + "action_chunk": action_chunk, + "padding_action": padding_action, + "action_horizon": action_horizon, + "action_dim": action_dim, + "input_ids": input_ids, + "moe_token_types": moe_token_types, + "positional_masks": positional_masks, + "dataset_names": dataset_names, + "output_attentions": output_attentions, + "img_mask": img_mask, + } - prefetch_start_time = time.time() + def _generate_flow_action_vanilla(self, ctx): + # Timing: prefill forward pass + all_attention_maps = [] + output_attentions = True if ctx["output_attentions"] else False prefetch_output = self.model( input_ids=None, - attention_mask=attention_mask, - position_ids=position_ids, + attention_mask=ctx["attention_mask"], + position_ids=ctx["position_ids"], past_key_values=None, - inputs_embeds=inputs_embeds, - moe_token_types=moe_token_types, - start_indices=start_indices, - end_indices=end_indices, - positional_masks=positional_masks, + inputs_embeds=ctx["inputs_embeds"], + moe_token_types=ctx["moe_token_types"], + start_indices=ctx["start_indices"], + end_indices=ctx["end_indices"], + positional_masks=ctx["positional_masks"], use_cache=True, - output_attentions=False, + output_attentions=output_attentions, output_hidden_states=False, return_dict=True, - adarms_conds=[None, adarms_cond], + adarms_conds=[None, ctx["adarms_cond"]], ) hidden_states = prefetch_output.last_hidden_state prefix_kv_cache = prefetch_output.past_key_values - - action_hidden_states = hidden_states[flow_action_mask].to(torch.float32) + if output_attentions: + all_attention_maps.append(prefetch_output.attentions) + action_hidden_states = hidden_states[ctx["flow_action_mask"]].to(torch.float32) action_pred = self.action_preprocessor.action_proj_back( action_hidden_states[:, : self.action_preprocessor.action_hidden_size] ) if getattr(self.config, "use_x_pred", False): - v_0 = action_pred - noise.reshape(-1, noise.shape[-1]) + v_0 = action_pred - ctx["noise"].reshape(-1, ctx["noise"].shape[-1]) else: v_0 = action_pred + v_0 = v_0.reshape(ctx["batch_size"], ctx["action_horizon"], ctx["action_dim"]) - if (not dof_mask.all()) and (padding_action is not None): - print("use padding action", flush=True) - v_padding = padding_action - noisy_action - v_0 = (v_padding) * (1 - dof_mask) + v_0 * dof_mask - - noisy_action = noisy_action + dt * v_0.reshape( - batch_size, action_horizon, action_dim - ) - - timing_results["prefetch_forward"] = time.time() - prefetch_start_time - - cache_prep_start_time = time.time() - - if prefix_length is None: - has_true = flow_action_mask.any(dim=1) - prefix_length = torch.argmax(flow_action_mask.float(), dim=1, keepdim=True) - prefix_length[~has_true] = flow_action_mask.shape[1] - prefix_length = prefix_length[0] - - # support different transformers version - if hasattr(prefix_kv_cache, "key_cache"): - for layer_i in range(len(prefix_kv_cache.key_cache)): - prefix_kv_cache.key_cache[layer_i] = prefix_kv_cache.key_cache[layer_i][ - :, :, :prefix_length, : - ] - prefix_kv_cache.value_cache[layer_i] = prefix_kv_cache.value_cache[ - layer_i - ][:, :, :prefix_length, :] - else: - for layer_i in range(len(prefix_kv_cache.layers)): - prefix_kv_cache.layers[layer_i].keys = prefix_kv_cache.layers[ - layer_i - ].keys[:, :, :prefix_length, :] - prefix_kv_cache.layers[layer_i].values = prefix_kv_cache.layers[ - layer_i - ].values[:, :, :prefix_length, :] - - postfix_position_ids = position_ids[:, :, prefix_length:] - postfix_inputs_embeds = inputs_embeds[:, prefix_length:, :] - postfix_attention_mask = attention_mask[:, prefix_length:] - postfix_moe_token_types = moe_token_types[:, prefix_length:] - postfix_input_ids = input_ids[:, prefix_length:] - - group_size = torch.zeros( - self.config.num_experts, dtype=torch.long, device="cpu" - ) - for i in range(self.config.num_experts): - group_size[i] = (postfix_moe_token_types == i).sum() - - # Calculate start and end indices for each expert group - postfix_start_indices = torch.cumsum(group_size, dim=0) - group_size - postfix_end_indices = torch.cumsum(group_size, dim=0) - - pad_token_id = self.processor.tokenizer.pad_token_id - padding_mask = input_ids == pad_token_id - - # prefix_length, postfix_length = prefix_indices.shape[0], postfix_indices.shape[0] - - postfix_length = input_ids.shape[-1] - prefix_length - _postfix_attention_mask = torch.ones( - (batch_size, postfix_length, prefix_length + postfix_length), - dtype=torch.bool, - device=postfix_attention_mask.device, - ) - - # Use a padding mask to set the corresponding rows and columns to false. - # Get the padding mask for the postfix portion. - postfix_padding_mask = padding_mask[ - :, prefix_length: - ] # [batch_size, postfix_length] - full_padding_mask = padding_mask # [batch_size, prefix_length + postfix_length] - - # causal mask for postfix attention - if self.config.causal_action_attention_mask: - _postfix_attention_mask[:, :, prefix_length:] = torch.tril( - torch.ones( - (postfix_length, postfix_length), - dtype=torch.bool, - device=postfix_attention_mask.device, - ) - ) - - for batch_idx in range(padding_mask.shape[0]): - # Set the rows corresponding to the padding positions to False (where the query position is padding). - _postfix_attention_mask[batch_idx, postfix_padding_mask[batch_idx], :] = ( - False - ) - # Set the columns corresponding to the padding positions to False (the key position is the padding). - _postfix_attention_mask[batch_idx, :, full_padding_mask[batch_idx]] = False - - timing_results["cache_preprocessing"] = time.time() - cache_prep_start_time + v_padding = None + if (not ctx["dof_mask"].all()) and (ctx["padding_action"] is not None): + v_padding = ctx["padding_action"] - ctx["noisy_action"] + v_0 = (v_padding) * (1 - ctx["dof_mask"]) + v_0 * ctx["dof_mask"] + dt = ctx["times"][1] - ctx["times"][0] + ctx["noisy_action"] = ctx["noisy_action"] + dt * v_0 - ode_start_time = time.time() + ( + postfix_position_ids, + postfix_inputs_embeds, + _postfix_attention_mask, + postfix_moe_token_types, + postfix_input_ids, + postfix_start_indices, + postfix_end_indices, + padding_mask, + ) = self._prepare_flow_postfix(ctx, prefix_kv_cache) def step_with_kvcache(timestep, noisy_action): action_mask = ( @@ -2297,10 +2178,11 @@ def step_with_kvcache(timestep, noisy_action): assert action_mask.any(), "No action token found in input_ids" timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0]) action_embed, adarms_cond = self.action_preprocessor.step( - timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask + timestep=timestep, noisy_action=noisy_action, dof_mask=ctx["dof_mask"] ) action_embed = action_embed.reshape(-1, postfix_inputs_embeds.shape[-1]) + # Create a temporary inputs_embeds copy for thread safety temp_inputs_embeds = postfix_inputs_embeds.clone() temp_inputs_embeds[action_mask] = action_embed.to(temp_inputs_embeds.dtype) transformer_outputs = self.model( @@ -2313,81 +2195,173 @@ def step_with_kvcache(timestep, noisy_action): start_indices=postfix_start_indices, end_indices=postfix_end_indices, use_cache=False, - output_attentions=False, + output_attentions=output_attentions, output_hidden_states=False, return_dict=True, adarms_conds=[None, adarms_cond], ) - + if output_attentions: + all_attention_maps.append(transformer_outputs.attentions) hidden_states = transformer_outputs.last_hidden_state action_hidden_states = hidden_states[action_mask].to(torch.float32) action_pred = self.action_preprocessor.action_proj_back( action_hidden_states[:, : self.action_preprocessor.action_hidden_size] ) if getattr(self.config, "use_x_pred", False): - v_t = action_pred - noise.reshape(-1, noise.shape[-1]) + # Align noisy_action (and timestep) to action_pred's shape + B, action_horizon, action_dim = noisy_action.shape + noisy_action_flat = noisy_action.reshape( + -1, action_dim + ) # [B * action_horizon, action_dim] + timestep_expand = timestep.view(-1, 1).repeat_interleave( + action_horizon, dim=0 + ) # [B * action_horizon, 1] + v_t = (action_pred - noisy_action_flat) / torch.clamp( + 1 - timestep_expand, min=0.05 + ) else: v_t = action_pred - return v_t.reshape(batch_size, action_horizon, action_dim) + v_t = v_t.reshape( + ctx["batch_size"], ctx["action_horizon"], ctx["action_dim"] + ) + + if (not ctx["dof_mask"].all()) and (ctx["padding_action"] is not None): + v_t = (v_padding) * (1 - ctx["dof_mask"]) + v_t * ctx["dof_mask"] + + return v_t action_trajectory = odeint( - step_with_kvcache, noisy_action, times[1:], method="euler" + step_with_kvcache, ctx["noisy_action"], ctx["times"][1:], method="euler" ) - timing_results["ode_integration"] = time.time() - ode_start_time - - postprocess_start_time = time.time() - predict_action = action_trajectory[-1] - if unnorm: - predict_action = ( - self.action_preprocessor.normalizer_action.unnormalize_data( - predict_action, dataset_names - ) - ) - output["predict_action"] = predict_action - # normalize action chunk to get gt_action - if action_chunk is not None: - output["gt_action"] = ( - self.action_preprocessor.normalizer_action.unnormalize_data( - action_chunk, dataset_names + attention_maps = None + if output_attentions and ctx["img_mask"] is not None: + attention_maps = [torch.stack(map, dim=0) for map in all_attention_maps] + attention_maps = torch.stack( + attention_maps, dim=0 + ) # [flow steps, layer depths, action tokens, all tokens] + attention_maps = attention_maps.mean(2) + image_mask_indices = ctx["img_mask"][0].nonzero(as_tuple=True)[0] + attention_maps = attention_maps[:, :, image_mask_indices] + return action_trajectory, attention_maps + + def _prepare_flow_postfix(self, ctx, prefix_kv_cache=None): + postfix_position_ids = ctx["position_ids"][:, :, ctx["prefix_length"] :] + postfix_inputs_embeds = ctx["inputs_embeds"][:, ctx["prefix_length"] :, :] + postfix_attention_mask = ctx["attention_mask"][:, ctx["prefix_length"] :] + postfix_moe_token_types = ctx["moe_token_types"][:, ctx["prefix_length"] :] + postfix_input_ids = ctx["input_ids"][:, ctx["prefix_length"] :] + + postfix_start_indices = None + postfix_end_indices = None + padding_mask = None + + if prefix_kv_cache is not None: + if ctx["prefix_length"] is None: + has_true = ctx["flow_action_mask"].any(dim=1) + prefix_length = torch.argmax( + ctx["flow_action_mask"].float(), dim=1, keepdim=True ) + prefix_length[~has_true] = ctx["flow_action_mask"].shape[1] + # check if prefix_length is the same for all batch + if not torch.all(prefix_length == prefix_length[0]): + raise ValueError( + "prefix_length differs across batch; batch prompts must align" + ) + prefix_length = int(prefix_length[0].item()) + ctx["prefix_length"] = prefix_length + + if hasattr(prefix_kv_cache, "key_cache"): + for layer_i in range(len(prefix_kv_cache.key_cache)): + prefix_kv_cache.key_cache[layer_i] = prefix_kv_cache.key_cache[ + layer_i + ][:, :, : ctx["prefix_length"], :] + prefix_kv_cache.value_cache[layer_i] = prefix_kv_cache.value_cache[ + layer_i + ][:, :, : ctx["prefix_length"], :] + else: + for layer_i in range(len(prefix_kv_cache.layers)): + prefix_kv_cache.layers[layer_i].keys = prefix_kv_cache.layers[ + layer_i + ].keys[:, :, : ctx["prefix_length"], :] + prefix_kv_cache.layers[layer_i].values = prefix_kv_cache.layers[ + layer_i + ].values[:, :, : ctx["prefix_length"], :] + + group_size = torch.zeros( + self.config.num_experts, dtype=torch.long, device="cpu" ) + for i in range(self.config.num_experts): + group_size[i] = (postfix_moe_token_types == i).sum() - timing_results["postprocessing"] = time.time() - postprocess_start_time - timing_results["total_time"] = time.time() - total_start_time + # Calculate start and end indices for each expert group + postfix_start_indices = torch.cumsum(group_size, dim=0) - group_size + postfix_end_indices = torch.cumsum(group_size, dim=0) - output["timing_results"] = timing_results + pad_token_id = self.processor.tokenizer.pad_token_id + padding_mask = ctx["input_ids"] == pad_token_id + else: + pad_token_id = self.processor.tokenizer.pad_token_id + padding_mask = ctx["input_ids"] == pad_token_id - return output + postfix_length = ctx["input_ids"].shape[-1] - ctx["prefix_length"] + _postfix_attention_mask = torch.ones( + (ctx["batch_size"], postfix_length, ctx["prefix_length"] + postfix_length), + dtype=torch.bool, + device=postfix_attention_mask.device, + ) - def forward( - self, mode: Optional[str] = None, predict_mode: Optional[str] = "text", **kwargs - ): - """ - Main forward pass dispatcher for different execution modes. + postfix_padding_mask = padding_mask[:, ctx["prefix_length"] :] + full_padding_mask = padding_mask - This method routes execution to appropriate forward functions based on the specified mode: - - No mode (None): Training step with gradient disabled - - 'predict': Prediction/inference mode - - 'train': Training mode with gradients enabled - - 'validate': Validation mode with gradients disabled + if self.config.causal_action_attention_mask: + _postfix_attention_mask[:, :, ctx["prefix_length"] :] = torch.tril( + torch.ones( + (postfix_length, postfix_length), + dtype=torch.bool, + device=postfix_attention_mask.device, + ) + ) - Args: - mode (str, optional): Execution mode. If None, defaults to training step without gradients - predict_mode (str, optional): Prediction mode for 'predict' mode ("text", "fast", or "diffusion") - **kwargs: Additional arguments passed to the selected forward function + # postfix_padding_mask: [B, postfix_length], True where padding (query rows) + # full_padding_mask: [B, prefix_length+postfix_length], True where padding (key cols) + _postfix_attention_mask.masked_fill_(postfix_padding_mask.unsqueeze(2), False) + _postfix_attention_mask.masked_fill_(full_padding_mask.unsqueeze(1), False) + + return ( + postfix_position_ids, + postfix_inputs_embeds, + _postfix_attention_mask, + postfix_moe_token_types, + postfix_input_ids, + postfix_start_indices, + postfix_end_indices, + padding_mask, + ) - Returns: - Model outputs appropriate for the selected mode + def _finalize_flow_action_output(self, action_trajectory, ctx): + predict_action = action_trajectory[-1] + predict_action = unnormalize_data_with_virtual_tail( + self.action_preprocessor.normalizer_action, + predict_action, + ctx["dataset_names"], + ctx.get("dof_mask"), + ) + output = {"predict_action": predict_action} + if ctx["action_chunk"] is not None: + output["gt_action"] = unnormalize_data_with_virtual_tail( + self.action_preprocessor.normalizer_action, + ctx["action_chunk"], + ctx["dataset_names"], + ctx.get("dof_mask"), + ) - Todo: - - Add support for distinguishing multi-modal data types in prediction mode - """ + return output + + def forward(self, mode: Optional[str] = None, **kwargs): if not mode: with torch.no_grad(): return self.train_step_forward(**kwargs) - elif mode == "predict": - return self.generate_flow_action(predict_mode=predict_mode, **kwargs) elif mode == "train": return self.train_step_forward(use_cache=False, **kwargs) elif mode == "validate": @@ -2403,6 +2377,8 @@ def prepare_inputs_for_generation( attention_mask=None, inputs_embeds=None, moe_token_types=None, + start_indices=None, + end_indices=None, cache_position=None, position_ids=None, use_cache=True, @@ -2417,95 +2393,60 @@ def prepare_inputs_for_generation( agent_pos_mask=None, **kwargs, ): - """ - Prepare inputs for autoregressive generation with multi-modal support. - - This method handles input preparation for generation, including proper slicing of inputs - based on cache position, MoE token type management, and multi-modal data handling. - Vision inputs are selectively forwarded only when needed during generation. - - Args: - input_ids: Input token IDs - past_key_values: Cached key-value pairs from previous generation steps - attention_mask: Attention mask for input tokens - inputs_embeds: Pre-computed input embeddings - moe_token_types: Token type assignments for MoE routing - cache_position: Current cache position for generation - position_ids: Position IDs for tokens - use_cache: Whether to use key-value caching - pixel_values: Image pixel values - pixel_values_videos: Video pixel values - image_grid_thw: Image grid dimensions - video_grid_thw: Video grid dimensions - second_per_grid_ts: Time interval per temporal grid - dataset_names: Dataset names for processing - proprioception: Proprioceptive sensor data - dof_mask: Degrees of freedom mask - agent_pos_mask: Agent position mask - **kwargs: Additional arguments - - Returns: - dict: Prepared model inputs for generation step + # Overwritten -- in specific circumstances we don't want to forward image inputs to the model - Todo: - - Test this function thoroughly with various input configurations + # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens + # Exception 1: when passing input_embeds, input_ids may be missing entries + # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here + # Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case. + # (we can't check exception 3 while compiling) + # Exception 4: If input_embeds are passed then slice it through `cache_position`, to keep only the unprocessed tokens and + # generate the first token for each sequence. Later use the generated Input ids for continuation. - Note: - This is an overridden method that handles specific cases for multi-modal generation: - - Slices input_ids through cache_position to keep only unprocessed tokens - - Handles special cases for input_embeds, generation methods, and GPU synchronization - - Manages vision inputs to avoid unnecessary forward passes - """ - # Initialize MoE token types if not provided if moe_token_types is None: moe_token_types = torch.zeros_like( input_ids - ) # FIXME: Handle case when input_embeds is used instead + ) # FIXME input_ids or input_embeds else: - # Ensure moe_token_types length matches input_ids + # Ensure moe_token_types has the same length as input_ids if moe_token_types.shape[1] < input_ids.shape[1]: - # Calculate required padding length + # Compute the required padding length pad_length = input_ids.shape[1] - moe_token_types.shape[1] - # Create padding tensor with default token type (0) + # Create the padding tensor using 0 as the default token type pad_tensor = torch.zeros( (moe_token_types.shape[0], pad_length), dtype=moe_token_types.dtype, device=moe_token_types.device, ) - # Concatenate padding to existing moe_token_types + # Append padding after the existing moe_token_types moe_token_types = torch.cat([moe_token_types, pad_tensor], dim=1) - # Handle input slicing based on cache state and special cases if past_key_values is not None: - if ( - inputs_embeds is not None and input_ids.shape[1] == 0 - ): # Exception 4: input_embeds case + if inputs_embeds is not None and input_ids.shape[1] == 0: # Exception 4 inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :] moe_token_types = moe_token_types[:, -cache_position.shape[0] :] - elif inputs_embeds is not None or ( # Exception 1: input_embeds provided + elif inputs_embeds is not None or ( # Exception 1 is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1] - ): # Exception 3: GPU sync edge case + ): # Exception 3 input_ids = input_ids[:, -cache_position.shape[0] :] moe_token_types = moe_token_types[:, -cache_position.shape[0] :] elif ( input_ids.shape[1] != cache_position.shape[0] - ): # Default case (Exception 2 is no-op) + ): # Default case (the "else", a no op, is Exception 2) cache_pos = cache_position.clone() input_ids = input_ids[:, cache_pos] moe_token_types = moe_token_types[:, cache_pos] - # Skip vision inputs for continuation steps (not initial generation) if cache_position[0] != 0: pixel_values = None pixel_values_videos = None - # Determine whether to use inputs_embeds or input_ids for this generation step + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]: model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None} else: model_inputs = {"input_ids": input_ids, "inputs_embeds": None} - # Prepare 4D causal attention mask for static cache if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2: if model_inputs["inputs_embeds"] is not None: batch_size, sequence_length, _ = inputs_embeds.shape @@ -2528,12 +2469,19 @@ def prepare_inputs_for_generation( ) ) - # Assemble all model inputs for generation + if ( + self.config._attn_implementation == "flash_attention_2_ki" + and self.config.attention_moe is True + ): + attention_mask = None # Pass None as the mask to use the official flash attention for autoregressive prediction + model_inputs.update( { "position_ids": position_ids, "past_key_values": past_key_values, "moe_token_types": moe_token_types, + "start_indices": start_indices, + "end_indices": end_indices, "use_cache": use_cache, "attention_mask": attention_mask, "pixel_values": pixel_values, @@ -2555,30 +2503,25 @@ def _get_image_nums_and_video_nums( input_ids: Optional[torch.LongTensor], ) -> Tuple[torch.Tensor, torch.Tensor]: """ - Get the number of images and videos for each sample to calculate tensor separation lengths. - - These parameters are computed directly from input_ids rather than being passed through - the processor to avoid unpredictable impacts from interface modifications. + Get the number of images and videos for each sample to calculate the separation length of the sample tensor. + These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications. Args: - input_ids (torch.LongTensor): Input token IDs of shape (batch_size, sequence_length) + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Returns: - tuple: - - image_nums (torch.LongTensor): Number of images per sample - - video_nums (torch.LongTensor): Number of videos per sample + image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`) + video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`) """ image_token_id = self.config.image_token_id video_token_id = self.config.video_token_id vision_start_token_id = self.config.vision_start_token_id - # Find vision start tokens and their following tokens vision_start_mask = input_ids == vision_start_token_id vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1) image_mask = input_ids == image_token_id video_mask = input_ids == video_token_id - - # Count images and videos following vision start tokens image_nums = torch.sum(vision_first_mask & image_mask, dim=1) video_nums = torch.sum(vision_first_mask & video_mask, dim=1) @@ -2591,28 +2534,14 @@ def _expand_inputs_for_generation( input_ids: Optional[torch.LongTensor] = None, **model_kwargs, ) -> Tuple[torch.LongTensor, Dict[str, Any]]: - """ - Expand inputs for generation with support for multi-modal tensors. - - This is an overridden method that supports expanding tensors without a standard batch - size dimension, specifically for vision-related tensors: - - pixel_values.shape[0] = sum(sequence_lengths for all image samples) - - image_grid_thw.shape[0] = sum(num_images for all samples) - - Similar patterns for video tensors + # Overwritten -- Support for expanding tensors without a batch size dimension + # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t + # pixel_values.shape[0] is sum(seqlen_images for samples) + # image_grid_thw.shape[0] is sum(num_images for samples) - Args: - expand_size (int): Factor by which to expand inputs (for beam search, etc.) - is_encoder_decoder (bool): Whether using encoder-decoder architecture - input_ids (torch.LongTensor, optional): Input token IDs - **model_kwargs: Additional model arguments to expand - - Returns: - tuple: (expanded_input_ids, expanded_model_kwargs) - """ if expand_size == 1: return input_ids, model_kwargs - # Define keys for vision-related tensors that need special handling visual_keys = [ "pixel_values", "image_grid_thw", @@ -2622,13 +2551,11 @@ def _expand_inputs_for_generation( ] def _expand_dict_for_generation_visual(dict_to_expand): - """Expand vision-related tensors based on image/video counts per sample.""" image_grid_thw = model_kwargs.get("image_grid_thw", None) video_grid_thw = model_kwargs.get("video_grid_thw", None) image_nums, video_nums = self._get_image_nums_and_video_nums(input_ids) def _repeat_interleave_samples(x, lengths, repeat_times): - """Split tensor by lengths and repeat each sample.""" samples = torch.split(x, lengths) repeat_args = [repeat_times] + [1] * (x.dim() - 1) result = torch.cat( @@ -2638,33 +2565,31 @@ def _repeat_interleave_samples(x, lengths, repeat_times): for key in dict_to_expand: if key == "pixel_values": - # Split images into samples and compute sequence lengths + # split images into samples samples = torch.split(image_grid_thw, list(image_nums)) + # compute the sequence length of images for each sample lengths = [torch.prod(sample, dim=1).sum() for sample in samples] dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "image_grid_thw": - # Expand based on number of images per sample + # get the num of images for each sample lengths = list(image_nums) dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "pixel_values_videos": - # Split videos into samples and compute sequence lengths samples = torch.split(video_grid_thw, list(video_nums)) lengths = [torch.prod(sample, dim=1).sum() for sample in samples] dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "video_grid_thw": - # Expand based on number of videos per sample lengths = list(video_nums) dict_to_expand[key] = _repeat_interleave_samples( dict_to_expand[key], lengths=lengths, repeat_times=expand_size ) elif key == "second_per_grid_ts": - # Handle list-type temporal grid data if not isinstance(dict_to_expand[key], list): raise TypeError( f"Expected value for key '{key}' to be a list, but got {type(dict_to_expand[key])} instead." @@ -2678,7 +2603,6 @@ def _repeat_interleave_samples(x, lengths, repeat_times): return dict_to_expand def _expand_dict_for_generation(dict_to_expand): - """Expand standard tensors using repeat_interleave.""" for key in dict_to_expand: if ( key != "cache_position" @@ -2691,19 +2615,16 @@ def _expand_dict_for_generation(dict_to_expand): ) return dict_to_expand - # Expand visual inputs only if input_ids is available for counting images/videos - # If input_ids is unavailable, visual inputs won't be used, so no expansion needed + # input_ids is required for expanding visual inputs + # If input_ids is unavailable, visual inputs will not be used; therefore, there is no need to expand visual inputs. if input_ids is not None and input_ids.numel() != 0: model_kwargs = _expand_dict_for_generation_visual(model_kwargs) - # Expand input_ids using standard repeat_interleave if input_ids is not None: input_ids = input_ids.repeat_interleave(expand_size, dim=0) - # Expand all other model arguments model_kwargs = _expand_dict_for_generation(model_kwargs) - # Handle encoder-decoder specific expansion if is_encoder_decoder: if model_kwargs.get("encoder_outputs") is None: raise ValueError( @@ -2714,3 +2635,431 @@ def _expand_dict_for_generation(dict_to_expand): ) return input_ids, model_kwargs + + def rename_vlm_weights_for_vla(self, merged_weights): + renamed = {} + + for key, value in merged_weights.items(): + + if key.startswith("model.layers") and "mlp." in key and self.config.mlp_moe: + layer_num = key.split(".layers.")[1].split(".mlp")[0] + new_key = key.replace( + f"layers.{layer_num}.mlp.", f"layers.{layer_num}.moe.experts.0." + ) + renamed[new_key] = value + continue + + if ( + key.startswith("model.layers") + and "self_attn." in key + and self.config.attention_moe + ): + layer_num = key.split(".layers.")[1].split(".self_attn")[0] + proj_types = ["q_proj", "k_proj", "v_proj", "o_proj"] + for proj in proj_types: + if proj in key: + new_key = key.replace( + f"layers.{layer_num}.self_attn.{proj}", + f"layers.{layer_num}.self_attn.{proj}_experts.0", + ) + renamed[new_key] = value + break + continue + + if self.config.norm_moe and ".input_layernorm." in key: + renamed[key.replace("input_layernorm", "input_layernorms.0")] = value + continue + if self.config.norm_moe and ".post_attention_layernorm." in key: + renamed[ + key.replace( + "post_attention_layernorm", "post_attention_layernorms.0" + ) + ] = value + continue + if self.config.norm_moe and ".norm." in key: + renamed[key.replace("norm", "norms.0")] = value + continue + + renamed[key] = value + + fused = Qwen2_5_VLMoEForAction.convert_to_fused(renamed) + + return fused + + @staticmethod + def fuse_gate_up( + fused, prefix, suffix_gate="gate_proj", suffix_up="up_proj", out="gate_up_proj" + ): + gate_w = fused.get(prefix + f"{suffix_gate}.weight") + up_w = fused.get(prefix + f"{suffix_up}.weight") + gate_b = fused.get(prefix + f"{suffix_gate}.bias") + up_b = fused.get(prefix + f"{suffix_up}.bias") + + # Skip fusion if gate_up_proj already exists + if prefix + f"{out}.weight" in fused: + return + + if gate_w is not None and up_w is not None: + fused[prefix + f"{out}.weight"] = torch.cat([gate_w, up_w], dim=0) + if gate_b is not None and up_b is not None: + fused[prefix + f"{out}.bias"] = torch.cat([gate_b, up_b], dim=0) + + # Remove the old modules + for n in [ + f"{suffix_gate}.weight", + f"{suffix_up}.weight", + f"{suffix_gate}.bias", + f"{suffix_up}.bias", + ]: + full = prefix + n + if full in fused: + del fused[full] + + @staticmethod + def fuse_qkv(fused, prefix, *, experts=False, expert_id=None): + """ + prefix: + - non-MoE: model.layers.X.self_attn. + - MoE: model.layers.X.self_attn. + + When experts=True, expert_id must be provided + """ + + if experts: + assert expert_id is not None, "expert_id must be provided for MoE attention" + + eid = expert_id + + q_w = fused.get(prefix + f"q_proj_experts.{eid}.weight") + k_w = fused.get(prefix + f"k_proj_experts.{eid}.weight") + v_w = fused.get(prefix + f"v_proj_experts.{eid}.weight") + + q_b = fused.get(prefix + f"q_proj_experts.{eid}.bias") + k_b = fused.get(prefix + f"k_proj_experts.{eid}.bias") + v_b = fused.get(prefix + f"v_proj_experts.{eid}.bias") + + out_w = f"qkv_proj_experts.{eid}.weight" + out_b = f"qkv_proj_experts.{eid}.bias" + + remove_list = [ + f"q_proj_experts.{eid}.weight", + f"k_proj_experts.{eid}.weight", + f"v_proj_experts.{eid}.weight", + f"q_proj_experts.{eid}.bias", + f"k_proj_experts.{eid}.bias", + f"v_proj_experts.{eid}.bias", + ] + + else: + q_w = fused.get(prefix + "q_proj.weight") + k_w = fused.get(prefix + "k_proj.weight") + v_w = fused.get(prefix + "v_proj.weight") + + q_b = fused.get(prefix + "q_proj.bias") + k_b = fused.get(prefix + "k_proj.bias") + v_b = fused.get(prefix + "v_proj.bias") + + out_w = "qkv_proj.weight" + out_b = "qkv_proj.bias" + + remove_list = [ + "q_proj.weight", + "k_proj.weight", + "v_proj.weight", + "q_proj.bias", + "k_proj.bias", + "v_proj.bias", + ] + + if q_w is None or k_w is None or v_w is None: + return + + fused[prefix + out_w] = torch.cat([q_w, k_w, v_w], dim=0) + + if q_b is not None and k_b is not None and v_b is not None: + fused[prefix + out_b] = torch.cat([q_b, k_b, v_b], dim=0) + + for n in remove_list: + full = prefix + n + if full in fused: + del fused[full] + + @staticmethod + def is_fused(state_dict): + # Check whether the model is already fused by looking for fused weight markers + return any(".moe.experts.0.gate_up_proj" in key for key in state_dict.keys()) + + @staticmethod + def convert_to_fused(state_dict): + """ + Convert an unfused checkpoint to a fused checkpoint + """ + fused = {} + + # ========================= + # Phase 1: copy everything first + # ========================= + for k, v in state_dict.items(): + fused[k] = v + + # ========================= + # Phase 2: fuse gate + up (MoE aware) + # ========================= + for key in list(fused.keys()): + # language MoE mlp (ANY expert) + m = re.match( + r"(model\.layers\.\d+\.moe\.experts\.\d+\.)gate_proj\.weight", + key, + ) + if m: + prefix = m.group(1) + Qwen2_5_VLMoEForAction.fuse_gate_up(fused, prefix) + continue + + # language non-MoE mlp + if re.match(r"(model\.layers\.\d+\.mlp\.)gate_proj\.weight", key): + prefix = key.replace("gate_proj.weight", "") + Qwen2_5_VLMoEForAction.fuse_gate_up(fused, prefix) + continue + + # visual mlp + if re.match(r"(visual\.blocks\.\d+\.mlp\.)gate_proj\.weight", key): + prefix = key.replace("gate_proj.weight", "") + Qwen2_5_VLMoEForAction.fuse_gate_up(fused, prefix) + + # ========================= + # Phase 3: fuse attention qkv (MoE aware) + # ========================= + for key in list(fused.keys()): + # language MoE attention: ANY expert + m = re.match( + r"(model\.layers\.\d+\.self_attn\.)q_proj_experts\.(\d+)\.weight", + key, + ) + if m: + prefix = m.group(1) + expert_id = int(m.group(2)) + Qwen2_5_VLMoEForAction.fuse_qkv( + fused, + prefix, + experts=True, + expert_id=expert_id, + ) + continue + + # language non-MoE attention + if re.match(r"(model\.layers\.\d+\.self_attn\.)q_proj\.weight", key): + prefix = key.replace("q_proj.weight", "") + Qwen2_5_VLMoEForAction.fuse_qkv( + fused, + prefix, + experts=False, + ) + continue + + # visual attention + if re.match(r"(visual\.blocks\.\d+\.attn\.)q_proj\.weight", key): + prefix = key.replace("q_proj.weight", "") + Qwen2_5_VLMoEForAction.fuse_qkv( + fused, + prefix, + experts=False, + ) + + # ========================= + # Phase 4: sanity check（optional） + # ========================= + for k in fused.keys(): + assert not any( + x in k + for x in [ + ".gate_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ] + ), f"Unfused key still exists: {k}" + + return fused + + def convert_to_mix_precision(self): + # Mix Precision + params_to_keep_float32 = [] + for name, _ in self.named_parameters(): + if any( + k in name + for k in [ + "input_layernorm", + "post_attention_layernorm", + "model.norm", + "action_preprocessor", + ] + ): + params_to_keep_float32.append(name) + for name, param in self.named_parameters(): + if name not in params_to_keep_float32: + param.data = param.data.to(torch.bfloat16) + if name in params_to_keep_float32: + param.data = param.data.to(torch.float32) + + def convert_to_fsdp( + self, + *, + mesh, + mp_policy, + offload_policy=None, + reshard_after_forward: bool = True, + use_dmuon: bool = False, + ): + """Wrap with FSDP2: per-decoder-layer + per-vision-block fully_shard + with bf16 ``mp_policy``; layernorms and ``action_preprocessor`` leaves + get their own nested ``fully_shard`` with an fp32 policy so their + forward stays in fp32 (matching the old FSDP1 NO_SHARD + fp32_policy + layout). + + ``mesh`` and ``mp_policy`` are required and produced by + ``FSDPStrategy._build_fsdp2_layout``. The pre-migration FSDP1 wrap + used ``ShardingStrategy.SHARD_GRAD_OP`` regardless of the yaml + config; the new path honors ``fsdp_sharding_strategy`` (default + ``full_shard`` → ``reshard_after_forward=True``). Active runs that + relied on the implicit SHARD_GRAD_OP throughput should set + ``distributed.fsdp_sharding_strategy: shard_grad_op`` explicitly. + + When ``use_dmuon=True``, DMuon owns selected trainable + matrix weights before the remaining parameters are wrapped by FSDP2. + """ + from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy + + if use_dmuon: + import dmuon + + # DMuon keeps fp32 master weights for dedicated params; symmetric + # FSDP2 params still use MixedPrecisionPolicy for bf16 compute. + for param in self.parameters(): + if param.requires_grad and param.dtype != torch.float32: + param.data = param.data.float() + + action_preprocessor_linear_ids = { + id(module) + for module in self.action_preprocessor.modules() + if isinstance(module, nn.Linear) + } + + def hook_boundary(module: nn.Module) -> bool: + # ActionProcessor helper methods call internal Linear layers + # directly, bypassing ActionProcessor.forward(). Those internal + # Linear layers need their own hooks. + if id(module) in action_preprocessor_linear_ids: + return True + return isinstance( + module, + ( + Qwen2_5_VLDecoderLayer_with_MoE, + Qwen2_5_VLVisionBlock, + Qwen2_5_VLPatchMerger, + type(self.action_preprocessor), + ), + ) + + assignment = dmuon.dedicate_params( + self, + mesh, + predicate=_is_qwen25_dmuon_target_param, + hook_boundary_predicate=hook_boundary, + hook_boundary_strict=True, + compute_dtype=torch.bfloat16, + reshard_after_forward=reshard_after_forward, + ) + logger.info( + "[Qwen2.5 FSDP2 + DMuon] dedicate_params: %d params assigned", + len(assignment), + ) + + fp32_mp = MixedPrecisionPolicy( + param_dtype=torch.float32, + reduce_dtype=torch.float32, + ) + shard_kwargs = dict( + mesh=mesh, + mp_policy=mp_policy, + reshard_after_forward=reshard_after_forward, + ) + if offload_policy is not None: + shard_kwargs["offload_policy"] = offload_policy + fp32_shard_kwargs = dict(shard_kwargs, mp_policy=fp32_mp) + + def fully_shard_fp32_leaf_or_container(name: str, module: nn.Module) -> None: + # FSDP2 cannot shard containers without forward(), e.g. ModuleList. + # When a norm match resolves to a container, shard only norm leaves. + def is_norm_leaf(leaf_name: str, leaf: nn.Module) -> bool: + return ( + "norm" in leaf_name.lower() or "norm" in type(leaf).__name__.lower() + ) + + if isinstance(module, (nn.ModuleList, nn.ModuleDict)): + for leaf_name, leaf in module.named_modules(): + if not leaf_name: + continue + if any(True for _ in leaf.children()): + continue + if not is_norm_leaf(leaf_name, leaf): + continue + logger.info( + "[FSDP2] fully_shard fp32 layernorm: %s.%s", + name, + leaf_name, + ) + fully_shard(leaf, **fp32_shard_kwargs) + return + + logger.info("[FSDP2] fully_shard fp32 layernorm: %s", name) + fully_shard(module, **fp32_shard_kwargs) + + # Layernorm leaves inside decoder layers / vision blocks / model.norm + # — wrap with fp32 mp_policy BEFORE wrapping their parents so the + # nested policy survives. Match-by-FQN mirrors the FSDP1 logic. + for module_name, module in list(self.named_modules()): + for child_name, child in list(module.named_children()): + if any( + k in child_name.lower() + for k in ("input_layernorm", "post_attention_layernorm") + ) or ("norm" in child_name.lower() and module_name.endswith("model")): + fully_shard_fp32_leaf_or_container( + f"{module_name}.{child_name}", child + ) + + # action_preprocessor leaves: same fp32 treatment as layernorms. + for name, module in list(self.named_modules()): + if "action_preprocessor" not in name.lower(): + continue + if any(True for _ in module.children()): + continue # only leaves + logger.info( + "[FSDP2] fully_shard fp32 action_preprocessor leaf: %s", + name, + ) + fully_shard(module, **fp32_shard_kwargs) + + # Each transformer decoder block: bf16 mp_policy from strategy. + for idx, layer in enumerate(self.model.layers): + fully_shard(layer, **shard_kwargs) + if idx == 0: + logger.info( + "[FSDP2] fully_shard model.layers.* (bf16, " + "reshard_after_forward=%s)", + reshard_after_forward, + ) + + # Each vision tower block: same. + if hasattr(self, "visual") and hasattr(self.visual, "blocks"): + for idx, block in enumerate(self.visual.blocks): + fully_shard(block, **shard_kwargs) + if idx == 0: + logger.info("[FSDP2] fully_shard visual.blocks.* (bf16)") + + # Root: covers lm_head, embed_tokens, and any params not in a + # nested fully_shard. Inner fp32 wraps maintain their own policy. + logger.info("[FSDP2] fully_shard root (bf16)") + fully_shard(self, **shard_kwargs) + return self diff --git a/wall_x/model/qact/tokenizer_mixin.py b/wall_x/model/qact/tokenizer_mixin.py new file mode 100644 index 00000000..a5bd3a65 --- /dev/null +++ b/wall_x/model/qact/tokenizer_mixin.py @@ -0,0 +1,963 @@ +"""Action tokenizer mixin for loading tokenizers and action mappings.""" + +import torch +import numpy as np +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, Tuple, Union +from transformers import AutoProcessor + +from wall_x.utils.constant import is_action_dataset_name + +# Delay imports so missing optional packages do not fail at import time +try: + from spatial_tokenizer.spatial_tokenizer import SpatialActionTokenizer +except ImportError: + SpatialActionTokenizer = None + + + +class ActionTokenizerMixin(ABC): + """Base class for action tokenizers""" + + def __init__(self): + self._tokenizer = None + self.action_normalizer = None + self._tokenizer_type: str = "" + self._action_mapper_cache: Optional[Dict] = None # Cached action_mapper + self.dllm = False + self.input_placeholder_flag = False + + @property + def tokenizer_type(self) -> str: + """Return the tokenizer type identifier""" + return self._tokenizer_type + + @property + def tokenizer(self): + """Return the underlying tokenizer instance""" + return self._tokenizer + + @property + def action_mapper(self) -> Optional[Dict]: + """Return the cached action_mapper""" + return self._action_mapper_cache + + @abstractmethod + def load_tokenizer(self, config: dict, normalizer, device: str = "cpu") -> Any: + """ + Load the tokenizer instance + + Args: + config: Configuration dictionary + normalizer: Normalizer + device: Device, usually "cpu" for training and "cuda" for inference + + Returns: + tokenizer instance + """ + pass + + @abstractmethod + def get_val_tokenizer(self, config: dict) -> Any: + """ + Get the tokenizer used for validation/inference + + Args: + config: Configuration dictionary + + Returns: + validation tokenizer instance + """ + pass + + @abstractmethod + def get_special_tokens(self) -> List[str]: + """ + Return special tokens to add to the vocabulary + + Returns: + list of token strings + """ + pass + + def get_all_special_tokens(self) -> Tuple[List[str], Optional[List[str]]]: + """ + Return all special tokens and keep <|action_token_0|> before AR action tokens + + Returns: + (new_tokens, special_tokens) tuple + """ + tokens, special_tokens = self.get_special_tokens() + if "<|action_token_0|>" not in tokens: + tokens.insert(0, "<|action_token_0|>") + return tokens, special_tokens + + @abstractmethod + def build_action_mapper(self, processor) -> Optional[Dict]: + """ + Build action_mapper + + Args: + processor: HuggingFace processor，used to convert tokens to IDs + + Returns: + action_mapper dictionary; format depends on tokenizer type + """ + pass + + @abstractmethod + def get_action_token_list(self, processor) -> List[int]: + """ + Get the action token ID list + + Args: + processor: HuggingFace processor + + Returns: + action token ID list + """ + pass + + @abstractmethod + def decode_action( + self, + output_ids: torch.Tensor, + action_mapper: Dict, + action_horizon: int, + action_dim: int, + device: torch.device, + proprioception: Optional[torch.Tensor] = None, + dof_mask: Optional[torch.Tensor] = None, + robot_type_id: Optional[int] = None, + state: Optional[torch.Tensor] = None, + ) -> Tuple[Optional[Union[np.ndarray, torch.Tensor]], bool]: + """ + Unified decoding interface + + Args: + output_ids: model output token IDs [1, seq_len] + action_mapper: action_mapper dictionary + action_horizon: action horizon + action_dim: action dimension + device: Device + proprioception: proprioception, normalized when required by a tokenizer + dof_mask: DOF mask when required by a tokenizer + robot_type_id: robot type ID when required by a tokenizer + state: state when required by fast/spatial tokenizers + + Returns: + (predict_action, decode_success) + - predict_action: decoded action [T, action_dim] or None + - decode_success: whether decoding succeeded + """ + pass + + @abstractmethod + def compute_accuracy( + self, + logits: torch.Tensor, + labels: torch.Tensor, + action_mapper: Dict, + action_token_id_set: Dict, + ) -> Dict[str, torch.Tensor]: + """ + Compute accuracy metrics + + Args: + logits: model output logits + labels: Labels + action_mapper: action_mapper dictionary + action_token_id_set: set of action token IDs + + Returns: + accuracy metric dictionary, such as {"action_accuracy": tensor, ...} + """ + pass + + @abstractmethod + def get_accuracy_keys(self) -> List[str]: + """ + Return accuracy metric keys for logging + + Returns: + list of key strings + """ + pass + + @property + @abstractmethod + def vocab_size(self) -> int: + """Return the vocabulary size""" + pass + + @property + @abstractmethod + def uses_dof_mask_for_unnorm(self) -> bool: + """ + Whether unnormalization requires dof_mask + + Returns: + True: dof_mask is required by fast/spatial tokenizers + False: dof_mask is not required and full dimensions are returned + """ + pass + + @property + @abstractmethod + def needs_action_crop(self) -> bool: + """ + Whether actions must be clipped before encoding + + Returns: + True: clip by chunk_size and dof_mask for fast/spatial tokenizers + False: no clipping; the encoder handles the full sequence internally + """ + pass + + @abstractmethod + def encode_to_tokens( + self, + actions: torch.Tensor, + obs_state: Optional[torch.Tensor] = None, + dof_mask: Optional[torch.Tensor] = None, + robot_type_ids: Optional[List[int]] = None, + is_train: bool = True, + ) -> List[List[str]]: + """ + Encode actions into token strings for training data processing + + Args: + actions: Normalized actions + - fast/spatial: List[Tensor], each [T, D] (clipped) + - v3.1 delta: Tensor [B, T, D] (full sequence) + obs_state: observation state when required by a tokenizer[B, obs_horizon, D] + dof_mask: DOF mask[B, T, D] + robot_type_ids: robot type ID list when required by a tokenizer + + Returns: + List[List[str]]: token string list for each sample + """ + pass + + def init_inference(self, robot_type: Optional[str] = None) -> None: + """ + Inference initialization hook; subclasses may override + + Args: + robot_type: robot type name + """ + pass + + def prepare_action_for_ar_encoding( + self, + ar_actionchunk: torch.Tensor, + dataset_names: List[str], + agent_pos: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Prepare action data before AR encoding; subclasses may override + + SpatialVLA needs waypoint selection; other tokenizers return normalized_action directly + + Args: + ar_actionchunk: [B, T, D] actions before normalization + dataset_names: [str] dataset name list + agent_pos: [B, T, D] agent positions required by SpatialVLA + + Returns: + prepared action data + """ + if self.action_normalizer is not None: + ar_actionchunk = self.action_normalizer.normalize_data( + ar_actionchunk, dataset_names + ) + return ar_actionchunk + + def get_robot_type_ids( + self, + uids: List[Optional[str]], + dataset_names: List[str], + ) -> Optional[List[int]]: + """ + Get robot_type_ids; subclasses may override + + Only the v3.1 delta tokenizer needs this; other tokenizers return None + + Args: + uids: UID list + dataset_names: dataset name list + + Returns: + robot_type_id list or None + """ + return None + + def modify_inputs_for_dllm( + self, + inputs: Dict[str, torch.Tensor], + processor, + sample_time: torch.Tensor, + dataset_names: List[str], + ): + return NotImplementedError + + +class FastTokenizerMixin(ActionTokenizerMixin): + """Fast tokenizer implementation""" + + def __init__(self): + super().__init__() + self._tokenizer_type = "fast" + + def load_tokenizer(self, config: dict, normalizer, device: str = "cpu") -> Any: + """Load the fast tokenizer""" + self._tokenizer = AutoProcessor.from_pretrained( + config["action_tokenizer_path"], trust_remote_code=True + ) + self.action_normalizer = normalizer + return self._tokenizer + + def get_val_tokenizer(self, config: dict) -> Any: + return self._tokenizer + + def get_special_tokens(self) -> List[str]: + """Return special tokens for the fast tokenizer""" + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + # tokens = ["<|ar_action|>", "<|ar_pad|>"] #temporary compatibility with existing checkpoints + tokens = [] + for i in range(self._tokenizer.vocab_size): + tokens.append(f"<|action_token_{i}|>") + return tokens, None + + def build_action_mapper(self, processor) -> Dict[int, int]: + """ + Build the cached action_mapper for the fast tokenizer + + Returns: + Dict[token_id, action_idx] + """ + # Return cached value if available + if self._action_mapper_cache is not None: + return self._action_mapper_cache + + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + + action_mapper = {} + for i in range(self._tokenizer.vocab_size): + token = f"<|action_token_{i}|>" + token_id = processor.tokenizer.convert_tokens_to_ids(token) + action_mapper[token_id] = i + + self._action_mapper_cache = action_mapper + return action_mapper + + def get_action_token_list(self, processor) -> List[int]: + """Get the action token ID list""" + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + + action_token_list = [] + for i in range(self._tokenizer.vocab_size): + token_id = processor.tokenizer.convert_tokens_to_ids( + f"<|action_token_{i}|>" + ) + action_token_list.append(token_id) + return action_token_list + + def decode_action( + self, + output_ids: torch.Tensor, + action_mapper: Dict, + action_horizon: int, + action_dim: int, + device: torch.device, + proprioception: Optional[torch.Tensor] = None, + dof_mask: Optional[torch.Tensor] = None, + robot_type_id: Optional[int] = None, + state: Optional[torch.Tensor] = None, + ) -> Tuple[Optional[Union[np.ndarray, torch.Tensor]], bool]: + """Fast tokenizer decoding""" + action_id = [] + for token_id_i in output_ids[0]: + if token_id_i.item() in action_mapper: + action_id.append(action_mapper[token_id_i.item()]) + + if len(action_id) == 0: + return np.zeros((action_horizon, action_dim)), False + + predict_action = self._tokenizer.decode( + [action_id], time_horizon=action_horizon, action_dim=action_dim + ) + + # Check whether decoding succeeded + decode_success = False + if isinstance(predict_action, np.ndarray): + decode_success = np.sum(predict_action) != 0 + elif isinstance(predict_action, torch.Tensor): + decode_success = predict_action.sum().item() != 0 + + return predict_action, decode_success + + def compute_accuracy( + self, + logits: torch.Tensor, + labels: torch.Tensor, + action_mapper: Dict, + action_token_id_set: Dict, + ) -> Dict[str, torch.Tensor]: + """Compute fast tokenizer accuracy""" + result = {} + + if len(action_token_id_set.get("action_token_list", [])) > 0: + shift_logits = logits[..., :-1, :].contiguous() + action_preds = shift_logits.argmax(dim=-1) + shift_labels = labels[..., 1:].contiguous() + action_mask = shift_labels > action_token_id_set["action_token_list"][0] + correct_preds = (action_preds == shift_labels) & action_mask + action_accuracy = correct_preds.sum().float() / action_mask.sum().float() + result["action_accuracy"] = action_accuracy + + return result + + def get_accuracy_keys(self) -> List[str]: + """Return fast tokenizer accuracy keys""" + return ["action_accuracy"] + + @property + def vocab_size(self) -> int: + if self._tokenizer is None: + return 0 + return self._tokenizer.vocab_size + + @property + def uses_dof_mask_for_unnorm(self) -> bool: + """fast tokenizer requires dof_mask""" + return True + + @property + def needs_action_crop(self) -> bool: + return True + + @property + def inference_ar_steps_for_dllm(self) -> int: + return self.max_length + + def encode_to_tokens( + self, + actions: List, + obs_state: Optional[torch.Tensor] = None, + dof_mask: Optional[torch.Tensor] = None, + robot_type_ids: Optional[List[int]] = None, + is_train: bool = True, + ) -> List[List[str]]: + """ + Fast tokenizer encoding + + Args: + actions: List[Tensor/ndarray], each [T, D] (clipped) + """ + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + + all_action_tokens = [] + for i in range(len(actions)): + action = actions[i] + if isinstance(action, torch.Tensor): + action = action.cpu().numpy() + token_id = self._tokenizer(action) + action_tokens = [f"<|action_token_{idx}|>" for idx in token_id[0]] + all_action_tokens.append(action_tokens) + return all_action_tokens + + def modify_inputs_for_dllm( + self, + inputs: Dict[str, torch.Tensor], + processor, + sample_time: torch.Tensor, + dataset_names: List[str], + ): + # Untested + input_ids = inputs["input_ids"] + labels = inputs["labels"] + prefix_length = inputs["prefix_length"] + bs, seqlen = input_ids.shape + + ar_token_length = self.max_length + ar_step_num = ar_token_length + + device = input_ids.device + dtype = input_ids.dtype + + placeholder_ids = torch.tensor( + processor.placeholder_seq, device=device, dtype=dtype + ) + + if not torch.is_tensor(sample_time): + sample_time = torch.tensor(sample_time, device=device, dtype=torch.float32) + else: + sample_time = sample_time.to(device=device, dtype=torch.float32) + + sample_time = sample_time.clamp(0.0, 1.0) + + noisy_steps_per_sample = ( + torch.ceil((1.0 - sample_time) * (ar_step_num + 1)).long() - 1 + ) + noisy_steps_per_sample = noisy_steps_per_sample.clamp(min=0, max=ar_step_num) + start = prefix_length - ar_token_length - 2 + end = prefix_length - 2 + + # Prepare the noise sequence for each sample + noise_seqs = placeholder_ids.repeat(bs, 1) + + ar_len = end - start + if noise_seqs.size(1) != ar_len: + # These should usually match; defensively truncate to ar_len + noise_seqs = noise_seqs[:, :ar_len] + rand = torch.rand(bs, ar_len, device=device) + perm = rand.argsort(dim=-1) + ranks = perm.argsort(dim=-1) + noisy_mask = ranks < noisy_steps_per_sample.view(-1, 1) + + ar_input = input_ids[:, start:end] + ar_labels = labels[:, start + 1 : end + 1] + ar_input[noisy_mask] = noise_seqs[noisy_mask] + ar_labels[~noisy_mask] = -100 + + inputs["input_ids"] = input_ids + inputs["labels"] = labels + return inputs + + def update_placeholder_mask(self, processor, prefix_length, input_ids): + # Untested + ar_action_mask = torch.zeros_like(input_ids) + + inc = torch.arange( + 1, + self.max_length + 1, + ) + inc = inc.unsqueeze(0).expand(ar_action_mask.size(0), -1) # [bs, ar_len] + ar_action_mask[ + :, + prefix_length - self.max_length - 2 : prefix_length - 2, + ] = inc + ar_action_mask[:, prefix_length - 2 : prefix_length] = -1 # eos + + return {"ar_action_mask": ar_action_mask} + + def get_placeholder_for_dllm(self): + placeholder_seq = ["<|ar_action|>"] * self.max_length + return placeholder_seq + + +class SpatialVLATokenizerMixin(ActionTokenizerMixin): + """SpatialVLA tokenizer implementation""" + + def __init__(self): + super().__init__() + self._tokenizer_type = "spatialvla" + + def load_tokenizer(self, config: dict, normalizer, device: str = "cpu") -> Any: + """Load the SpatialVLA tokenizer""" + if SpatialActionTokenizer is None: + raise ImportError( + "SpatialActionTokenizer is not installed. " + "Please install spatial_tokenizer package." + ) + self._tokenizer = SpatialActionTokenizer( + normalizer=normalizer, + augment_ratio=config.get("augment_ratio", 0.0), + max_waypoints=config.get("max_waypoints", 5), + with_gripper=config.get("with_gripper", True), + single_arm=config.get("single_arm", False), + ) + self._val_tokenizer = None + self.config = config + self.dllm = config.get("dllm", False) + self.input_placeholder_flag = config.get("input_placeholder_flag", False) + self.action_normalizer = normalizer + self.with_gripper = config.get("with_gripper", True) + return self._tokenizer + + def get_placeholder_for_dllm(self): + if self.with_gripper: + placeholder_seq = [ + "<|left_xyz|>", + "<|left_rpy|>", + "<|left_gripper|>", + "<|right_xyz|>", + "<|right_rpy|>", + "<|right_gripper|>", + ] + else: + placeholder_seq = [ + "<|left_xyz|>", + "<|left_rpy|>", + "<|right_xyz|>", + "<|right_rpy|>", + ] + if self._tokenizer.single_arm: + placeholder_seq = placeholder_seq[len(placeholder_seq) // 2 :] + placeholder_seq = placeholder_seq * self._tokenizer.max_waypoints + return placeholder_seq + + def get_val_tokenizer(self, config: dict) -> Any: + if self._val_tokenizer: + return self._val_tokenizer + + self._val_tokenizer = SpatialActionTokenizer( + normalizer=self.action_normalizer, + augment_ratio=0, + max_waypoints=self.config.get("max_waypoints", 5), + with_gripper=self.config.get("with_gripper", True), + single_arm=self.config.get("single_arm", False), + ) + return self._val_tokenizer + + def get_special_tokens(self) -> List[str]: + """Return special tokens for the SpatialVLA tokenizer""" + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + tokens = [ + "<|step|>", + "<|left|>", + "<|right|>", + "<|move|>", + ] # only for compatibility with existing checkpoints + if self.input_placeholder_flag: + special_tokens = [ + "<|left_xyz|>", + "<|left_rpy|>", + "<|left_gripper|>", + "<|right_xyz|>", + "<|right_rpy|>", + "<|right_gripper|>", + ] + tokens += special_tokens + if not self._tokenizer.with_gripper: + indices = [0, 1, 3, 4] + special_tokens = [special_tokens[i] for i in indices] + if self._tokenizer.single_arm: + special_tokens = special_tokens[len(special_tokens) // 2 :] + for i in range(self._tokenizer.vocab_size): + tokens.append(f"<|action_token_{i}|>") + + return tokens, special_tokens + + def build_action_mapper(self, processor) -> Dict[int, int]: + """ + Build the cached action_mapper for the SpatialVLA tokenizer + + Returns: + Dict[token_id, action_idx] + """ + # Return cached value if available + if self._action_mapper_cache is not None: + return self._action_mapper_cache + + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + + action_mapper = {} + for i in range(self._tokenizer.vocab_size): + token = f"<|action_token_{i}|>" + token_id = processor.tokenizer.convert_tokens_to_ids(token) + action_mapper[token_id] = i + + self._action_mapper_cache = action_mapper + return action_mapper + + def get_action_token_list(self, processor) -> List[int]: + """Get the action token ID list""" + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + + action_token_list = [] + for i in range(self._tokenizer.vocab_size): + token_id = processor.tokenizer.convert_tokens_to_ids( + f"<|action_token_{i}|>" + ) + action_token_list.append(token_id) + return action_token_list + + def decode_action( + self, + output_ids: torch.Tensor, + action_mapper: Dict, + action_horizon: int, + action_dim: int, + device: torch.device, + proprioception: Optional[torch.Tensor] = None, + dof_mask: Optional[torch.Tensor] = None, + robot_type_id: Optional[int] = None, + state: Optional[torch.Tensor] = None, + ) -> Tuple[Optional[Union[np.ndarray, torch.Tensor]], bool]: + """SpatialVLA tokenizer decoding""" + action_id = [] + for token_id_i in output_ids[0]: + if token_id_i.item() in action_mapper: + action_id.append(action_mapper[token_id_i.item()]) + + if len(action_id) == 0: + return np.zeros((action_horizon, action_dim)), False + + if state is not None: + predict_action = self._tokenizer.decode( + [action_id], + state=state[0, 0, :action_dim], + time_horizon=action_horizon, + action_dim=action_dim, + ) + else: + predict_action = self._tokenizer.decode( + [action_id], time_horizon=action_horizon, action_dim=action_dim + ) + + # Check whether decoding succeeded + decode_success = False + if isinstance(predict_action, np.ndarray): + decode_success = np.sum(predict_action) != 0 + elif isinstance(predict_action, torch.Tensor): + decode_success = predict_action.sum().item() != 0 + + return predict_action, decode_success + + def compute_accuracy( + self, + logits: torch.Tensor, + labels: torch.Tensor, + action_mapper: Dict, + action_token_id_set: Dict, + ) -> Dict[str, torch.Tensor]: + """Compute SpatialVLA tokenizer accuracy, same as fast""" + result = {} + + if len(action_token_id_set.get("action_token_list", [])) > 0: + shift_logits = logits[..., :-1, :].contiguous() + action_preds = shift_logits.argmax(dim=-1) + shift_labels = labels[..., 1:].contiguous() + action_mask = shift_labels > action_token_id_set["action_token_list"][0] + correct_preds = (action_preds == shift_labels) & action_mask + action_accuracy = correct_preds.sum().float() / action_mask.sum().float() + result["action_accuracy"] = action_accuracy + + return result + + def get_accuracy_keys(self) -> List[str]: + """Return SpatialVLA tokenizer accuracy keys""" + return ["action_accuracy"] + + @property + def vocab_size(self) -> int: + if self._tokenizer is None: + return 0 + return self._tokenizer.vocab_size + + @property + def uses_dof_mask_for_unnorm(self) -> bool: + """SpatialVLA tokenizer requires dof_mask""" + return True + + @property + def needs_action_crop(self) -> bool: + return False + + @property + def inference_ar_steps_for_dllm(self) -> int: + return self._tokenizer.max_waypoints + + def prepare_action_for_ar_encoding( + self, + ar_actionchunk: torch.Tensor, + dataset_names: List[str], + agent_pos: Optional[torch.Tensor] = None, + ) -> List[torch.Tensor]: + step = ar_actionchunk.shape[1] // self._tokenizer.max_waypoints + indices = np.arange(0, ar_actionchunk.shape[1], step)[ + : self._tokenizer.max_waypoints + ] + ar_action = ar_actionchunk[:, indices, :] + ar_action = self.action_normalizer.normalize_data(ar_action, dataset_names) + return ar_action + + def encode_to_tokens( + self, + actions: List, + obs_state: Optional[torch.Tensor] = None, + dof_mask: Optional[torch.Tensor] = None, + robot_type_ids: Optional[List[int]] = None, + is_train: bool = True, + ) -> List[List[str]]: + """ + SpatialVLA tokenizer encoding + + Args: + actions: List[Tensor/ndarray], each [T, D] (clipped) + """ + if self._tokenizer is None: + raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") + + if isinstance(actions, torch.Tensor): + # Convert torch tensors to numpy arrays + actions = actions.cpu().numpy() + tokenizer = self._tokenizer if is_train else self.get_val_tokenizer({}) + token_ids_group = tokenizer.batch_encode(actions) + all_action_tokens = [] + for i in range(len(token_ids_group)): + token_ids = np.array(token_ids_group[i]).reshape(-1) + action_token = [f"<|action_token_{i}|>" for i in token_ids] + all_action_tokens.append(action_token) + return all_action_tokens + + def modify_inputs_for_dllm( + self, + inputs: Dict[str, torch.Tensor], + processor, + sample_time: torch.Tensor, + dataset_names: List[str], + ): + """ + Prepare AR DLLM inputs by encoding sample_time into input_ids and labels + + Noise injection strategy: + - Split [0, 1] into N equal parts (N = ar_step_num) + - For each sample, compute the number of noisy steps k in [1, N] + k = clamp(ceil((1 - t) * N), 1, N) + -> smaller t means more noise; values closer to 1 mean less noise, with at least one noised step + - seq is the noise sequence formed by concatenating ar_step_num placeholder_seq blocks + Each step maps to len(placeholder_seq) consecutive tokens + - Noise injection overwrites matching tokens in input_ids / labels step by step from seq + """ + + input_ids = inputs["input_ids"] + labels = inputs["labels"] + prefix_length = inputs["prefix_length"] # int or tensor(1,) + bs, seqlen = input_ids.shape + + # Total number of AR steps + ar_step_num = self._tokenizer.max_waypoints # N + step_token_len = len(processor.placeholder_seq) # Number of tokens per step + ar_token_length = ar_step_num * step_token_len # Total AR token length + + device = input_ids.device + dtype = input_ids.dtype + + placeholder_ids = torch.tensor( + processor.placeholder_seq, device=device, dtype=dtype + ) # (step_token_len,) + noise_seq = placeholder_ids.repeat(ar_step_num) # (ar_token_length,) + + if not torch.is_tensor(sample_time): + sample_time = torch.tensor(sample_time, device=device, dtype=torch.float32) + else: + sample_time = sample_time.to(device=device, dtype=torch.float32) + + sample_time = sample_time.clamp(0.0, 1.0) + + noisy_steps_per_sample = ( + torch.ceil((1.0 - sample_time) * (ar_step_num + 1)).long() - 1 + ) + noisy_steps_per_sample = noisy_steps_per_sample.clamp(min=0, max=ar_step_num) + start = prefix_length - ar_token_length - 2 + end = prefix_length - 2 + + action_idx = 0 # action sample counter + for i in range(bs): + if not is_action_dataset_name(dataset_names[i]): + continue + k = noisy_steps_per_sample[action_idx].item() + action_idx += 1 + perm = torch.randperm(ar_step_num, device=device) + chosen_steps = perm[:k] # (k,) + + step_mask = torch.zeros(ar_step_num, dtype=torch.bool, device=device) + step_mask[chosen_steps] = True + + token_mask = step_mask.repeat_interleave( + step_token_len + ) # (ar_token_length,) + ignore_mask = ~token_mask + # Take a view of the current sample's AR span for mask-based replacement + cur_input_view = input_ids[i, start:end] + cur_label_view = labels[ + i, start + 1 : end + 1 + ] # shift placeholders one position to the right + + # Overwrite selected tokens with noise + cur_input_view[token_mask] = noise_seq[token_mask] + cur_label_view[ignore_mask] = -100 + + inputs["input_ids"] = input_ids + inputs["labels"] = labels + return inputs + + def update_positional_masks_for_dllm( + self, positional_masks, inputs, processor, visible_predict_ar_ratio=1 + ): + if self.dllm and self.input_placeholder_flag: + mask = self.update_placeholder_mask( + processor, + inputs["prefix_length"], + inputs["input_ids"], + ) + positional_masks.update(mask) + if ( + np.random.rand() < visible_predict_ar_ratio + ): # FIXME ar_visible is temporarily decided per batch; mixed settings are untested and need optimization. + positional_masks["ar_visible"] = True + if "ar_predict_token_positions" in positional_masks: + del positional_masks["ar_predict_token_positions"] + # positional_masks["ar_predict_token_positions"] = None + else: + positional_masks["ar_visible"] = False + return positional_masks + + def update_placeholder_mask(self, processor, prefix_length, input_ids): + ar_action_mask = torch.zeros_like(input_ids) + current_index = 1 + step_len = len(processor.placeholder_seq) + + for b, seq in enumerate(input_ids): + for i in range(self._tokenizer.max_waypoints): + start_idx = ( + prefix_length - (self._tokenizer.max_waypoints - i) * step_len - 2 + ) + end_idx = start_idx + step_len + ar_action_mask[b, start_idx:end_idx] = current_index + current_index += 1 + ar_action_mask[b, end_idx:prefix_length] = -1 # eos + + return {"ar_action_mask": ar_action_mask} + + + +# ============================================================ +# Factory function +# ============================================================ + +_TOKENIZER_REGISTRY: Dict[str, type] = { + "fast": FastTokenizerMixin, + "spatialvla": SpatialVLATokenizerMixin, +} + + +def get_action_tokenizer_mixin(tokenizer_type: str) -> ActionTokenizerMixin: + """ + Get the mixin instance for a tokenizer type + + Args: + tokenizer_type: tokenizer type, supports "fast", "spatialvla" + + Returns: + ActionTokenizerMixin instance + + Raises: + ValueError: Unsupported tokenizer type + """ + if tokenizer_type not in _TOKENIZER_REGISTRY: + raise ValueError( + f"Unsupported action tokenizer type: {tokenizer_type}. " + f"Supported types: {list(_TOKENIZER_REGISTRY.keys())}" + ) + return _TOKENIZER_REGISTRY[tokenizer_type]() diff --git a/wall_x/model/registry.py b/wall_x/model/registry.py new file mode 100644 index 00000000..ab7820ad --- /dev/null +++ b/wall_x/model/registry.py @@ -0,0 +1,64 @@ +"""Decorator-based model adapter registry.""" + +from __future__ import annotations + +from typing import Dict, Type + +# Global registry mapping model_type string → adapter class +_MODEL_REGISTRY: Dict[str, Type] = {} + + +def register_model(model_type: str): + """Decorator to register an adapter class for a given model_type. + + Can be stacked to register the same class for multiple model_types: + + @register_model("qwen2_5") + class QActAdapter: ... + """ + + def decorator(cls): + if model_type in _MODEL_REGISTRY: + existing = _MODEL_REGISTRY[model_type] + if existing is not cls: + raise ValueError( + f"model_type '{model_type}' already registered to " + f"{existing.__name__}, cannot re-register to {cls.__name__}" + ) + _MODEL_REGISTRY[model_type] = cls + return cls + + return decorator + + +def get_adapter(model_type: str, **kwargs): + """Instantiate the adapter registered for model_type. + + Args: + model_type: Registered model type string (e.g. "qwen2_5"). + **kwargs: Passed to the adapter constructor. + + Returns: + An adapter instance. + + Raises: + KeyError: If model_type is not registered. + """ + if model_type not in _MODEL_REGISTRY: + available = ", ".join(sorted(_MODEL_REGISTRY.keys())) or "(none)" + raise KeyError( + f"Unknown model_type '{model_type}'. " + f"Available: {available}. " + f"Did you forget to import the adapter module?" + ) + return _MODEL_REGISTRY[model_type](**kwargs) + + +def list_registered_models() -> list: + """Return sorted list of registered model_type strings.""" + return sorted(_MODEL_REGISTRY.keys()) + + +def clear_registry(): + """Clear all registrations. Intended for testing only.""" + _MODEL_REGISTRY.clear() diff --git a/wall_x/model/vla_mixin.py b/wall_x/model/vla_mixin.py deleted file mode 100644 index 1c98aa2b..00000000 --- a/wall_x/model/vla_mixin.py +++ /dev/null @@ -1,987 +0,0 @@ -import torch -import torch.nn as nn -import torch.utils.checkpoint as cp - -from torch.distributed.fsdp import MixedPrecision as MP -from torch.distributed.fsdp import FullyShardedDataParallel as FSDP - -from wall_x.fusions import ops - -from peft import LoraConfig, get_peft_model -from typing import Optional, Union, Dict -from packaging import version - -from transformers import GenerationMixin -from transformers.activations import ACT2FN -from transformers.modeling_utils import AttentionInterface - -from transformers.utils import logging, is_torch_xla_available - -from wall_x.model.action_head import ActionProcessor -from wall_x.model.model_utils import find_first_last_ones - -ALL_ATTENTION_FUNCTIONS: AttentionInterface = AttentionInterface() -logger = logging.get_logger(__name__) - - -X2ROBOT_ATTENTION_FUNCTIONS = [] -ATTENTION_TYPES_WITH_2D_MASK = [ - "sdpa", -] -ATTENTION_TYPES_WITH_FLASH_MASK = [] - - -class TokenTypeRouter(nn.Module): - def __init__(self, num_experts: int): - super().__init__() - self.num_experts = num_experts - - def forward(self, token_types: torch.Tensor) -> torch.Tensor: - """ - Assigns tokens to different experts based on `token_type`. - Args: - token_types (torch.Tensor): A tensor of shape (batch_size, seq_length) representing the type of each token. - - Returns: - experts_indices (torch.Tensor): A tensor of shape (batch_size, seq_length) representing the expert index assigned to each token. - """ - experts_indices = token_types % self.num_experts - return experts_indices - - -class BlockSparseMLP(nn.Module): - def __init__(self, config, use_selective_recompute: bool = False): - super().__init__() - self.hidden_size = config["hidden_size"] - self.intermediate_size = config["intermediate_size"] - self.hidden_act = config["hidden_act"] - - self.use_selective_recompute = use_selective_recompute - - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) - - self.act_fn = ACT2FN[self.hidden_act] - - def _full_mlp(self, hidden_state): - gate_out = self.gate_proj(hidden_state) - up_out = self.up_proj(hidden_state) - act_out = self.act_fn(gate_out) * up_out - return self.down_proj(act_out) - - def forward(self, hidden_state): - if self.use_selective_recompute: - # Perform checkpoint recalculation for the entire expert MLP. - return cp.checkpoint( - self._full_mlp, - hidden_state, - use_reentrant=False, - ) - else: - return self._full_mlp(hidden_state) - - -class SparseMoeBlock(nn.Module): - def __init__(self, config, num_experts: int, use_selective_recompute: bool = False): - super().__init__() - self.num_experts = num_experts - self.use_selective_recompute = use_selective_recompute - - # Pass the `use_selective_recompute` parameter to each expert. - self.experts = nn.ModuleList( - [ - BlockSparseMLP( - config.experts[i], use_selective_recompute=use_selective_recompute - ) - for i in range(num_experts) - ] - ) - - if not hasattr(config, "dim_inputs") or not config.dim_inputs: - raise ValueError("Configuration must contain a valid dim_inputs") - - self.dim_inputs = config.dim_inputs - self.permuted = config.mot_opt - - def forward( - self, - hidden_states: torch.Tensor, - experts_indices: torch.Tensor, - start_indices: torch.Tensor, - end_indices: torch.Tensor, - ) -> torch.Tensor: - - if self.permuted: - permuted_inputs = hidden_states - else: - batch_size, seq_length, hidden_dim = hidden_states.shape - - flat_hidden = hidden_states.reshape(-1, hidden_dim) - experts_indices = experts_indices.reshape(-1) - probs = torch.ones_like(experts_indices, dtype=torch.float32).reshape(-1, 1) - permuted_inputs, row_id_map = ops.permute(flat_hidden, experts_indices) - - # buffer - final_output = torch.zeros_like(permuted_inputs) - - # Expert forward contain selective recompute - for expert_idx, expert in enumerate(self.experts): - start, end = start_indices[expert_idx], end_indices[expert_idx] - if start == end: - continue - - dim_input = self.dim_inputs[expert_idx] - expert_input = permuted_inputs[start:end, :dim_input] - - partial_output = expert(expert_input) - final_output[start:end, :dim_input] = partial_output[:, :dim_input] - - if self.permuted: - return final_output - else: - final_output = ops.unpermute(final_output, row_id_map, probs) - return final_output.reshape(batch_size, seq_length, hidden_dim) - - -class ActionModelMixMin: - # config: Qwen2_5_VLConfig - action_preprocessor: ActionProcessor - router: TokenTypeRouter - moe: SparseMoeBlock - - def __init__(self, config, action_preprocessor, router, moe): - self.config = config - self.action_preprocessor = action_preprocessor - self.router = router - self.moe = moe - self._mot_opt_warned = False - - def set_normalizer(self, normalizer_action, normalizer_propri): - if hasattr(self, "action_preprocessor"): - self.action_preprocessor.set_normalizer( - normalizer_action, normalizer_propri - ) - else: - logger.warning( - "ActionModelMixMin.set_normalizer is called but action_preprocessor is not set" - ) - - def _apply_mlp_moe(self, hidden_states, token_types, start_indices, end_indices): - if self.config.mlp_moe: - hidden_states = self.moe( - hidden_states, token_types, start_indices, end_indices - ) - else: - hidden_states = self.mlp(hidden_states) - return hidden_states - - def _apply_norm_moe( - self, - hidden_states, - token_types, - adarms_conds, - norms, # list of norm layers (expert-wise) - norm, # shared norm if not norm_moe - start_indices=None, - end_indices=None, - use_selective_recompute=False, - ): - """ - MoE-aware LayerNorm with optional selective activation recomputation. - - Only activation math is recomputed. No GEMM is recomputed. - Safe for FSDP (use_reentrant=False). - """ - - gate = None - gate_mask = None - - # ------------------------- - # Case 1: norm_moe=True (expert-wise norm) - # ------------------------- - if self.config.norm_moe: - - # --------------------------------------------------------- - # Case 1A: mot_opt=True (segments assigned by start/end) - # --------------------------------------------------------- - if self.config.mot_opt: - new_hidden_states = torch.zeros_like(hidden_states) - - for expert_idx, expert_norm in enumerate(norms): - start = start_indices[expert_idx] - end = end_indices[expert_idx] - if start == end: - continue - - dim_input = self.config.dim_inputs[expert_idx] - selected = hidden_states[start:end] # [K, D] - - # ====== reshape if adarms on flow expert ====== - if self.config.use_adarms and expert_idx == 1: - selected = selected.view( - -1, - self.config.action_horizon_flow, - selected.shape[-1], - ) - input_slice = selected[:, :, :dim_input] - cond = adarms_conds[expert_idx] - else: - input_slice = selected[:, :dim_input] - cond = adarms_conds[expert_idx] - - if use_selective_recompute: - - def norm_chunk(t_x, t_cond, expert_norm=expert_norm): - if t_cond is None or ( - isinstance(t_cond, torch.Tensor) and t_cond.numel() == 0 - ): - out, _ = expert_norm(t_x) - else: - out, _ = expert_norm(t_x, t_cond) - return out - - cond_for_cp = ( - cond - if cond is not None - else torch.empty(0, device=input_slice.device) - ) - processed = cp.checkpoint( - norm_chunk, - input_slice, - cond_for_cp, - use_reentrant=False, - ) - else: - processed, gate = expert_norm(input_slice, cond) - - # reshape back if needed - if self.config.use_adarms and expert_idx == 1: - processed = processed.view(-1, dim_input) - - new_hidden_states[start:end, :dim_input] = processed.to( - hidden_states.dtype - ) - - hidden_states = new_hidden_states - - # --------------------------------------------------------- - # Case 1B: mot_opt=False (token-level mask) - # --------------------------------------------------------- - else: - - new_hidden_states = torch.zeros_like(hidden_states) - B, S, D = hidden_states.shape - - for expert_idx, expert_norm in enumerate(norms): - mask = token_types == expert_idx - if mask.sum() == 0: - continue - - dim_input = self.config.dim_inputs[expert_idx] - selected = hidden_states[mask] # [K, D] - - if self.config.use_adarms and expert_idx == 1: - gate_mask = mask - selected = selected.view( - -1, - self.config.action_horizon_flow, - selected.shape[-1], - ) - input_slice = selected[:, :, :dim_input] - cond = adarms_conds[expert_idx] - else: - input_slice = selected[:, :dim_input] - cond = adarms_conds[expert_idx] - - if use_selective_recompute: - - def norm_chunk(t_x, t_cond, expert_norm=expert_norm): - if t_cond is None or ( - isinstance(t_cond, torch.Tensor) and t_cond.numel() == 0 - ): - out, _ = expert_norm(t_x) - else: - out, _ = expert_norm(t_x, t_cond) - return out - - cond_for_cp = ( - cond - if cond is not None - else torch.empty(0, device=input_slice.device) - ) - - processed = cp.checkpoint( - norm_chunk, - input_slice, - cond_for_cp, - use_reentrant=False, - ) - else: - processed, gate = expert_norm(input_slice, cond) - - if self.config.use_adarms and expert_idx == 1: - processed = processed.view(-1, dim_input) - - # scatter back - b_id, s_id = torch.where(mask) - new_hidden_states[b_id, s_id, :dim_input] = processed.to( - hidden_states.dtype - ) - - hidden_states = new_hidden_states - - # ------------------------- - # Case 2: norm_moe=False (single LN) - # ------------------------- - else: - - def norm_chunk_shared(t_x, dummy, norm_module=norm): - out, _ = norm_module(t_x) - return out - - if use_selective_recompute: - dummy = torch.empty(0, device=hidden_states.device) - hidden_states = cp.checkpoint( - norm_chunk_shared, - hidden_states, - dummy, - use_reentrant=False, - ) - else: - hidden_states, gate = norm(hidden_states) - - return hidden_states, gate, gate_mask - - def _gated_residual(self, x, y, gate, start_indices=None, end_indices=None): - """ - Applies gated residual connection with optional gate parameter. - - Args: - x: Input tensor (residual) - y: Output tensor to be added - gate: Optional gate tensor to modulate the addition - - Returns: - x + y if gate is None, otherwise x + y * gate - """ - if x is None and y is None: - return None - if x is None or y is None: - return x if x is not None else y - if gate is None: - return x + y - - new_y = y.clone() - selected_y = y[start_indices[1] : end_indices[1]] - selected_y = selected_y.view( - -1, self.config.action_horizon_flow, selected_y.shape[-1] - )[:, :, : self.config.dim_inputs[1]] - selected_y = selected_y.to(torch.float32) * gate - new_y[start_indices[1] : end_indices[1], : self.config.dim_inputs[1]] = ( - selected_y.view(-1, self.config.dim_inputs[1]).to(new_y.dtype) - ) - - return x + new_y - - def scatter_proprioception_embeddings( - self, input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask - ): - if ( - proprioception is not None - and not self.config.use_state_string_representation - ): - proprioception = proprioception.to(inputs_embeds.device).to( - inputs_embeds.dtype - ) - agent_pos_mask = agent_pos_mask.to(inputs_embeds.device).to( - inputs_embeds.dtype - ) - proprioception = self.action_preprocessor.proprioception_proj( - proprioception, - dataset_names, - agent_pos_mask, - use_history=proprioception.shape[1] > 1, - ) - mask = input_ids == self.action_token_id_set["propri_token_id"] - mask_unsqueezed = mask.unsqueeze(-1) - mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) - proprioception_mask = mask_expanded.to(inputs_embeds.device) - - proprioception = proprioception.to( - inputs_embeds.device, inputs_embeds.dtype - ) - inputs_embeds = inputs_embeds.masked_scatter( - proprioception_mask, proprioception - ) - - return inputs_embeds - - def scatter_flow_action_embeddings( - self, input_ids, inputs_embeds, action_chunk, dataset_names, dof_mask - ): - if not self.config.use_flow_action_expert: - return inputs_embeds, None, None - adarms_cond, flow = None, None - if action_chunk is not None: - action_chunk = action_chunk.to(inputs_embeds.device) - dof_mask = dof_mask.to(inputs_embeds.device) - noisy_action_emb, flow, adarms_cond = self.action_preprocessor( - action_chunk, dataset_names, dof_mask - ) - mask = input_ids == self.action_token_id_set["action_token_id"] - mask_unsqueezed = mask.unsqueeze(-1) - mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) - action_mask = mask_expanded.to(inputs_embeds.device) - - noisy_action_emb = noisy_action_emb.to( - inputs_embeds.device, inputs_embeds.dtype - ) - inputs_embeds = inputs_embeds.masked_scatter(action_mask, noisy_action_emb) - - return inputs_embeds, flow, adarms_cond - - @staticmethod - def _update_position_ids( - position_ids, - moe_token_types, - positional_masks, - ): - if ( - positional_masks is None - or "ar_predict_token_positions" not in positional_masks - ): - return position_ids - - new_position_ids = position_ids.clone() - ar_predict_token_positions = positional_masks["ar_predict_token_positions"] - flow_mask = moe_token_types == 1 - - start_ar_pos, end_ar_pos = find_first_last_ones(ar_predict_token_positions) - start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) - - for bs_i in range(position_ids.shape[1]): - if start_ar_pos[bs_i] != -1 and end_ar_pos[bs_i] != -1: - start_ar_ids = new_position_ids[:, bs_i, start_ar_pos[bs_i]] - start_flow_ids = new_position_ids[:, bs_i, start_flow_pos[bs_i]] - diff = start_flow_ids - start_ar_ids - new_position_ids[:, bs_i, start_flow_pos[bs_i] :] = position_ids[ - :, bs_i, start_flow_pos[bs_i] : - ] - diff.unsqueeze(-1) - - return new_position_ids - - def _update_joint_attention_mask_2d( - self, - attention_mask, - moe_token_types, - positional_masks, - ): - if attention_mask.dim() == 3: # bs, seq_len, seq_len - return attention_mask - - bs, seq_len = moe_token_types.shape[0], moe_token_types.shape[1] - # Create a lower triangular matrix as a causal mask. - causal_mask = torch.tril( - torch.ones( - (seq_len, seq_len), dtype=torch.bfloat16, device=moe_token_types.device - ) - ) - # Extended to the batch dimension. - attention_mask = causal_mask.unsqueeze(0).expand(bs, -1, -1) - - if positional_masks is not None and "padding_positions" in positional_masks: - padding_positions = positional_masks["padding_positions"] - # The padding is set to zero. - attention_mask = torch.where( - padding_positions[:, None, :], - torch.zeros_like(attention_mask), - attention_mask, - ) - # The padding is set to zero. - attention_mask = torch.where( - padding_positions[:, :, None], - torch.zeros_like(attention_mask), - attention_mask, - ) - - # Set all values ​​in the moe1 section to 1, and disable the fast section. - moe1_mask = (moe_token_types[:, :, None]) & (moe_token_types[:, None, :]) - - if ( - not self.config.causal_action_attention_mask - ): # If a causal action attention mask is not used, then all elements in the moe1 section are set to 1. - attention_mask = torch.where( - moe1_mask, torch.ones_like(attention_mask), attention_mask - ) - - if ( - positional_masks is not None - and "ar_predict_token_positions" in positional_masks - ): - ar_predict_token_positions = positional_masks["ar_predict_token_positions"] - moe1_mask = (moe_token_types[:, :, None]) & ( - ar_predict_token_positions[:, None, :] - ) - attention_mask = torch.where( - moe1_mask, torch.zeros_like(attention_mask), attention_mask - ) - - if ( - positional_masks is not None - and "valid_flow_action_positions" in positional_masks - ): - # true in moe_token_types but false in valid_flow_action_positions - nonvalid_flow_action_positions = ( - moe_token_types & ~positional_masks["valid_flow_action_positions"] - ) - attention_mask = torch.where( - nonvalid_flow_action_positions[:, None, :], - torch.zeros_like(attention_mask), - attention_mask, - ) - attention_mask = torch.where( - nonvalid_flow_action_positions[:, :, None], - torch.zeros_like(attention_mask), - attention_mask, - ) - - return attention_mask - - def _update_joint_attention_flash_mask( - self, - attention_mask, - moe_token_types, - positional_masks, - debug=False, - ): - device = moe_token_types.device - B, S = moe_token_types.shape - i32 = torch.int32 - - # ---- Return vector initialization ---- - LTS = torch.ones((B, S), device=device, dtype=i32) * S - UTE = ( - torch.arange(S, device=device, dtype=i32).unsqueeze(0).expand(B, S).clone() - ) - - # Handling padding positions - if positional_masks is not None and "padding_positions" in positional_masks: - padding_positions = positional_masks["padding_positions"] - LTS[padding_positions] = 0 - UTE[padding_positions] = S - - # Handling ar predict tokens - if ( - positional_masks is not None - and "ar_predict_token_positions" in positional_masks - ): - start_ar_pos, end_ar_pos = find_first_last_ones( - positional_masks["ar_predict_token_positions"] - ) - for bs_i in range(B): - if end_ar_pos[bs_i] != -1: - LTS[bs_i, positional_masks["ar_predict_token_positions"][bs_i]] = ( - end_ar_pos[bs_i].to(i32) + 1 - ) - - # Handling flow action bidirectional mask - flow_mask = moe_token_types == 1 - if not self.config.causal_action_attention_mask: - start_flow_pos, end_flow_pos = find_first_last_ones(flow_mask) - for bs_i in range(B): - if start_flow_pos[bs_i] != -1: - UTE[bs_i, flow_mask[bs_i]] = start_flow_pos[bs_i].to(i32) - - # Handling validate flow - if ( - positional_masks is not None - and "valid_flow_action_positions" in positional_masks - ): - flow_mask = moe_token_types == 1 - nonvalid_flow_action_positions = ( - flow_mask & ~positional_masks["valid_flow_action_positions"] - ) - if nonvalid_flow_action_positions.any(): - LTS[nonvalid_flow_action_positions] = 0 - UTE[nonvalid_flow_action_positions] = S - - LTS = LTS.unsqueeze(-1) - UTE = UTE.unsqueeze(-1) - - startend_row_indices = torch.cat([LTS, UTE], dim=-1) - # startend_row_indices = LTS - - # add num_heads dimension - startend_row_indices = startend_row_indices.unsqueeze(1) - - return startend_row_indices - - -class ActionGenerationMixin(GenerationMixin): - action_preprocessor: ActionProcessor - - def to_bfloat16_for_selected_params(self, fsdp_plugin=None, accelerator=None): - """ - Keep some model parameters as float32, and convert others to bfloat16. - - If `fsdp_plugin` exists, use FSDP v1's `mixed_precision` wrapper. - - Otherwise, directly modify the parameter dtype. - """ - - def _assign_child(root_module, dotted_name: str, new_child): - parts = dotted_name.split(".") - parent = root_module - for p in parts[:-1]: - parent = getattr(parent, p) - setattr(parent, parts[-1], new_child) - - if fsdp_plugin: - fsdp_version = getattr(fsdp_plugin, "fsdp_version", None) - if fsdp_version != 1: - raise RuntimeError("Only FSDP v1 is supported (fsdp_version=1).") - - device = getattr( - accelerator, "device", torch.device("cuda", torch.cuda.current_device()) - ) - if isinstance(device, torch.device) and device.type == "cuda": - if device.index is not None: - torch.cuda.set_device(device.index) - device_id = device.index - - # move model to device - self = self.to(device) - - # Define the mixed-precision strategy. - bf16_policy = MP( - param_dtype=torch.bfloat16, - reduce_dtype=torch.float32, - buffer_dtype=torch.bfloat16, - cast_forward_inputs=False, - cast_root_forward_inputs=False, - ) - - fp32_policy = MP( - param_dtype=torch.float32, - reduce_dtype=torch.float32, - buffer_dtype=torch.float32, - cast_forward_inputs=False, - cast_root_forward_inputs=False, - ) - - # Step 1️⃣: Identify the top-level ActionProcessor module and wrap it separately with FSDP (FP32). - for name, module in list(self.named_modules()): - if isinstance(module, nn.Module) and any( - k in name.lower() for k in ["action_preprocessor"] - ): - if any(True for _ in module.children()): - continue - if getattr(module, "_fsdp_wrapped", False): - continue - - print(f"[FSDP v1] wrapping module in FP32: {name}") - wrapped = FSDP( - module, - mixed_precision=fp32_policy, - sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, - backward_prefetch="BACKWARD_PRE", - device_id=device_id, - use_orig_params=True, - ) - _assign_child(self, name, wrapped) - setattr(wrapped, "_fsdp_wrapped", True) - - # Step 2️⃣: The outermost layer uses unified FSDP (BF16 strategy). - print("[FSDP v1] wrapping root model with bf16 mixed precision...") - self = FSDP( - self, - mixed_precision=bf16_policy, - sharding_strategy=torch.distributed.fsdp.ShardingStrategy.SHARD_GRAD_OP, - backward_prefetch="BACKWARD_PRE", - device_id=device_id, - use_orig_params=True, - ) - - return self - - # ----------------- Non-FSDP scenarios ----------------- - else: - print("[INFO] Running manual dtype conversion (no FSDP).") - self.to(dtype=torch.float32) - - params_to_keep_float32 = [] - for name, _ in self.named_parameters(): - if any( - k in name - for k in [ - "input_layernorm", - "post_attention_layernorm", - "model.norm", - "action_preprocessor", - ] - ): - params_to_keep_float32.append(name) - - for name, param in self.named_parameters(): - if name not in params_to_keep_float32: - param.data = param.data.to(torch.bfloat16) - - return self - - def define_action_token_id(self): - action_token_list = [] - if self.action_tokenizer_type: - for i in range(self.action_tokenizer.vocab_size): - action_token_id = self.processor.tokenizer.convert_tokens_to_ids( - f"<|action_token_{i}|>" - ) - action_token_list.append(action_token_id) - - action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>") - propri_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|propri|>") - self.action_token_id_set = { - "action_token_list": action_token_list, - "propri_token_id": propri_token_id, - "action_token_id": action_token_id, - } - - def add_lora( - self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1 - ): - """Add LoRA adapter""" - config = LoraConfig( - r=r, - lora_alpha=lora_alpha, - target_modules=target_modules, - lora_dropout=lora_dropout, - bias="none", - task_type="CAUSAL_LM", - ) - self.model = get_peft_model(self.model, config) - # Print trainable parameter information. - self.model.print_trainable_parameters() - - def compute_loss( - self, - hidden_states, - logits, - input_ids=None, - dataset_names=None, - labels=None, - action_chunk=None, - dof_mask=None, - flow=None, - flow_loss_mask=None, - **kwargs, - ): - if input_ids is not None: - batch_size, seq_length = input_ids.shape - - loss = 0 - cross_entropy_loss, flow_loss = None, None - - # if dataset_names is not None: - # unique_datasets_name = list(set(dataset_names)) - # channel_loss_dict = { - # dataset_name: torch.tensor(0.0, device=logits.device) - # for dataset_name in _ACTION_DATASET_NAMES + _MULTIMODAL_DATASET_NAMES - # } - # channel_loss_count_dict = { - # dataset_name: torch.tensor(0, device=logits.device) - # for dataset_name in _ACTION_DATASET_NAMES + _MULTIMODAL_DATASET_NAMES - # } - # else: - unique_datasets_name, channel_loss_dict, channel_loss_count_dict = ( - None, - None, - None, - ) - - if labels is not None: - action_accuracy = 0 - - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - non_ignored_mask = shift_labels != -100 - _cross_entropy_loss = self.loss_fct(shift_logits, shift_labels) - cross_entropy_loss = ( - _cross_entropy_loss[non_ignored_mask].mean() - if non_ignored_mask.any() - else torch.tensor(0.0, device=shift_logits.device) - ) - - # compute channel loss - _cross_entropy_loss = _cross_entropy_loss.view(batch_size, seq_length - 1) - non_ignored_mask = non_ignored_mask.view(batch_size, seq_length - 1) - for dataset_name_i in unique_datasets_name: - dataset_mask = torch.tensor( - [name == dataset_name_i for name in dataset_names], - device=logits.device, - ) - combined_mask = dataset_mask.unsqueeze(1) & non_ignored_mask - channel_loss_dict[dataset_name_i] = ( - _cross_entropy_loss[combined_mask].sum() - if combined_mask.any() - else torch.tensor(0.0, device=shift_logits.device) - ) - channel_loss_count_dict[dataset_name_i] += combined_mask.sum() - - if not torch.isnan(cross_entropy_loss): - loss += cross_entropy_loss - else: - with torch.no_grad(): - cross_entropy_loss.detach() - - # compute action token accuracy - if len(self.action_token_id_set["action_token_list"]) > 0: - shift_logits = logits[..., :-1, :].contiguous() - action_preds = shift_logits.argmax(dim=-1) - shift_labels = labels[..., 1:].contiguous() - action_mask = ( - shift_labels > self.action_token_id_set["action_token_list"][0] - ) - correct_preds = (action_preds == shift_labels) & action_mask - action_accuracy = ( - correct_preds.sum().float() / action_mask.sum().float() - ) - channel_loss_dict["action_accuracy"] = action_accuracy - - if action_chunk is not None: - action_mask = input_ids == self.action_token_id_set["action_token_id"] - if action_mask.any(): - action_hidden_states = hidden_states[action_mask].to(torch.float32) - flow = flow.reshape(-1, flow.shape[-1]) - _flow_loss = self.action_preprocessor.flow_loss( - action_hidden_states, flow, action_chunk, dof_mask, flow_loss_mask - ) - if isinstance(_flow_loss, torch.Tensor): - flow_loss = _flow_loss.mean() - loss += flow_loss * self.config.flow_loss_weight - _flow_loss = _flow_loss.view( - dof_mask.shape[0], dof_mask.shape[1], dof_mask.shape[2] - ) - - return ( - loss, - cross_entropy_loss, - flow_loss, - channel_loss_dict, - channel_loss_count_dict, - ) - - -class AttentionsSelectorMixin: - - @classmethod - def _autoset_attn_implementation( - cls, - config, - use_flash_attention_2: bool = False, - torch_dtype: Optional[torch.dtype] = None, - device_map: Optional[Union[str, Dict[str, int]]] = None, - check_device_map: bool = True, - ): - """ - Automatically checks and dispatches to a default attention implementation. In order of priority: - 1. An implementation specified in `config._attn_implementation` (due for example to the argument attn_implementation="sdpa" in from_pretrained). - 2. DEPRECATED: if use_flash_attention_2 is set to `True` and `flash_attn` is available, flash attention. (`LlamaFlashAttention` for example) - 3. SDPA implementation, if available and supported by the model type. (`LlamaSdpaAttention` for example) - 4. The default model's implementation otherwise (`LlamaAttention` for example) . - """ - # Here we use config._attn_implementation_internal to check whether the attention implementation was explicitly set by the user. - # The property `PretrainedConfig._attn_implementation` is never `None`, for backward compatibility (always fall back on "eager"). - # The `hasattr` here is used as some Transformers tests for some reason do not call PretrainedConfig __init__ (e.g. test_no_super_init_config_and_model) - requested_attn_implementation = None - if ( - hasattr(config, "_attn_implementation_internal") - and config._attn_implementation_internal is not None - ): - if ( - config._attn_implementation != "flash_attention_2" - and use_flash_attention_2 - ): - raise ValueError( - f'Both attn_implementation="{config._attn_implementation}" and `use_flash_attention_2=True` were used when loading the model, which are not compatible.' - ' We recommend to just use `attn_implementation="flash_attention_2"` when loading the model.' - ) - - if ( - not isinstance(config._attn_implementation, dict) - and config._attn_implementation - not in ["eager"] - + ALL_ATTENTION_FUNCTIONS.valid_keys() - + X2ROBOT_ATTENTION_FUNCTIONS - ): - message = f'Specified `attn_implementation="{config._attn_implementation}"` is not supported. The only possible arguments are `attn_implementation="eager"` (manual attention implementation)' - if cls._supports_flash_attn_2: - message += ', `"attn_implementation=flash_attention_2"` (implementation using flash attention 2)' - if cls._supports_sdpa: - message += ', `"attn_implementation=sdpa"` (implementation using torch.nn.functional.scaled_dot_product_attention)' - if cls._supports_flex_attn: - message += ', `"attn_implementation=flex_attention"` (implementation using torch\'s flex_attention)' - raise ValueError(message + ".") - - # If a config is passed with a preset attn_implementation, we skip the automatic dispatch and use the user-provided config, with hard checks that the requested attention implementation is available. - requested_attn_implementation = config._attn_implementation_internal - - if use_flash_attention_2: - logger.warning_once( - 'The model was loaded with use_flash_attention_2=True, which is deprecated and may be removed in a future release. Please use `attn_implementation="flash_attention_2"` instead.' - ) - config._attn_implementation = "flash_attention_2" - - if config._attn_implementation == "flash_attention_2": - cls._check_and_enable_flash_attn_2( - config, - torch_dtype=torch_dtype, - device_map=device_map, - hard_check_only=False, - check_device_map=check_device_map, - ) - elif requested_attn_implementation == "flex_attention": - config = cls._check_and_enable_flex_attn(config, hard_check_only=True) - elif ( - requested_attn_implementation in [None, "sdpa"] - and not is_torch_xla_available() - ): - # use_flash_attention_2 takes priority over SDPA, hence SDPA treated in this elif. - config = cls._check_and_enable_sdpa( - config, - hard_check_only=( - False if requested_attn_implementation is None else True - ), - ) - - if ( - torch.version.hip is not None - and config._attn_implementation == "sdpa" - and torch.cuda.device_count() > 1 - and version.parse(torch.__version__) < version.parse("2.4.1") - ): - logger.warning_once( - "Using the `SDPA` attention implementation on multi-gpu setup with ROCM may lead to performance issues due to the FA backend. Disabling it to use alternative backends." - ) - torch.backends.cuda.enable_flash_sdp(False) - elif requested_attn_implementation in ALL_ATTENTION_FUNCTIONS.valid_keys(): - config._attn_implementation = requested_attn_implementation - elif isinstance(requested_attn_implementation, dict): - config._attn_implementation = None - elif config._attn_implementation in X2ROBOT_ATTENTION_FUNCTIONS: - pass - else: - config._attn_implementation = "eager" - - config._attn_implementation_autoset = True - return config - - def _check_and_adjust_attn_implementation( - self, attn_implementation: Optional[str], is_init_check: bool = False - ) -> str: - assert ( - attn_implementation - in ["eager", "flash_attention_2", "sdpa"] + X2ROBOT_ATTENTION_FUNCTIONS - ) - return attn_implementation diff --git a/wall_x/serving/README.md b/wall_x/serving/README.md deleted file mode 100644 index febfe3aa..00000000 --- a/wall_x/serving/README.md +++ /dev/null @@ -1,263 +0,0 @@ -# Wall-X Model Serving - -This directory contains scripts for serving Wall-X models via a websocket server, allowing remote clients to connect and get action predictions from observations. - -## Overview - -The serving infrastructure consists of three main components: - -1. **WebsocketPolicyServer** (`wall_x/serving/websocket_policy_server.py`): Generic websocket server that can serve any policy implementing the `BasePolicy` interface -2. **WallXPolicy** (`wall_x/serving/policy/wall_x_policy.py`): Policy wrapper that adapts the Wall-X model to the `BasePolicy` interface -3. **launch_serving.py**: Main script for starting the server with various configurations - -## Quick Start - -### Basic Usage - -Serve a model with default LIBERO configuration: - -```bash -cd /x2robot_v2/vincent/workspace/opensource -python -m wall_x.serving.launch_serving \ - --env libero \ - --model-config.model-path /path/to/libero_model_stuff \ - --model-config.action-tokenizer-path /path/to/fast/ \ - --model-config.train-config-path /path/to/config.yml -``` - -### Specify Environment - -Serve with a specific environment preset: - -```bash -# LIBERO (single arm, 7 DOF) -python -m wall_x.serving.launch_serving --env libero - -# ALOHA (dual arm, 14 DOF) -python -m wall_x.serving.launch_serving --env aloha -``` - -### Custom Configuration - -Serve with custom model paths and settings: - -```bash -python -m wall_x.serving.launch_serving \ - --model-config.model-path /path/to/model \ - --model-config.action-tokenizer-path /path/to/tokenizer \ - --model-config.train-config-path /path/to/train_config.yml \ - --model-config.action-dim 7 \ - --model-config.state-dim 8 \ - --model-config.pred-horizon 32 \ - --model-config.camera-key front_view left_wrist_view \ - --port 8000 -``` - -## Command Line Arguments - -### Basic Arguments - -- `--env {libero,aloha}`: Environment mode (default: libero) -- `--port PORT`: Port to serve on (default: 8000) -- `--host HOST`: Host to bind to (default: 0.0.0.0) -- `--default-prompt TEXT`: Default text prompt if not provided in observation -- `--debug`: Enable debug logging - -### Model Configuration - -All model configuration arguments use the `--model-config.` prefix: - -- `--model-config.model-path PATH`: Path to pretrained model checkpoint (required) -- `--model-config.action-tokenizer-path PATH`: Path to action tokenizer (required) -- `--model-config.train-config-path PATH`: Path to train config YAML file (required) -- `--model-config.action-dim INT`: Action space dimension (default: 7) -- `--model-config.state-dim INT`: Robot state dimension (default: 8) -- `--model-config.pred-horizon INT`: Prediction horizon (default: 32) -- `--model-config.device {cuda,cpu}`: Device to run on (default: cuda) -- `--model-config.dtype {bfloat16,float16,float32}`: Model dtype (default: bfloat16) -- `--model-config.predict-mode {fast,diffusion}`: Prediction mode (default: fast) -- `--model-config.camera-key KEY1 KEY2 ...`: Camera keys for observation images - -### Camera Keys - -The `camera-key` parameter specifies which camera views are expected in the observation dictionary. This is **critical** for proper operation: - -- Keys must match between server configuration and client observations -- Order matters: keys are processed in the order specified -- Common keys: `front_view`, `left_wrist_view`, `right_wrist_view`, `face_view` - -Example: -```bash ---model-config.camera-key front_view left_wrist_view -``` - -Client must send observations with matching keys: -```python -obs = { - "front_view": image1, # Must match camera-key[0] - "left_wrist_view": image2, # Must match camera-key[1] - "prompt": "task description", - "state": robot_state, -} -``` - -## Default Configurations - -### LIBERO (Single Arm) - -```python -ModelConfig( - model_path="/path/to/model", - action_tokenizer_path="/path/to/action_tokenizer", - train_config_path="/path/to/train_config", - state_dim=8, - action_dim=7, - pred_horizon=32, - device="cuda", - dtype="bfloat16", - predict_mode="fast", - camera_key=["front_view", "left_wrist_view"], -) -``` - -### ALOHA (Dual Arm) - -```python -ModelConfig( - model_path="/path/to/model", - action_tokenizer_path="/path/to/action_tokenizer", - train_config_path="/path/to/train_config", - state_dim=14, - action_dim=14, - pred_horizon=32, - device="cuda", - dtype="bfloat16", - predict_mode="fast", - camera_key=["face_view", "left_wrist_view", "right_wrist_view"], -) -``` - -## Server Protocol - -### Connection Flow - -1. Client connects to `ws://host:port` -2. Server sends metadata JSON with policy information -3. Client sends observation (msgpack-encoded) -4. Server responds with action prediction (msgpack-encoded) -5. Repeat steps 3-4 for each inference - -### Observation Format - -Observations must be a dictionary with camera keys matching server configuration: - -```python -obs = { - # Image observations - keys must match server's camera_key configuration - "front_view": np.ndarray, # (H, W, 3) uint8 or float - "left_wrist_view": np.ndarray, # (H, W, 3) uint8 or float - - # Required fields - "prompt": str, # Task description - "dataset_names": List[str], # Dataset/robot name, e.g., ["physical-intelligence/libero"] - "state": np.ndarray, # Robot proprioception state (state_dim,) -} -``` - -**Important**: The image keys (`front_view`, `left_wrist_view`, etc.) must exactly match the `camera_key` parameter configured on the server. - -### Action Response Format - -Actions are returned as a dictionary: - -```python -{ - "action": np.ndarray, # Predicted action [pred_horizon, action_dim] - "server_timing": { - "infer_ms": float, # Inference time in milliseconds - "prev_total_ms": float, # Total time for previous request - } -} -``` - -### Server Metadata - -When connecting, the server sends metadata: - -```python -{ - "action_dim": int, # Action space dimension - "pred_horizon": int, # Number of future actions predicted - "device": str, # Device model runs on - "predict_mode": str, # Prediction mode (fast/diffusion) - "env": str, # Environment name -} -``` - -### Health Check - -HTTP health check endpoint available at: -``` -http://host:port/healthz -``` - -Returns `200 OK` if the server is running. - -## Client Example - -### Synchronous Python Client - -For synchronous usage, see `wall_x/serving/client.py`: - -```python -from wall_x.serving.client import WallXClient - -# Create and connect -client = WallXClient(uri="ws://localhost:8000") -client.connect_sync() - -# Prepare observation -obs = { - "front_view": image1, - "left_wrist_view": image2, - "prompt": "task description", - "state": robot_state, - "dataset_names": ["physical-intelligence/libero"], -} - -# Get prediction -response = client.predict_sync(obs) -action = response["action"] - -# Close connection -client.close_sync() -``` - -## Architecture - -### WebsocketPolicyServer - -Generic websocket server that: -- Handles websocket connections with msgpack serialization -- Tracks inference timing and performance metrics -- Provides health check endpoint -- Handles errors gracefully with proper logging -- Supports concurrent client connections - -### WallXPolicy - -Policy wrapper that: -- Loads and manages the Wall-X model from pretrained checkpoint -- Processes multi-camera observations -- Handles image preprocessing (smart resize, normalization) -- Manages device placement and dtype conversion -- Provides policy metadata to clients -- Supports both fast tokenizer and diffusion prediction modes - -### Image Processing Pipeline - -1. **Camera Key Matching**: Extracts images from observation dict using configured camera keys -2. **Format Conversion**: Converts numpy arrays to PIL Images -3. **Smart Resize**: Applies Qwen's smart resize algorithm based on min/max pixels -4. **Vision Token Formatting**: Inserts vision tokens in text prompt -5. **Batch Preparation**: Creates model-ready BatchFeature input diff --git a/wall_x/serving/client.py b/wall_x/serving/client.py deleted file mode 100644 index 8bb60a92..00000000 --- a/wall_x/serving/client.py +++ /dev/null @@ -1,397 +0,0 @@ -#!/usr/bin/env python3 -""" -Example client for Wall-X model server with sync support. - -This script demonstrates how to connect to a Wall-X server and request -action predictions from observations in both sync and async contexts. -""" - -import asyncio -import logging -from typing import Dict, List -import numpy as np -import threading -import yaml -import torch -import matplotlib.pyplot as plt -import os - -from wall_x.data.utils import update_action_statistics -from wall_x.utils.constant import action_statistic_dof -from wall_x.model.action_head import Normalizer - -try: - import msgpack - import msgpack_numpy as m - - m.patch() -except ImportError: - print("Please install msgpack-numpy: pip install msgpack-numpy") - exit(1) - -try: - import websockets -except ImportError: - print("Please install websockets: pip install websockets") - exit(1) - -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - - -class WallXClient: - """Client for connecting to Wall-X model server.""" - - def __init__( - self, - config_path: str, - uri: str = "ws://localhost:8000", - norm_stats_path: str = "x2_norm_stats.json", - ): - """Initialize client. - - Args: - config_path: Path to train config file - uri: WebSocket URI of the server (e.g., ws://localhost:8000) - norm_stats_path: Path to normalization stats file - """ - self.uri = uri - self.websocket = None - self.metadata = None - self._loop = None - self._thread = None - self.norm_stats_path = norm_stats_path - - with open(config_path, "r") as f: - self.train_config = yaml.load(f, Loader=yaml.FullLoader) - - self.init_normalizer(self.train_config) - - async def connect(self): - """Connect to the server and receive metadata.""" - logger.info(f"Connecting to {self.uri}...") - self.websocket = await websockets.connect( - self.uri, - ping_interval=None, - ping_timeout=None, - max_size=None, - ) - - self.metadata = msgpack.unpackb(await self.websocket.recv()) - logger.info(f"Connected! Server metadata: {self.metadata}") - - async def predict(self, obs: Dict) -> Dict: - """Get action prediction from observation. - - Args: - obs: Observation dictionary containing: - - 'image': Image array (H, W, C) - - 'prompt': Optional text prompt - - 'state': Optional robot state - - Returns: - Dictionary with: - - 'action': Predicted action array - - 'server_timing': Timing information - """ - if self.websocket is None: - raise RuntimeError("Not connected. Call connect() first.") - - await self.websocket.send(msgpack.packb(obs)) - response = msgpack.unpackb(await self.websocket.recv()) - return response - - async def close(self): - """Close the connection.""" - if self.websocket: - await self.websocket.close() - logger.info("Connection closed") - - async def reset(self): - """Reset the policy (if supported).""" - pass - - # ============ Synchronous methods (using independent thread event loop) ============ - - def _start_background_loop(self): - """Start event loop in background thread.""" - self._loop = asyncio.new_event_loop() - asyncio.set_event_loop(self._loop) - self._loop.run_forever() - - def _ensure_loop(self): - """Ensure background event loop is running.""" - if self._loop is None or not self._loop.is_running(): - self._thread = threading.Thread( - target=self._start_background_loop, daemon=True - ) - self._thread.start() - # Wait for loop to start - import time - - while self._loop is None: - time.sleep(0.01) - - def _run_async(self, coro): - """Run coroutine in background event loop.""" - self._ensure_loop() - future = asyncio.run_coroutine_threadsafe(coro, self._loop) - return future.result() - - def connect_sync(self): - """Synchronously connect to server.""" - return self._run_async(self.connect()) - - def norm_state( - self, - state: np.ndarray, - dataset_names: List[str], - state_mask: torch.Tensor = None, - ) -> np.ndarray: - """Normalize state.""" - return self.normalizer_propri.normalize_data(state, dataset_names, state_mask) - - def predict_sync(self, obs: Dict) -> Dict: - """Synchronous prediction method. - - Args: - obs: Observation dictionary - - Returns: - Prediction result dictionary - """ - return self._run_async(self.predict(obs)) - - def close_sync(self): - """Synchronously close connection.""" - result = self._run_async(self.close()) - # Stop event loop - if self._loop: - self._loop.call_soon_threadsafe(self._loop.stop) - return result - - def init_normalizer(self, train_config): - # Define default configurations - dof_config = {"biarm_eed_with_base": 20} - - agent_pos_config = {"biarm_eed_with_base": 20} - - update_action_statistics( - action_statistic_dof=action_statistic_dof, - norm_stats_path=self.norm_stats_path, - repo_id="x2", - dof_config=dof_config, - agent_pos_config=agent_pos_config, - ) - - self.normalizer_action = Normalizer(action_statistic_dof, dof_config) - self.normalizer_propri = Normalizer(action_statistic_dof, agent_pos_config) - - print("Normalizer initialized") - - -def prepare_batch_sync(data, normalizer_action, normalizer_propri, dataset_names): - """Synchronous version of prepare_batch.""" - image = (data["image"].permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy() - wrist_image = ( - (data["wrist_image"].permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy() - ) - prompt = data["task"] - - state = data["state"].to("cuda") - if state.dim() == 1: - state = state.unsqueeze(0) - - state_mask = torch.ones([1, 32, 20]).to("cuda") - state_mask[:, :, 8:] = 0 - - state = normalizer_propri.normalize_data(state, dataset_names, state_mask) - state = state.cpu().numpy().astype(np.float32) - - obs = { - "front_view": image, - "left_wrist_view": wrist_image, - "prompt": prompt, - "state": state, - "dataset_names": dataset_names, - } - return obs - - -def init_serving_sample_dataset(train_config): - from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata - - repo_id = train_config["data"]["lerobot_config"]["repo_id"] - - meta_info = LeRobotDatasetMetadata(repo_id) - dataset_fps = meta_info.fps - delta_timestamps = { - "actions": [t / dataset_fps for t in range(32)], - } - dataset = LeRobotDataset( - repo_id, - episodes=[0], - delta_timestamps=delta_timestamps, - video_backend="pyav", - ) - - return dataset, repo_id - - -# ============ Synchronous version of main function ============ - - -def main_sync(args): - """Synchronous version of main function.""" - - # Create client and connect - client = WallXClient( - args.config_path, uri=args.uri, norm_stats_path=args.norm_stats_path - ) - client.connect_sync() - - dataset, repo_id = init_serving_sample_dataset(client.train_config) - - total_frames = len(dataset) - gt_traj = np.zeros((total_frames, args.action_dim)) - pred_traj = np.zeros((total_frames, args.action_dim)) - import torch - - dof_mask = torch.ones([1, 32, 20]).to("cuda") - dof_mask[:, :, args.action_dim :] = 0 - - # Synchronous processing - for idx, data in enumerate(dataset): - if idx % args.pred_horizon == 0 and idx + args.pred_horizon < total_frames: - print(f"Processing frame {idx}") - obs = prepare_batch_sync( - data, - client.normalizer_action, - client.normalizer_propri, - dataset_names=[repo_id], - ) - response = client.predict_sync(obs) - pred_action = response["action"] - pred_traj[idx : idx + args.pred_horizon] = pred_action - gt_traj[idx : idx + args.pred_horizon] = data["actions"] - - # Draw plot - timesteps = gt_traj.shape[0] - fig, axs = plt.subplots( - args.action_dim, 1, figsize=(15, 5 * args.action_dim), sharex=True - ) - fig.suptitle("Action Comparison for lerobot", fontsize=16) - - for i in range(args.action_dim): - axs[i].plot(range(timesteps), gt_traj[:, i], label="Ground Truth") - axs[i].plot(range(timesteps), pred_traj[:, i], label="Prediction") - axs[i].set_ylabel(f"Action Dim {i+1}") - axs[i].legend() - axs[i].grid(True) - - axs[-1].set_xlabel("Timestep") - plt.tight_layout(rect=[0, 0.03, 1, 0.95]) - os.makedirs(args.save_dir, exist_ok=True) - save_path = os.path.join(args.save_dir, "lerobot_comparison_serving.png") - plt.savefig(save_path) - print(f"Saved plot to {save_path}") - plt.close() - - # Close connection - client.close_sync() - - -# ============ Asynchronous version of main function (keep original functionality) ============ - - -async def main(args): - client = WallXClient( - args.config_path, uri=args.uri, norm_stats_path=args.norm_stats_path - ) - await client.connect() - dataset, repo_id = init_serving_sample_dataset(client.train_config) - - total_frames = len(dataset) - gt_traj = np.zeros((total_frames, args.action_dim)) - pred_traj = np.zeros((total_frames, args.action_dim)) - - for idx, data in enumerate(dataset): - if idx % args.pred_horizon == 0 and idx + args.pred_horizon < total_frames: - print(f"Processing frame {idx}") - obs = prepare_batch_sync( - data, - client.normalizer_action, - client.normalizer_propri, - dataset_names=[repo_id], - ) - response = await client.predict(obs) - pred_action = response["action"] - print(pred_action.shape) - pred_traj[idx : idx + args.pred_horizon] = pred_action - gt_traj[idx : idx + args.pred_horizon] = data["actions"] - - timesteps = gt_traj.shape[0] - - fig, axs = plt.subplots( - args.action_dim, 1, figsize=(15, 5 * args.action_dim), sharex=True - ) - fig.suptitle("Action Comparison for lerobot", fontsize=16) - - for i in range(args.action_dim): - axs[i].plot(range(timesteps), gt_traj[:, i], label="Ground Truth") - axs[i].plot(range(timesteps), pred_traj[:, i], label="Prediction") - axs[i].set_ylabel(f"Action Dim {i+1}") - axs[i].legend() - axs[i].grid(True) - - axs[-1].set_xlabel("Timestep") - plt.tight_layout(rect=[0, 0.03, 1, 0.95]) - os.makedirs(args.save_dir, exist_ok=True) - save_path = os.path.join(args.save_dir, "lerobot_comparison_serving.png") - plt.savefig(save_path) - print(f"Saved plot to {save_path}") - plt.close() - - -if __name__ == "__main__": - """Asynchronous version of main function.""" - import argparse - - parser = argparse.ArgumentParser(description="Wall-X client examples") - parser.add_argument( - "--example", - choices=["single", "multiple", "benchmark"], - default="single", - help="Example to run", - ) - parser.add_argument( - "--uri", - default="ws://localhost:8000", - help="Server URI", - ) - parser.add_argument( - "--pred_horizon", type=int, default=32, help="Prediction horizon" - ) - parser.add_argument("--action_dim", type=int, default=7, help="Action dimension") - parser.add_argument( - "--config_path", - default="config_from_qwen_libero.yml", - help="Train config path", - ) - parser.add_argument( - "--save_dir", - default="libero", - help="Save directory", - ) - parser.add_argument( - "--norm_stats_path", - default="x2_norm_stats.json", - help="Normalization stats path", - ) - args = parser.parse_args() - - # Synchronous mode - main_sync(args) - - # Asynchronous mode - # asyncio.run(main(args)) diff --git a/wall_x/serving/launch_serving.py b/wall_x/serving/launch_serving.py deleted file mode 100644 index c675a6ab..00000000 --- a/wall_x/serving/launch_serving.py +++ /dev/null @@ -1,220 +0,0 @@ -#!/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. - -Based on the OpenPI serve_policy.py script structure. -""" - -import dataclasses -from dataclasses import field -import enum -import logging -import socket -import sys -import yaml -from pathlib import Path -from typing import List - -import tyro - -from wall_x.serving.policy.wall_x_policy import WallXPolicy -from wall_x.serving.websocket_policy_server import WebsocketPolicyServer - -logger = logging.getLogger(__name__) - - -class EnvMode(enum.Enum): - """Supported environments/datasets.""" - - LIBERO = "libero" - ALOHA = "aloha" - - -@dataclasses.dataclass -class ModelConfig: - """Configuration for loading a Wall-X model.""" - - # Path to the pretrained model checkpoint - model_path: str - # Path to the action tokenizer - action_tokenizer_path: str - # Path to train config yaml - train_config_path: str - # Action dimension for the environment - action_dim: int = 7 - # State dimension for the environment - state_dim: int = 8 - # Prediction horizon (number of future actions to predict) - pred_horizon: int = 32 - # Device to run model on - device: str = "cuda" - # Model dtype (bfloat16, float16, float32) - dtype: str = "bfloat16" - # Prediction mode (fast or slow) - predict_mode: str = "fast" - # Camera key for the environment - camera_key: List[str] = field( - default_factory=lambda: ["front_view", "left_wrist_view", "right_wrist_view"] - ) - - -@dataclasses.dataclass -class Args: - """Arguments for the serve_wall_x script.""" - - # Environment mode (used for default configurations) - env: EnvMode = EnvMode.LIBERO - - # Model configuration. If not provided, uses default config for the environment - model_config: ModelConfig | 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 = 8000 - - # Host to bind the server to - host: str = "0.0.0.0" - - # Enable debug logging - debug: bool = False - - -# Default model configurations for each environment -DEFAULT_CONFIGS: dict[EnvMode, ModelConfig] = { - EnvMode.LIBERO: ModelConfig( - model_path="/path/to/model", - action_tokenizer_path="/path/to/action_tokenizer", - train_config_path="/path/to/train_config", - state_dim=8, - action_dim=7, - pred_horizon=32, - device="cuda", - dtype="bfloat16", - predict_mode="fast", - camera_key=["front_view", "left_wrist_view"], - ), - EnvMode.ALOHA: ModelConfig( - model_path="/path/to/model", - action_tokenizer_path="/path/to/action_tokenizer", - train_config_path="/path/to/train_config", - state_dim=14, - action_dim=14, - pred_horizon=32, - device="cuda", - dtype="bfloat16", - predict_mode="fast", - camera_key=["face_view", "left_wrist_view", "right_wrist_view"], - ), -} - - -def get_model_config(args: Args) -> ModelConfig: - """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) -> WallXPolicy: - """Create a Wall-X policy from the given arguments.""" - config = get_model_config(args) - logger.info(f"Creating Wall-X policy with config: {config}") - - # Validate paths - if not Path(config.model_path).exists(): - logger.warning(f"Model path does not exist: {config.model_path}") - - if not Path(config.action_tokenizer_path).exists(): - logger.warning( - f"Action tokenizer path does not exist: {config.action_tokenizer_path}" - ) - - with open(config.train_config_path, "r") as f: - train_config = yaml.load(f, Loader=yaml.FullLoader) - - policy = WallXPolicy( - model_path=config.model_path, - train_config=train_config, - action_tokenizer_path=config.action_tokenizer_path, - action_dim=config.action_dim, - agent_pos_dim=config.state_dim, - pred_horizon=config.pred_horizon, - device=config.device, - dtype=config.dtype, - predict_mode=config.predict_mode, - default_prompt=args.default_prompt, - camera_key=config.camera_key, - ) - - 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", - ) - - logger.info("Starting Wall-X model server") - logger.info(f"Environment: {args.env.value}") - logger.info(f"Port: {args.port}") - logger.info(f"Host: {args.host}") - - # Create policy - try: - policy = create_policy(args) - except Exception as e: - logger.error(f"Failed to create policy: {e}") - sys.exit(1) - - # Get policy metadata - policy_metadata = policy.metadata - policy_metadata["env"] = args.env.value - - # 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") - - # Create and start server - server = WebsocketPolicyServer( - policy=policy, - host=args.host, - port=args.port, - metadata=policy_metadata, - ) - - logger.info("Starting server...") - try: - server.serve_forever() - except KeyboardInterrupt: - logger.info("Server stopped by user") - except Exception as e: - logger.error(f"Server error: {e}") - sys.exit(1) - - -if __name__ == "__main__": - main(tyro.cli(Args)) diff --git a/wall_x/serving/policy/utils.py b/wall_x/serving/policy/utils.py deleted file mode 100644 index d13d7683..00000000 --- a/wall_x/serving/policy/utils.py +++ /dev/null @@ -1,265 +0,0 @@ -from typing import Dict, List -import logging -import numpy as np -from wall_x.data.utils import preprocesser_call -from qwen_vl_utils.vision_process import smart_resize -import torch -from PIL import Image -from transformers import BatchFeature - -logger = logging.getLogger(__name__) - - -def prepare_batch( - obs: Dict, - processor, - normalizer_propri, - camera_key: List[str], - agent_pos_dim, - action_dim, - pred_horizon, - fixed_action_dim, - max_length, - image_factor: int, - min_pixels: int, - max_pixels: int, - predict_mode: str = "fast", - device: str = "cuda", -) -> BatchFeature: - """Prepare observation into model input format. - - Args: - obs: Dictionary containing: - - 'camera_key_0' : image 0 - - 'camera_key_1' : image 1 - ... - - 'prompt': Text prompt - - 'state': Robot state/proprioception - - 'dataset_names': Dataset names - - Returns: - BatchFeature object ready for model input - """ - # Handle images - can be single image, list of images, or dict of images - images = [] - images = [obs[key] for key in camera_key] - # Convert numpy arrays to PIL Images - processed_images = [] - for img in images: - if isinstance(img, np.ndarray): - # Debug: Log the shape and dtype - logger.debug(f"Image shape: {img.shape}, dtype: {img.dtype}") - - # Handle unexpected dimensions - squeeze if needed - if img.ndim > 3: - logger.warning( - f"Image has {img.ndim} dimensions, squeezing extra dimensions" - ) - img = np.squeeze(img) - - # Verify shape is valid for PIL - if img.ndim == 2: - # Grayscale image - pass - elif img.ndim == 3: - # Check if channel dimension is first or last - if img.shape[0] == 3 or img.shape[0] == 1: - # Channels first, transpose to channels last - img = np.transpose(img, (1, 2, 0)) - elif img.shape[2] == 3 or img.shape[2] == 1: - # Already channels last - pass - else: - raise ValueError( - f"Unexpected image shape: {img.shape}. Expected (H, W, C) or (C, H, W)" - ) - else: - raise ValueError( - f"Invalid image dimensions: {img.ndim}. Expected 2 or 3 dimensions, got shape {img.shape}" - ) - - # Convert to PIL Image - if img.dtype == np.uint8: - img = Image.fromarray(img) - else: - img = Image.fromarray((img * 255).astype(np.uint8)) - processed_images.append(img) - - # print("processed_images:",processed_images) - # Apply smart resize to images - resized_images = process_images( - processed_images, image_factor, min_pixels, max_pixels - ) - - # Handle text prompt - format with vision tokens - instruction = obs["prompt"] - formatted_text = format_text_with_vision_tokens( - instruction, camera_key, predict_mode, pred_horizon - ) - - # Use processor to prepare inputs - inputs = preprocesser_call( - processor=processor, - text=[formatted_text], - images=[resized_images], - videos=None, - padding=True, - truncation=True, - return_tensors="pt", - max_length=max_length, - ) - - action_token_id = processor.tokenizer.convert_tokens_to_ids("<|action|>") - moe_token_types = inputs.input_ids == action_token_id - inputs["moe_token_types"] = torch.tensor(moe_token_types) - - # obs["dataset_names"]="libero_all" - - # Handle robot state/proprioception if available - if "state" in obs: - state = obs["state"] - if isinstance(state, np.ndarray): - state = torch.from_numpy(state).float() - elif not isinstance(state, torch.Tensor): - state = torch.tensor(state, dtype=torch.float32) - - # Add batch dimension if needed - if state.dim() == 1: - state = state.unsqueeze(0) - if state.dim() == 2: - state = state.unsqueeze(1) # [batch, 1, state_dim] - - # Pad to 20 dimensions if needed (same as training) - # if state.shape[-1] < 20: - # padding = torch.zeros(state.shape[0], state.shape[1], 20 - state.shape[-1]) - # state = torch.cat([state, padding], dim=-1) - - # Create mask for valid dimensions - agent_pos_mask = torch.ones_like(state) - if state.shape[-1] > agent_pos_dim: - agent_pos_mask[:, :, agent_pos_dim:] = 0 - - normalizer_propri.normalize_data(state, [obs["dataset_names"]] * state.shape[0]) - - inputs["proprioception"] = state - inputs["agent_pos_mask"] = agent_pos_mask - - # Add dataset name (required by model) - inputs["dataset_names"] = [obs["dataset_names"]] * state.shape[0] - - # Move all tensors to device - for key in inputs: - if isinstance(inputs[key], torch.Tensor): - inputs[key] = inputs[key].to(device) - - dof_mask = torch.ones([state.shape[0], pred_horizon, fixed_action_dim]) - dof_mask[:, :, action_dim:] = 0 - - inputs["dof_mask"] = dof_mask - - # Convert to BatchFeature to maintain consistency with training pipeline - return BatchFeature(data=dict(inputs)).to(device) - - -def process_images( - images: List[Image.Image], image_factor: int, min_pixels: int, max_pixels: int -) -> List[Image.Image]: - """Process images with smart resize following the data loading pattern. - - Args: - images: List of PIL Images - - Returns: - List of resized PIL Images - """ - resized_images = [] - for img_pil in images: - - orig_width, orig_height = img_pil.size - target_size = 256 - if target_size != -1: - # Maintain aspect ratio logic - if orig_width > orig_height: # Landscape image - new_width = target_size - new_height = int(target_size * orig_height / orig_width) - else: # Portrait image - new_height = target_size - new_width = int(target_size * orig_width / orig_height) - img_pil = img_pil.resize((new_width, new_height)) - - # Apply smart scaling (Qwen logic) - current_width, current_height = img_pil.size - resized_height, resized_width = smart_resize( - current_height, - current_width, - factor=image_factor, - min_pixels=min_pixels, - max_pixels=max_pixels, - ) - - resized_img = img_pil.resize((resized_width, resized_height)) - resized_images.append(resized_img) - - return resized_images - - -def format_text_with_vision_tokens( - instruction: str, - camera_key: List[str], - predict_mode: str = "diffusion", - pred_horizon: int = 32, -) -> str: - """Format text prompt with vision tokens for the model. - - Args: - instruction: Task instruction text - camera_key: List of camera names - - Returns: - Formatted text with special tokens - """ - # Special tokens for formatting - role_start_symbol = "<|im_start|>" - role_end_symbol = "<|im_end|>" - vision_start_symbol = "<|vision_start|>" - vision_end_symbol = "<|vision_end|>" - image_pad_symbol = "<|image_pad|>" - propri_symbol = "<|propri|>" - action_symbol = "<|action|>" - action_fast_symbol = "<|action_fast|>" - - # Camera name mapping - camera_name_mapping = { - "front_view": "front view", - "face_view": "front view", - "left_wrist_view": "left wrist view", - "right_wrist_view": "right wrist view", - "top_view": "top view", - "wall_view": "wall view", - } - - # System prologue - prologue = ( - f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n" - ) - - # User request with observation - user_request = f"{role_start_symbol}user\nObservation:" - if camera_key: - for cam_name in camera_key: - view_name = camera_name_mapping.get(cam_name, cam_name) - user_request += f" {view_name}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}" - user_request += "\nInstruction:" - - text_prompt = ( - f"\nPredict the next action in robot action.\nProprioception: {propri_symbol}\n" - ) - user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n" - assistant_output = ( - f"{role_start_symbol}assistant\n{action_fast_symbol}{role_end_symbol}\n" - ) - if predict_mode == "diffusion": - assistant_output = f"{role_start_symbol}assistant\n{action_symbol * pred_horizon}{role_end_symbol}\n" - complete_text = prologue + user_message + assistant_output - - return complete_text diff --git a/wall_x/serving/policy/wall_x_policy.py b/wall_x/serving/policy/wall_x_policy.py deleted file mode 100644 index e2d6fa98..00000000 --- a/wall_x/serving/policy/wall_x_policy.py +++ /dev/null @@ -1,182 +0,0 @@ -import logging -from typing import Dict, Any, List -import torch -import copy -import numpy as np -from wall_x.serving.websocket_policy_server import BasePolicy -from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction -from wall_x.serving.policy.utils import prepare_batch -from wall_x.model.model_utils import load_wallx_processors, register_normalizers - -logger = logging.getLogger(__name__) - - -class WallXPolicy(BasePolicy): - """Policy wrapper for Wall-X model that implements the BasePolicy interface.""" - - def __init__( - self, - model_path: str, - train_config: dict, - action_tokenizer_path: str, - action_dim: int, - agent_pos_dim: int, - pred_horizon: int, - camera_key: List[str], - device: str = "cuda", - dtype: str = "bfloat16", - predict_mode: str = "diffusion", - default_prompt: str | None = None, - min_pixels: int = 4 * 28 * 28, - max_pixels: int = 16384 * 28 * 28, - image_factor: int = 28, - max_length: int = 2048, - ): - """Initialize the Wall-X policy. - - Args: - model_path: Path to the pretrained model checkpoint - action_tokenizer_path: Path to the action tokenizer - action_dim: Dimension of action space - pred_horizon: Prediction horizon for actions - device: Device to run model on ('cuda' or 'cpu') - dtype: Data type for model ('bfloat16', 'float16', or 'float32') - predict_mode: Prediction mode ('fast' or 'slow') - default_prompt: Default text prompt for the model - min_pixels: Minimum pixels for image resizing - max_pixels: Maximum pixels for image resizing - image_factor: Factor for smart resize - max_length: Maximum sequence length for text - """ - logger.info(f"Loading Wall-X model from {model_path}") - - self.normalizer_action, self.normalizer_propri = register_normalizers( - train_config, model_path - ) - - self.model = Qwen2_5_VLMoEForAction.from_pretrained( - model_path, - train_config=train_config, - action_tokenizer_path=action_tokenizer_path, - ) - self.model.set_normalizer( - copy.deepcopy(self.normalizer_action), copy.deepcopy(self.normalizer_propri) - ) - self.model.eval() - self.model = self.model.to(device) - self.model.to_bfloat16_for_selected_params() - - # hard code the action dim to 20 for align to wall-x configuration - self.fixed_action_dim = action_dim - - self.action_dim = action_dim - self.agent_pos_dim = action_dim - self.pred_horizon = pred_horizon - self.device = device - self.predict_mode = predict_mode - self.default_prompt = default_prompt - self.camera_key = camera_key - - # Image preprocessing config - self.min_pixels = min_pixels - self.max_pixels = max_pixels - self.image_factor = image_factor - self.max_length = max_length - - print("predict_mode", predict_mode) - print("camera_key", camera_key) - - # Load processor - logger.info("Loading processor and tokenizer...") - - processors_dict = load_wallx_processors(train_config) - self.processor = processors_dict["processor"] - - # Action buffer for multi-step predictions - self.action_buffer = [] - self.buffer_index = 0 - - logger.info( - f"Model loaded successfully. Device: {device}, Action dim: {action_dim}, Horizon: {pred_horizon}" - ) - - @property - def metadata(self) -> Dict[str, Any]: - """Return metadata about the policy.""" - return { - "action_dim": self.action_dim, - "pred_horizon": self.pred_horizon, - "device": self.device, - "predict_mode": self.predict_mode, - } - - def reset(self) -> None: - """Reset the policy state.""" - self.action_buffer = [] - self.buffer_index = 0 - logger.debug("Policy reset") - - def infer(self, obs: Dict) -> Dict: - """Infer action from observation. - - Args: - obs: Dictionary containing: - - 'image': Image observation (numpy array or PIL Image) - - 'prompt': Optional text prompt - - 'state': Optional robot state - - Other modality-specific observations - - Returns: - Dictionary containing: - - 'action': Predicted action (numpy array) - - Additional metadata - """ - try: - # Need to predict new actions - input_batch = prepare_batch( - obs, - self.processor, - self.normalizer_propri, - self.camera_key, - self.agent_pos_dim, - self.action_dim, - self.pred_horizon, - self.fixed_action_dim, - self.max_length, - self.image_factor, - self.min_pixels, - self.max_pixels, - self.predict_mode, - self.device, - ) - - with torch.no_grad(): - outputs = self.model( - **input_batch, - action_dim=( - self.action_dim - if self.predict_mode == "fast" - else self.fixed_action_dim - ), - action_horizon=self.pred_horizon, - mode="predict", - predict_mode=self.predict_mode, - ) - - if outputs["predict_action"] is None: - predicted_actions = np.zeros( - [1, self.pred_horizon, self.action_dim] - ).astype(np.float32) - - predicted_actions = ( - outputs["predict_action"][:, :, : self.action_dim] - .detach() - .cpu() - .to(torch.float32) - .numpy() - ) - return {"predict_action": predicted_actions} - - except Exception as e: - logger.error(f"Error during inference: {e}") - raise diff --git a/wall_x/trainer/adapters/__init__.py b/wall_x/trainer/adapters/__init__.py new file mode 100644 index 00000000..85921ca2 --- /dev/null +++ b/wall_x/trainer/adapters/__init__.py @@ -0,0 +1,129 @@ +"""Public model adapter registry.""" + +from importlib import import_module + +from wall_x.trainer.adapters.base_adapter import ModelAdapter +from wall_x.trainer.adapters.vla_model_adapter import VLAdapter + + +_ADAPTER_SPECS = { + "qwen2_5": ("wall_x.model.qact.qwen2_5.adapter", "Qwen2_5Adapter"), +} +_ADAPTER_CLASS_EXPORTS = { + "Qwen2_5Adapter": "qwen2_5", +} +ADAPTER_IMPORT_ERRORS: dict[str, str] = {} +_ADAPTER_LOADING: set[str] = set() + + +class _AdapterRegistry(dict[str, type[ModelAdapter]]): + def get(self, model_type: str, default=None): + _load_public_adapter(model_type) + return dict.get(self, model_type, default) + + def __contains__(self, model_type: object) -> bool: + if isinstance(model_type, str): + _load_public_adapter(model_type) + return dict.__contains__(self, model_type) + + def __getitem__(self, model_type: str) -> type[ModelAdapter]: + _load_public_adapter(model_type) + return dict.__getitem__(self, model_type) + + def __iter__(self): + _load_all_public_adapters() + return dict.__iter__(self) + + def keys(self): + _load_all_public_adapters() + return dict.keys(self) + + def items(self): + _load_all_public_adapters() + return dict.items(self) + + def values(self): + _load_all_public_adapters() + return dict.values(self) + + +ADAPTER_REGISTRY: _AdapterRegistry = _AdapterRegistry() + + +def _register_adapter(model_type: str, adapter_cls: type[ModelAdapter]) -> None: + existing = dict.get(ADAPTER_REGISTRY, model_type) + if existing is not None and existing is not adapter_cls: + raise ValueError( + f"model_type {model_type!r} already registered to " + f"{existing.__name__}, cannot re-register to {adapter_cls.__name__}" + ) + declared_model_type = getattr(adapter_cls, "MODEL_TYPE", model_type) + if declared_model_type and declared_model_type != model_type: + raise ValueError( + f"adapter {adapter_cls.__name__} declares MODEL_TYPE=" + f"{declared_model_type!r}, but is registered as {model_type!r}" + ) + dict.__setitem__(ADAPTER_REGISTRY, model_type, adapter_cls) + + +def _load_public_adapter(model_type: str) -> type[ModelAdapter] | None: + if dict.__contains__(ADAPTER_REGISTRY, model_type): + return dict.__getitem__(ADAPTER_REGISTRY, model_type) + spec = _ADAPTER_SPECS.get(model_type) + if spec is None or model_type in _ADAPTER_LOADING: + return None + module_name, class_name = spec + _ADAPTER_LOADING.add(model_type) + try: + module = import_module(module_name) + adapter_cls = getattr(module, class_name) + except (ImportError, AttributeError) as exc: + ADAPTER_IMPORT_ERRORS[model_type] = f"{type(exc).__name__}: {exc}" + return None + finally: + _ADAPTER_LOADING.discard(model_type) + _register_adapter(model_type, adapter_cls) + return adapter_cls + + +def _load_all_public_adapters() -> None: + for model_type in _ADAPTER_SPECS: + _load_public_adapter(model_type) + + +def format_adapter_error(model_type: str) -> str: + msg = ( + f"Unsupported model type: {model_type}. " + f"Registered: {sorted(ADAPTER_REGISTRY)}" + ) + if model_type in ADAPTER_IMPORT_ERRORS: + msg += f". Import failed: {ADAPTER_IMPORT_ERRORS[model_type]}" + return msg + + +def resolve_adapter(model_type: str) -> type[ModelAdapter]: + """Return a registered adapter class, loading lazily when needed.""" + adapter_cls = ADAPTER_REGISTRY.get(model_type) + if adapter_cls is None: + raise ValueError(format_adapter_error(model_type)) + return adapter_cls + + +def __getattr__(name: str): + model_type = _ADAPTER_CLASS_EXPORTS.get(name) + if model_type is not None: + adapter_cls = _load_public_adapter(model_type) + if adapter_cls is not None: + return adapter_cls + raise AttributeError(name) + + +__all__ = [ + "ModelAdapter", + "VLAdapter", + "Qwen2_5Adapter", + "ADAPTER_REGISTRY", + "ADAPTER_IMPORT_ERRORS", + "format_adapter_error", + "resolve_adapter", +] diff --git a/wall_x/trainer/adapters/base_adapter.py b/wall_x/trainer/adapters/base_adapter.py new file mode 100644 index 00000000..269da803 --- /dev/null +++ b/wall_x/trainer/adapters/base_adapter.py @@ -0,0 +1,495 @@ +import dataclasses +import logging +import os +from abc import ABC, abstractmethod +from collections import defaultdict + +import torch +import torch.distributed as dist +from tqdm import tqdm + +from wall_x.trainer.utils import move_batch_to_device + +logger = logging.getLogger(__name__) + + +def load_trainer_data_config(cfg): + """Return optional backend-specific trainer data config. + + Public backends build datasets directly from TrainConfig. A backend that + needs an additional trainer data config may expose + ``load_trainer_data_config(cfg)``. + """ + from wall_x.data import data_backend + + backend = data_backend() + if backend.supports("load_trainer_data_config"): + return backend.load_trainer_data_config(cfg) + return None + + +class ModelAdapter(ABC): + """Pure strategy that encapsulates model-type-specific behavior. + + The adapter is stateless w.r.t. the trainer: every method declares + its inputs as parameters and communicates results via return values. + The trainer orchestrates calls and manages its own state. + + Abstract methods (subclasses MUST implement): + load_processor -- load processor / tokenizer artifacts + build_model_config -- build model-specific config object + create_model -- instantiate the model + load_weights -- load pretrained weights & set normalizers + get_transformer_layer_cls -- layer classes for FSDP auto-wrap + load_dataset -- build dataset & training dataloader + forward -- run a forward pass + extract_loss -- extract scalar loss from model outputs + + Concrete methods with default behaviour (override when needed): + load_step_and_epoch -- restore global_step / epoch from checkpoint + log_model_info -- log model-specific info (no-op) + collect_output_stats -- collect auxiliary loss / accuracy stats (no-op) + collect_param_norms -- compute per-group param L2 norms (no-op) + get_fwd_flops -- optional forward FLOPs hook + get_output_field -- retrieve a named field from outputs (static) + """ + + def __init__(self, *, cfg=None, logger=None, model_type=None): + self.cfg = cfg + self.logger = logger + + # ---- processor ---- + @abstractmethod + def load_processor(self, action_statistic_dof): + """Load processor and associated tokenizer artifacts. + + Reads model/data config from ``self.cfg``. + + Args: + action_statistic_dof: normalizer statistics dict. + + Returns: + dict with keys: + "processor", "data_config", "tokenizer_mixin", + "train_action_tokenizer", "val_action_tokenizer", + "action_mapper", "num_added_tokens" + """ + ... + + # ---- model ---- + @abstractmethod + def create_model(self, processor, tokenizer_mixin, model_config): + """Instantiate the model (before weight loading / FSDP wrap). + + Reads model config from ``self.cfg``. + """ + ... + + @abstractmethod + def build_model_config(self): + """Build the model-specific config object from ``self.cfg.model``.""" + ... + + @abstractmethod + def load_weights(self, model, normalizer_action, normalizer_propri, **kwargs): + """Load pretrained weights, resize embeddings, set normalizers, etc.""" + ... + + # ---- FSDP wrapping ---- + @abstractmethod + def get_transformer_layer_cls(self): + """Return layer classes for FSDP transformer_auto_wrap_policy.""" + ... + + # ---- dataset ---- + @abstractmethod + def load_dataset(self, data_config, processor, rank, world_size, **kwargs): + """Load dataset and build the training dataloader. + + Args: + data_config: parsed data config from load_processor(). + processor: processor instance. + rank: current process rank. + world_size: total number of processes. + **kwargs: extra model-specific args. + + Returns: + (dataset, train_dataloader, train_num) tuple. + """ + ... + + def load_step_and_epoch(self, checkpoint_path: str, is_incomplete_epoch: bool): + global_step = 0 + start_epoch = 0 + + global_step_path = os.path.join(checkpoint_path, "global_step.pth") + if os.path.exists(global_step_path): + global_step = torch.load(global_step_path)["global_step"] + + current_epoch_path = os.path.join(checkpoint_path, "current_epoch.pth") + if os.path.exists(current_epoch_path): + start_epoch = torch.load(current_epoch_path)["current_epoch"] + if not is_incomplete_epoch: + start_epoch = start_epoch + 1 + + return {"global_step": global_step, "start_epoch": start_epoch} + + # ---- forward / loss ---- + @abstractmethod + def forward(self, model, batch, **kwargs): + """Run a forward pass. + + Args: + model: the (possibly wrapped) model. + batch: dict of tensors already on device. + **kwargs: extra context (e.g. global_step, mode). + + Returns: + Raw model outputs (dict or object). + """ + ... + + @abstractmethod + def extract_loss(self, outputs): + """Extract the scalar training loss from model outputs. + + Returns: + torch.Tensor (scalar). + """ + ... + + # ---- optional hooks (override when needed) ---- + def log_model_info(self, model, log_fn): + """Log model-specific info (e.g. attention implementation). + + Args: + model: the unwrapped model. + log_fn: callable(str) for logging. + """ + pass + + def collect_output_stats( + self, + outputs, + step_stats, + reduce_tensor_fn, + true_gather_fn, + tokenizer_mixin=None, + ): + """Collect model-output-specific statistics into step_stats. + + Extracts auxiliary losses (cross_entropy_loss, flow_loss), + per-dataset channel losses, and action accuracy metrics from + the model outputs and merges them into *step_stats* in-place. + + This is an optional hook — subclasses override it when their + output format carries extra fields. The default implementation + is a no-op. + + Args: + outputs: raw model outputs (dict or object). + step_stats: dict being built by the trainer; mutated in-place. + reduce_tensor_fn: callable(tensor, average=True) → scalar tensor, + wraps dist.all_reduce for the current parallelism. + true_gather_fn: callable(tensor) → scalar | None, + gathers a value to rank-0 (returns None on + non-main ranks or when input is None). + tokenizer_mixin: optional ActionTokenizerMixin used to retrieve + extra accuracy keys (e.g. RVQ layer accuracy). + """ + pass + + def collect_param_norms( + self, model, step_stats, device, reduce_tensor_fn, params_sharded=False + ): + """Compute per-group parameter L2 norms and merge into step_stats. + + Groups parameters by name keyword: + - "visual" → visual_param_norm + - "action" → action_expert_param_norm + - (others) → org_vlm_param_norm + - all → total_param_norm + + When parameters are sharded across ranks (FSDP FULL_SHARD / + HYBRID_SHARD), the local squared-sum is partial and must be + all_reduce'd (sum, not avg) before taking the square root. + When parameters are replicated (DDP, or FSDP with SHARD_GRAD_OP / + NO_SHARD), all_reduce would inflate the result by world_size. + + This is an optional hook — the default implementation is a no-op. + + Args: + model: the (possibly FSDP/DDP-wrapped) model. + step_stats: dict being built by the trainer; mutated in-place. + device: torch.device for tensor allocation. + reduce_tensor_fn: callable(tensor, average=False) → tensor, + wraps dist.all_reduce for the current parallelism. + params_sharded: True when parameters are actually sharded across + ranks (FSDP FULL_SHARD / HYBRID_SHARD). + """ + pass + + def collect_grad_norms( + self, model, step_stats, device, reduce_tensor_fn, params_sharded=False + ): + """Compute per-component gradient L2 norms and merge into step_stats. + + Must be called while gradients are still available (before + optimizer.zero_grad). Under FSDP each rank holds a gradient shard, + so local squared-sums must be all_reduce'd (sum, not avg) before + taking the square root. + + Default is a no-op; subclasses override for model-specific groupings. + """ + pass + + def get_fwd_flops(self, autocast_context, model, batch, global_step): + """Forward FLOPs profiling is disabled in the public package.""" + return None + + @staticmethod + def get_output_field(outputs, key, default=None): + """Retrieve a named field from model outputs (dict or object).""" + if isinstance(outputs, dict): + return outputs.get(key, default) + return getattr(outputs, key, default) + + # ---- console logging fields ---- + def console_fields(self, tokenizer_mixin=None) -> list: + """Per-step console fields as (stat_key, pretty_label, fmt_spec) triples. + + MetricsLogger renders them in order, skipping any key not present + in step_stats. Public adapters may override this method to add + task-specific fields. + """ + fields = [ + ("video_loss", "vid_loss", ".6f"), + ("action_loss", "act_loss", ".6f"), + ("action_accuracy", "accuracy", ".4f"), + ("flow_loss", "flow_loss", ".6f"), + ] + if tokenizer_mixin is not None: + for key in tokenizer_mixin.get_accuracy_keys(): + if key != "action_accuracy": + short_key = key.replace("action_accuracy_", "acc_") + fields.append((key, short_key, ".4f")) + fields.extend( + [ + ("action_grad_norm", "act_gnorm", ".4f"), + ("video_grad_norm", "vid_gnorm", ".4f"), + ] + ) + return fields + + # ---- distribution hooks ---- + def convert_to_mix_precision_hint( + self, + model, + *, + device, + use_fsdp: bool, + log_fn=None, + ): + """Prepare *model* dtype + placement before FSDP/DDP wrap. + + FSDP path: cast to fp32 on device; MixedPrecision policy handles the + bf16 compute cast at forward time. DDP path: delegate to the model's + ``convert_to_mix_precision`` method (manual bf16 conversion, since + DDP has no equivalent of MixedPrecision), then move to device. + + Adapters that handle placement themselves may override this hook. + """ + _log = log_fn or (lambda _msg: None) + if use_fsdp: + model.to(dtype=torch.float32) + _log("FSDP mode: all params fp32, MixedPrecision handles bf16 compute") + model.to(device) + else: + model.convert_to_mix_precision() + _log("DDP mode: params converted to bf16 (no MixedPrecision policy)") + model.to(device) + + # ---- validation ---- + def run_validation( + self, + *, + model, + val_dataloader, + rank, + world_size, + device, + autocast_context, + reduce_fn, + gather_fn, + logger, + global_step, + output_path, + tokenizer_mixin=None, + ): + """Default validation: forward-pass loss + adapter.collect_output_stats. + + Public VLA adapters may override this for task-specific metrics. + + ``output_path`` is accepted for interface uniformity but not used + in this default. + """ + del output_path + model.eval() + log_dict = defaultdict(float) + pbar = tqdm( + val_dataloader, + desc="Validating", + total=len(val_dataloader), + disable=rank != 0, + ) + + for batch in pbar: + batch = move_batch_to_device(batch, device) + with autocast_context(): + outputs = self.forward(model, batch, global_step=0) + + loss = outputs["total_loss"] if "total_loss" in outputs else outputs["loss"] + log_dict["val_loss"] += reduce_fn(loss.detach()).item() + + val_ce_loss = self.get_output_field(outputs, "cross_entropy_loss") + if val_ce_loss is not None: + log_ce_loss = gather_fn(val_ce_loss) + if log_ce_loss is not None and rank == 0: + log_dict["val_cross_entropy_loss"] += log_ce_loss + val_flow_loss = self.get_output_field(outputs, "flow_loss") + if val_flow_loss is not None: + log_flow_loss = gather_fn(val_flow_loss) + if log_flow_loss is not None and rank == 0: + log_dict["val_flow_loss"] += log_flow_loss + + for k, v in collect_channel_loss_stats( + outputs, + prefix="val_", + tokenizer_mixin=tokenizer_mixin, + ).items(): + log_dict[k] += v + + log_dict = {k: v / len(val_dataloader) for k, v in log_dict.items()} + + if rank == 0: + if self.logger is not None: + self.logger.info(f"[FSDP Val] Step {global_step}: {log_dict}") + if logger is not None and hasattr(logger, "log"): + logger.log(log_dict, step=global_step) + + val_info = f"[Validation] step {global_step}" + if "val_loss" in log_dict: + val_info += f" | val_loss {log_dict['val_loss']:.4f}" + if "val_cross_entropy_loss" in log_dict: + val_info += f" | val_ce_loss {log_dict['val_cross_entropy_loss']:.4f}" + if "val_flow_loss" in log_dict: + val_info += f" | val_flow_loss {log_dict['val_flow_loss']:.6f}" + if "val_action_accuracy" in log_dict: + val_info += f" | val_accuracy {log_dict['val_action_accuracy']:.4f}" + if tokenizer_mixin is not None: + for key in tokenizer_mixin.get_accuracy_keys(): + if key != "action_accuracy": + val_key = f"val_{key}" + if val_key in log_dict: + short_key = key.replace("action_accuracy_", "acc_") + val_info += f" | val_{short_key} {log_dict[val_key]:.4f}" + if self.logger is not None: + self.logger.info(val_info) + + model.train() + return log_dict + + def init_validation(self, normalizer_action, normalizer_propri): + """Optional pre-training-loop setup for validation. Default is a no-op.""" + del normalizer_action, normalizer_propri + + # ---- prediction / inference ---- + def predict( + self, + prediction_type: str, + *, + model, + val_dataloader, + rank, + world_size, + device, + processor, + tokenizer_mixin, + logger, + current_step, + max_iteration=None, + save_dir=None, + max_samples=None, + ): + """Run an inference loop for the given ``prediction_type``. + + Subclasses implement the dispatch table (flow_action / ar_action / + dllm_action / text / ...). Default raises — only VLAdapter has + meaningful predict paths today. + """ + raise NotImplementedError( + f"{type(self).__name__} does not support predict(prediction_type=...)" + ) + + # ---- optimizer-related defaults ---- + @property + def default_action_lr_keywords(self) -> list[str]: + """Keyword list used to identify action-expert parameters. + + The trainer splits parameters into two groups when + ``action_expert_learning_rate`` is set: names containing any of these + keywords go into the action-LR group. Subclasses override this when + the architecture uses different naming conventions. + """ + return [ + "action_preprocessor", + "moe.experts.1", + "qkv_proj_experts.1", + "o_proj_experts.1", + "input_layernorms.1", + "post_attention_layernorms.1", + "model.norms.1", + ] + + +def collect_channel_loss_stats(outputs, prefix="", tokenizer_mixin=None): + """Gather per-dataset channel losses + accuracy metrics via dist.all_reduce. + + Mirrors the legacy fsdp_trainer._collect_channel_loss_stats so training + and validation stats come out with identical keys. Returns {} when the + model output dict lacks ``channel_loss_dict``. + """ + stats = {} + channel_loss_dict = ( + outputs.get("channel_loss_dict") + if isinstance(outputs, dict) + else getattr(outputs, "channel_loss_dict", None) + ) + if channel_loss_dict is None: + return stats + channel_loss_count_dict = ( + outputs.get("channel_loss_count_dict") + if isinstance(outputs, dict) + else getattr(outputs, "channel_loss_count_dict", None) + ) + for dataset_name_i in channel_loss_dict: + count_tensor = channel_loss_count_dict[dataset_name_i].clone() + loss_tensor = channel_loss_dict[dataset_name_i].detach().clone() + dist.all_reduce(count_tensor, op=dist.ReduceOp.SUM) + dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM) + count_sum = count_tensor.item() + if count_sum >= 0.5: + stats[f"{prefix}channel_loss_{dataset_name_i}"] = ( + loss_tensor.item() / count_sum + ) + if "action_accuracy" in channel_loss_dict and tokenizer_mixin is not None: + acc_tensor = channel_loss_dict["action_accuracy"].detach().clone() + dist.all_reduce(acc_tensor, op=dist.ReduceOp.SUM) + world_size = dist.get_world_size() if dist.is_initialized() else 1 + stats[f"{prefix}action_accuracy"] = acc_tensor.item() / world_size + for key in tokenizer_mixin.get_accuracy_keys(): + if key != "action_accuracy" and key in channel_loss_dict: + rvq_acc_tensor = channel_loss_dict[key].detach().clone() + dist.all_reduce(rvq_acc_tensor, op=dist.ReduceOp.SUM) + stats[f"{prefix}{key}"] = rvq_acc_tensor.item() / world_size + return stats diff --git a/wall_x/trainer/adapters/vla_model_adapter.py b/wall_x/trainer/adapters/vla_model_adapter.py new file mode 100644 index 00000000..70c1037e --- /dev/null +++ b/wall_x/trainer/adapters/vla_model_adapter.py @@ -0,0 +1,990 @@ +import os + +import torch +import torch.distributed as dist +from tqdm import tqdm + +from wall_x.trainer.adapters.base_adapter import ModelAdapter +from wall_x.trainer.trainer_utils import ( + compute_action_metrics, + load_wallx_processors_from_cfg, + load_qwen_pretrain_weight, + save_text_results_to_file, +) +from wall_x.trainer.optimizer.dmuon import is_dmuon_model + + +class VLAdapter(ModelAdapter): + """Base adapter for public Qwen2.5 VLA models. + + Concrete subclasses define the model/config classes, FSDP wrap layers, + and optional task-specific logging hooks. + """ + + #: Subclass must set this. Identifies the variant in ``ADAPTER_REGISTRY``. + MODEL_TYPE: str = "" + + def __init__(self, *, cfg=None, logger=None, model_type=None): + # Concrete subclasses set MODEL_TYPE; tests may pass model_type explicitly + # (e.g. to swap a stub class into a different slot). + resolved = model_type or self.MODEL_TYPE + super().__init__(cfg=cfg, logger=logger, model_type=resolved) + if not resolved: + raise ValueError( + f"{type(self).__name__} requires a MODEL_TYPE class attribute " + "or an explicit model_type kwarg; got neither." + ) + self.model_type = resolved + self._dmuon_module = None + self._dmuon_import_checked = False + self._dmuon_named_param_cache = None + + def _get_dmuon_module(self): + if not self._dmuon_import_checked: + try: + import dmuon + except ImportError: + dmuon = None + self._dmuon_module = dmuon + self._dmuon_import_checked = True + return self._dmuon_module + + def _get_named_dmuon_dedicated_params(self, model): + cache = self._dmuon_named_param_cache + model_id = id(model) + if cache is not None and cache[0] == model_id: + return cache[1] + + dmuon = self._get_dmuon_module() + if dmuon is None: + named_params = [] + else: + module_to_name = { + id(module): module_name for module_name, module in model.named_modules() + } + named_params = [] + for dparam in dmuon.get_dedicated_params(model): + prefix = module_to_name.get(id(dparam.module), "") + name = f"{prefix}.{dparam.param_name}" if prefix else dparam.param_name + named_params.append((name, dparam)) + + self._dmuon_named_param_cache = (model_id, named_params) + return named_params + + # ---- variant hooks (subclass overrides) ---- + @classmethod + def model_class(cls): + """Return the training-side model class for this variant.""" + raise NotImplementedError(f"{cls.__name__} must override model_class()") + + @classmethod + def config_class(cls): + """Return the HF PretrainedConfig class for this variant.""" + raise NotImplementedError(f"{cls.__name__} must override config_class()") + + @classmethod + def inference_model_class(cls): + """Return the model class to use for inference.""" + return cls.model_class() + + # ---- helpers ---- + def _get_model_and_config_class(self): + """Resolve (ModelClass, ConfigClass) for this variant via classmethods.""" + cls = type(self) + return cls.model_class(), cls.config_class() + + # ---- processor ---- + def _build_processor_dict(self) -> dict: + """Flat-dict shape expected by legacy ``load_wallx_processors`` / + ``update_model_config`` / prediction loops. + + Derived from typed TrainConfig. Kept as the one-place adapter layer + between typed configs and legacy dict-based APIs; drop me when those + downstream APIs are typed-ified. + """ + import dataclasses + + flat = dataclasses.asdict(self.cfg.model) + flat["model_type"] = self.cfg.model_type + flat["data"] = dict(self.cfg._raw_data or {}) + 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: + flat["ar_dof_config"] = self.cfg.task.ar_dof_config + flat["batch_size_per_gpu"] = self.cfg.hyperparams.batch_size_per_gpu + return flat + + def load_processor(self, action_statistic_dof): + processors_dict = load_wallx_processors_from_cfg( + self.cfg, + normalizer=getattr(self, "normalizer_action", None), + action_statistic_dof=action_statistic_dof, + ) + self.logger.info( + f"processor vocab size: {len(processors_dict['processor'].tokenizer.vocab)}" + ) + self.logger.info( + f"num added tokens to processor: {processors_dict['num_added_tokens']}" + ) + return { + "processor": processors_dict["processor"], + "data_config": self._build_processor_dict(), + "tokenizer_mixin": processors_dict.get("tokenizer_mixin"), + "train_action_tokenizer": processors_dict["train_action_tokenizer"], + "val_action_tokenizer": processors_dict["val_action_tokenizer"], + "action_mapper": processors_dict["action_mapper"], + "num_added_tokens": processors_dict["num_added_tokens"], + } + + # ---- model ---- + def build_model_config(self): + _, ConfigClass = self._get_model_and_config_class() + qwen_vl_act_config_path = self.cfg.model.config_path + if qwen_vl_act_config_path.endswith(".json"): + model_config = ConfigClass.from_json_file(qwen_vl_act_config_path) + else: + model_config = ConfigClass.from_pretrained(qwen_vl_act_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}." + ) + + model_config.update_model_config(self._build_processor_dict()) + return model_config + + 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, + ) + + def load_weights(self, model, normalizer_action, normalizer_propri, **kwargs): + import copy + + processor = kwargs.get("processor") + + if self.cfg.model.pretrained_path: + model, err = load_qwen_pretrain_weight( + model, self.cfg.model.pretrained_path + ) + + if processor is not None: + model.resize_token_embeddings(len(processor.tokenizer)) + + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + elif ( + hasattr(model, "get_input_embeddings") + and model.get_input_embeddings() is not None + ): + + def _make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook( + _make_inputs_require_grad + ) + + if hasattr(model, "set_normalizer"): + model.set_normalizer( + copy.deepcopy(normalizer_action), + copy.deepcopy(normalizer_propri), + ) + + return model + + # ---- FSDP wrapping ---- + def get_transformer_layer_cls(self): + """Return the FSDP transformer-wrap layer classes for this variant. + + Subclass override — see the per-variant adapter for the actual + layer classes. + """ + raise NotImplementedError( + f"{type(self).__name__} must override get_transformer_layer_cls()" + ) + + # ---- dataset ---- + def load_dataset(self, data_config, processor, rank, world_size, **kwargs): + """Dispatch to the backend named by ``cfg.data.dataset_type``. + + The backend-specific wiring (resume indices, pool offsets, + processor wrapping) lives inside each backend's ``build()``; + this method only assembles a :class:`BuildContext` and forwards. + """ + import copy + + from wall_x.data import BuildContext, build_data + + resume_state = None + resume_batches = kwargs.get("resume_batches", 0) + indices = kwargs.get("resume_indices") + if indices is None and self.cfg.checkpoint.resume_from: + checkpoint_path = self.cfg.checkpoint.resume_from + if os.path.isfile(checkpoint_path): + checkpoint_path = os.path.dirname(checkpoint_path) + if os.path.isdir(checkpoint_path): + indices = self._load_episode_indices(checkpoint_path, rank) + + if resume_batches or indices or kwargs.get("episode_container_checkpoint_path"): + resume_state = { + "resume_batches": resume_batches, + "indices": indices, + "episode_container_checkpoint_path": kwargs.get( + "episode_container_checkpoint_path" + ), + } + + ctx = BuildContext( + rank=rank, + world_size=world_size, + processor=processor, + tokenizer_mixin=kwargs.get("tokenizer_mixin"), + normalizer_action=copy.deepcopy(kwargs.get("normalizer_action")), + normalizer_propri=copy.deepcopy(kwargs.get("normalizer_propri")), + model_config=kwargs.get("model_config"), + resume_state=resume_state, + ) + + bundle = build_data(self.cfg, ctx) + return bundle.dataset, bundle.train_loader, bundle.train_iters + + # ---- forward / loss ---- + def forward(self, model, batch, **kwargs): + self._last_batch_info = self._extract_batch_info(batch) + mode = kwargs.get("mode", "train") + return model(**batch, mode=mode) + + def extract_loss(self, outputs): + return outputs.loss + + def collect_output_stats( + self, + outputs, + step_stats, + reduce_tensor_fn, + true_gather_fn, + tokenizer_mixin=None, + ): + """Collect VL-specific auxiliary losses and accuracy metrics. + + Handles: + - cross_entropy_loss / flow_loss (all-reduce + true_gather) + - per-dataset channel losses (all-reduce via count-weighted avg) + - action_accuracy and extra RVQ layer accuracies + """ + import torch.distributed as dist + + cross_entropy_loss = self.get_output_field(outputs, "cross_entropy_loss") + if cross_entropy_loss is not None: + step_stats["cross_entropy_loss"] = reduce_tensor_fn( + cross_entropy_loss + ).item() + + flow_loss = self.get_output_field(outputs, "flow_loss") + if flow_loss is not None: + step_stats["flow_loss"] = reduce_tensor_fn(flow_loss).item() + + log_ce_loss = true_gather_fn( + self.get_output_field(outputs, "cross_entropy_loss") + ) + if log_ce_loss is not None: + step_stats["cross_entropy_loss"] = log_ce_loss + + log_flow_loss = true_gather_fn(self.get_output_field(outputs, "flow_loss")) + if log_flow_loss is not None: + step_stats["flow_loss"] = log_flow_loss + + channel_loss_dict = self.get_output_field(outputs, "channel_loss_dict") + channel_loss_count_dict = self.get_output_field( + outputs, "channel_loss_count_dict" + ) + if channel_loss_dict is not None: + for dataset_name_i in channel_loss_dict: + count_tensor = channel_loss_count_dict[dataset_name_i].clone() + loss_tensor = channel_loss_dict[dataset_name_i].detach().clone() + + dist.all_reduce(count_tensor, op=dist.ReduceOp.SUM) + dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM) + + cout_sum = count_tensor.item() + if cout_sum >= 0.5: + step_stats[f"channel_loss_{dataset_name_i}"] = ( + loss_tensor.item() / cout_sum + ) + + if "action_accuracy" in channel_loss_dict and tokenizer_mixin is not None: + acc_tensor = channel_loss_dict["action_accuracy"].detach().clone() + dist.all_reduce(acc_tensor, op=dist.ReduceOp.SUM) + world_size = dist.get_world_size() if dist.is_initialized() else 1 + step_stats["action_accuracy"] = acc_tensor.item() / world_size + + for key in tokenizer_mixin.get_accuracy_keys(): + if key != "action_accuracy" and key in channel_loss_dict: + rvq_acc_tensor = channel_loss_dict[key].detach().clone() + dist.all_reduce(rvq_acc_tensor, op=dist.ReduceOp.SUM) + step_stats[key] = rvq_acc_tensor.item() / world_size + + def collect_param_norms( + self, model, step_stats, device, reduce_tensor_fn, params_sharded=False + ): + """Compute per-group parameter L2 norms for VL models. + + Groups: + "visual" → visual_param_norm + "action" → action_expert_param_norm + others → org_vlm_param_norm + all → total_param_norm + """ + total_sq = torch.tensor(0.0, dtype=torch.float32, device=device) + visual_sq = torch.tensor(0.0, dtype=torch.float32, device=device) + action_sq = torch.tensor(0.0, dtype=torch.float32, device=device) + org_vlm_sq = torch.tensor(0.0, dtype=torch.float32, device=device) + dedicated_sq = torch.tensor(0.0, dtype=torch.float32, device=device) + dedicated_params = [] + + def add_named_sq(name, sq): + nonlocal total_sq, visual_sq, action_sq, org_vlm_sq + total_sq += sq + if "visual" in name: + visual_sq += sq + elif "action" in name: + action_sq += sq + else: + org_vlm_sq += sq + + with torch.no_grad(): + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + if param.numel() == 0 or hasattr(param, "_dedicated_owner_rank"): + continue + tensor = param.detach() + if hasattr(tensor, "to_local"): + tensor = tensor.to_local() + if tensor.numel() == 0: + continue + sq = torch.sum(tensor.float() ** 2) + add_named_sq(name, sq) + + if is_dmuon_model(model): + dedicated_params = self._get_named_dmuon_dedicated_params(model) + # DMuon replaces dedicated params with placeholders in + # named_parameters(). Count the authoritative dedicated + # storage here so total/category norms cover the full model. + for name, dparam in dedicated_params: + if getattr(dparam, "_dmuon_route", None) == "sharded_adamw": + replicate_group = getattr(dparam, "replicate_group", None) + # replicate_group=None is 1D shard-only mode: each rank + # contributes a distinct shard and the later reduce sums + # them into the global norm. + if ( + replicate_group is not None + and replicate_group.rank() + != getattr(dparam, "owner_replicate", 0) + ): + continue + tensor = getattr(dparam, "_sharded_adamw_data", None) + if tensor is None: + continue + valid_numel = int( + getattr( + dparam, + "_sharded_adamw_valid_numel", + tensor.numel(), + ) + ) + tensor = tensor[:valid_numel] + else: + if not bool(getattr(dparam, "is_owner", False)): + continue + tensor = getattr(dparam, "_owned_data", None) + if tensor is None: + continue + + if tensor.numel() == 0: + continue + sq = torch.sum(tensor.detach().float() ** 2) + dedicated_sq += sq + add_named_sq(name, sq) + + if params_sharded: + total_sq = reduce_tensor_fn(total_sq, average=False) + visual_sq = reduce_tensor_fn(visual_sq, average=False) + action_sq = reduce_tensor_fn(action_sq, average=False) + org_vlm_sq = reduce_tensor_fn(org_vlm_sq, average=False) + dedicated_sq = reduce_tensor_fn(dedicated_sq, average=False) + + step_stats["total_param_norm"] = torch.sqrt(total_sq).item() + step_stats["visual_param_norm"] = torch.sqrt(visual_sq).item() + step_stats["action_expert_param_norm"] = torch.sqrt(action_sq).item() + step_stats["org_vlm_param_norm"] = torch.sqrt(org_vlm_sq).item() + if dedicated_params: + step_stats["dmuon_dedicated_param_norm"] = torch.sqrt(dedicated_sq).item() + + # ---- MFU computation ---- + @staticmethod + def _extract_batch_info(batch): + """Extract token counts from a training batch.""" + info = {} + input_ids = batch.get("input_ids") + if input_ids is not None: + info["batch_size"] = input_ids.shape[0] + info["seq_length"] = input_ids.shape[1] + + moe_token_types = batch.get("moe_token_types") + if moe_token_types is not None: + info["num_lang_tokens"] = ( + int((moe_token_types == 0).sum().item()) // info["batch_size"] + ) + info["num_action_tokens"] = ( + int((moe_token_types == 1).sum().item()) // info["batch_size"] + ) + else: + info["num_lang_tokens"] = info.get("seq_length", 0) + info["num_action_tokens"] = 0 + + pixel_values = batch.get("pixel_values") + info["vision_seq_length"] = ( + pixel_values.shape[0] if pixel_values is not None else 0 + ) + + labels = batch.get("labels") + if labels is not None: + info["num_loss_tokens"] = int((labels[..., 1:] != -100).sum().item()) + else: + info["num_loss_tokens"] = None + + return info + + def _compute_detailed_flops(self, model, config): + """Compute detailed per-module FLOPs for one training step. + + Uses the same formulas as scripts/profile_forward.py:compute_module_flops, + with fwd+bwd multipliers: + - Frozen modules (e.g. ViT if frozen): 1x forward + - Trainable modules: 3x forward (fwd + 2x bwd) + """ + + batch_info = getattr(self, "_last_batch_info", None) + if not batch_info: + return None + + model_config = model.config if hasattr(model, "config") else None + if model_config is None: + return None + + B = batch_info["batch_size"] + S = batch_info["seq_length"] + num_lang = batch_info["num_lang_tokens"] + num_act = batch_info["num_action_tokens"] + N_lang = B * num_lang + N_act = B * num_act + N_total = N_lang + N_act + Nv = batch_info["vision_seq_length"] + num_loss_tokens = batch_info["num_loss_tokens"] + + H = model_config.hidden_size + num_layers = model_config.num_hidden_layers + num_heads = model_config.num_attention_heads + num_kv = model_config.num_key_value_heads + vocab_size = getattr(model_config, "padded_vocab_size", model_config.vocab_size) + + use_mot = getattr(model_config, "attention_moe", False) + use_moe_mlp = getattr(model_config, "mlp_moe", False) + dim_inputs = getattr(model_config, "dim_inputs", (H, H)) + dim_lang, dim_act = dim_inputs + + grad_accum = self.cfg.hyperparams.gradient_accumulation_steps + + vit_frozen = True + if hasattr(model, "visual"): + for p in model.visual.parameters(): + if p.requires_grad: + vit_frozen = False + break + vit_mult = 1.0 if vit_frozen else 3.0 + + vit_fwd_flops = 0 + if hasattr(model_config, "vision_config") and Nv > 0: + vcfg = model_config.vision_config + Hv = vcfg.hidden_size + Iv = vcfg.intermediate_size + out_hidden = vcfg.out_hidden_size + depth_v = vcfg.depth + + F_qkv_v = 6 * Nv * Hv * Hv + F_o_v = 2 * Nv * Hv * Hv + F_mlp_v = 6 * Nv * Hv * Iv + 2 * Nv * Iv + F_linear_per_layer = F_qkv_v + F_o_v + F_mlp_v + + fullatt_set = set(getattr(vcfg, "fullatt_block_indexes", [])) + num_images = max(Nv // 768, 1) + si = Nv // num_images + sum_si_sq = num_images * si * si + sum_wi_sq = sum_si_sq // 16 + + for i in range(depth_v): + if i in fullatt_set: + vit_fwd_flops += F_linear_per_layer + 4 * Hv * sum_si_sq + else: + vit_fwd_flops += F_linear_per_layer + 4 * Hv * sum_wi_sq + + merge_unit = getattr(vcfg, "spatial_merge_size", 2) ** 2 + merger_hidden = Hv * merge_unit + Nv_merged = Nv // merge_unit + vit_fwd_flops += ( + 2 * Nv_merged * merger_hidden * merger_hidden + + 2 * Nv_merged * merger_hidden * out_hidden + ) + + kv_ratio = num_kv / num_heads + F_matmul_fwd = 4 * B * (S**2) * H + + if not use_mot: + F_attn_fwd = ( + 2 * N_total * H * H + + 4 * N_total * H * H * kv_ratio + + 2 * N_total * H * H + + F_matmul_fwd + ) + else: + F_attn_fwd = ( + N_lang * dim_lang * H * (4 + 4 * kv_ratio) + + N_act * dim_act * H * (4 + 4 * kv_ratio) + + F_matmul_fwd + ) + + if not use_moe_mlp: + ffn_hidden = model_config.intermediate_size + F_mlp_fwd = 6 * N_total * H * ffn_hidden + 2 * N_total * ffn_hidden + else: + hid_lang = model_config.experts[0]["intermediate_size"] + hid_act = model_config.experts[1]["intermediate_size"] + F_mlp_fwd = (6 * N_lang * dim_lang * hid_lang + 2 * N_lang * hid_lang) + ( + 6 * N_act * dim_act * hid_act + 2 * N_act * hid_act + ) + + decoder_fwd_flops = num_layers * (F_attn_fwd + F_mlp_fwd) + + N_lm = num_loss_tokens if num_loss_tokens is not None else N_total + lm_head_fwd_flops = 2 * N_lm * H * vocab_size + + total_flops = ( + vit_mult * vit_fwd_flops + 3.0 * decoder_fwd_flops + 3.0 * lm_head_fwd_flops + ) * grad_accum + + return { + "total_flops": total_flops, + "vit_fwd_flops": vit_fwd_flops, + "decoder_fwd_flops": decoder_fwd_flops, + "lm_head_fwd_flops": lm_head_fwd_flops, + "vit_mult": vit_mult, + } + + def compute_mfu(self, model, step_time_seconds): + """Compute detailed Model FLOPs Utilization for VL transformer training.""" + if step_time_seconds <= 0: + return None + try: + config = self._build_processor_dict() + flops_info = self._compute_detailed_flops(model, config) + if flops_info is None: + return None + + gpu_peak_tflops = 312.0 # TODO: move to DebugConfig if needed + peak_flops = gpu_peak_tflops * 1e12 + + # Data-parallel: each GPU independently processes its own micro-batch, + # so per-GPU MFU = per_gpu_flops / (step_time * per_gpu_peak). + # No division by num_gpus needed. + total_flops = flops_info["total_flops"] + mfu = total_flops / (step_time_seconds * peak_flops) + + return { + "mfu": mfu, + "flops_per_step_T": total_flops / 1e12, + } + except Exception: + return None + + # ---- optional hooks ---- + @staticmethod + def log_attention_implementation(logger, model): + """Log the attention implementation name. Variant-specific because the + model layout (model.model vs model.model.language_model) differs.""" + raise NotImplementedError( + "VLAdapter subclasses must override log_attention_implementation()" + ) + + def _load_episode_indices(self, checkpoint_path: str, rank: int): + """Read backend-managed per-rank resume offsets if supported.""" + from wall_x.data import data_backend + + backend = data_backend() + if not backend.supports("load_episode_indices"): + return None + return backend.load_episode_indices(checkpoint_path, rank) + + # ---- prediction / inference ---- + def predict( + self, + prediction_type: str, + *, + model, + val_dataloader, + rank, + world_size, + device, + processor, + tokenizer_mixin, + logger, + current_step, + max_iteration=None, + save_dir=None, + max_samples=None, + ): + """Dispatch VLA inference by prediction_type. + + flow_action / dllm_action return per-rank L1 action metrics; ar_action + is reserved (was unimplemented pre-refactor too); text runs text + generation with optional point-L1 distance scoring. + """ + config = self._build_processor_dict() + if prediction_type == "flow_action": + return self._predict_flow_action( + model=model, + val_dataloader=val_dataloader, + rank=rank, + world_size=world_size, + device=device, + config=config, + logger=logger, + current_step=current_step, + max_iteration=max_iteration, + ) + if prediction_type == "ar_action": + return None # matches pre-refactor behaviour (was pass) + if prediction_type == "dllm_action": + return self._predict_dllm_action( + model=model, + val_dataloader=val_dataloader, + rank=rank, + world_size=world_size, + device=device, + config=config, + processor=processor, + tokenizer_mixin=tokenizer_mixin, + logger=logger, + current_step=current_step, + max_iteration=max_iteration, + ) + if prediction_type == "text": + return self._predict_text( + model=model, + val_dataloader=val_dataloader, + rank=rank, + device=device, + logger=logger, + current_step=current_step, + max_samples=max_samples, + save_dir=save_dir, + ) + raise ValueError(f"Unsupported prediction type: {prediction_type}") + + @torch.no_grad() + def _predict_flow_action( + self, + *, + model, + val_dataloader, + rank, + world_size, + device, + config, + logger, + current_step, + max_iteration, + ): + if dist.is_initialized(): + dist.barrier() + + total_num = len(val_dataloader) + if max_iteration: + total_num = min(max_iteration, total_num) + model.eval() + + all_preds, all_actions = [], [] + pepoch = tqdm( + total=total_num, + desc=f"Predicting ckpt at step {current_step}", + disable=rank != 0, + ) + for batch_idx, batch in enumerate(val_dataloader): + if batch_idx >= total_num: + break + batch = _move_batch(batch, device) + + model_output = model.generate_flow_action( + action_horizon=config["data"]["action_horizon_flow"], + action_dim=model.action_preprocessor.action_dim, + **batch, + ) + pred_action, gt_action = ( + model_output["predict_action"], + model_output["gt_action"], + ) + + pred_list = [torch.zeros_like(pred_action) for _ in range(world_size)] + gt_list = [torch.zeros_like(gt_action) for _ in range(world_size)] + dist.all_gather(pred_list, pred_action) + dist.all_gather(gt_list, gt_action) + + if dist.is_initialized(): + dist.barrier() + if rank == 0: + all_preds.append(torch.cat(pred_list, dim=0).cpu()) + all_actions.append(torch.cat(gt_list, dim=0).cpu()) + pepoch.update(1) + + pepoch.close() + + if rank == 0: + step_log = {} + if all_preds: + all_preds = torch.cat(all_preds, dim=0) + all_actions = torch.cat(all_actions, dim=0) + step_log = compute_action_metrics( + all_preds, all_actions, config=config, step_log=step_log + ) + if logger and step_log: + logger.log(step_log, step=current_step) + if self.logger is not None: + self.logger.info( + f"Step {current_step}, Validation L1 Loss: {step_log.get('val_action_l1', -1)}" + ) + return step_log + return None + + @torch.no_grad() + def _predict_dllm_action( + self, + *, + model, + val_dataloader, + rank, + world_size, + device, + config, + processor, + tokenizer_mixin, + logger, + current_step, + max_iteration, + ): + if dist.is_initialized(): + dist.barrier() + + total_num = len(val_dataloader) + if max_iteration: + total_num = min(max_iteration, total_num) + model.eval() + + all_preds, all_actions = [], [] + pepoch = tqdm( + total=total_num, + desc=f"Predicting ckpt at step {current_step}", + disable=rank != 0, + ) + for batch_idx, batch in enumerate(val_dataloader): + if batch_idx >= total_num: + break + batch = _move_batch(batch, device) + batch = self._preprocess_dllm_batch(batch, processor, tokenizer_mixin) + + if hasattr(model.action_tokenizer, "max_waypoints"): + total_ar_step = model.action_tokenizer.max_waypoints + elif hasattr(model.action_tokenizer, "max_length"): + total_ar_step = model.action_tokenizer.max_length + else: + raise ValueError( + "Unknown action_tokenizer type for dllm action inference" + ) + model_output = model.generate_dllm_action( + action_dim=7, + action_horizon=config["data"]["action_horizon_flow"], + use_ar_action=False, + total_ar_step=total_ar_step, + **batch, + ) + pred_action, gt_action = ( + model_output["predict_action"], + model_output["gt_action"], + ) + + pred_list = [torch.zeros_like(pred_action) for _ in range(world_size)] + gt_list = [torch.zeros_like(gt_action) for _ in range(world_size)] + dist.all_gather(pred_list, pred_action) + dist.all_gather(gt_list, gt_action) + + if dist.is_initialized(): + dist.barrier() + if rank == 0: + all_preds.append(torch.cat(pred_list, dim=0).cpu()) + all_actions.append(torch.cat(gt_list, dim=0).cpu()) + pepoch.update(1) + + pepoch.close() + + if rank == 0: + step_log = {} + if all_preds: + all_preds = torch.cat(all_preds, dim=0) + all_actions = torch.cat(all_actions, dim=0) + step_log = compute_action_metrics( + all_preds, all_actions, config=config, step_log=step_log + ) + if logger and step_log: + logger.log(step_log, step=current_step) + if self.logger is not None: + self.logger.info( + f"Step {current_step}, Validation L1 Loss: {step_log.get('val_action_l1', -1)}" + ) + return step_log + return None + + @staticmethod + def _preprocess_dllm_batch(batch, processor, tokenizer_mixin): + """Patch the dllm action placeholders into input_ids (in place).""" + input_ids = batch["input_ids"] + prefix_length = batch["prefix_length"] + placeholder_str = tokenizer_mixin.get_placeholder_for_dllm() + placeholder_seq = torch.tensor( + processor.tokenizer.convert_tokens_to_ids(placeholder_str) + ) + placeholder_len = placeholder_seq.shape[0] + input_ids[:, prefix_length - placeholder_len - 2 : prefix_length - 2] = ( + placeholder_seq + ) + batch.update({"input_ids": input_ids}) + return batch + + @torch.no_grad() + def _predict_text( + self, + *, + model, + val_dataloader, + rank, + device, + logger, + current_step, + max_samples, + save_dir, + ): + from wall_x._vendor.x2robot_utils.grounding import calculate_point_l1_distance + + if dist.is_initialized(): + dist.barrier() + + total_num = ( + len(val_dataloader) + if max_samples is None + else min(max_samples, len(val_dataloader)) + ) + model.eval() + + all_input_texts, all_gt_texts, all_pred_texts = [], [], [] + all_point_l1_distances = [] + + pepoch = tqdm( + total=total_num, + desc=f"Predicting text at step {current_step}", + disable=rank != 0, + ) + + for batch_idx, batch in enumerate(val_dataloader): + if batch_idx >= total_num: + break + batch = _move_batch(batch, device) + + model_output = model.generate_text( + input_ids=batch.get("input_ids"), + attention_mask=batch.get("attention_mask"), + moe_token_types=batch.get("moe_token_types"), + pixel_values=batch.get("pixel_values"), + image_grid_thw=batch.get("image_grid_thw"), + proprioception=batch.get("proprioception"), + dataset_names=batch.get("dataset_names"), + dof_mask=batch.get("dof_mask"), + agent_pos_mask=batch.get("agent_pos_mask"), + prefix_length=batch.get("prefix_length"), + positional_masks=batch.get("positional_masks"), + re_generate=False, + ) + + input_texts = model_output["input_text"] + gt_texts = model_output["gt_output_text"] + pred_texts = model_output["predict_output_text"] + + if rank == 0: + all_input_texts.extend(input_texts) + all_gt_texts.extend(gt_texts) + all_pred_texts.extend(pred_texts) + for gt_text, pred_text in zip(gt_texts, pred_texts): + gt_clean = gt_text[0] if isinstance(gt_text, list) else gt_text + pred_clean = ( + pred_text[0] if isinstance(pred_text, list) else pred_text + ) + point_l1_dist = calculate_point_l1_distance(gt_clean, pred_clean) + if point_l1_dist is not None: + all_point_l1_distances.append(point_l1_dist) + + pepoch.update(1) + + pepoch.close() + + if rank == 0 and all_pred_texts: + if all_point_l1_distances: + avg_point_l1 = sum(all_point_l1_distances) / len(all_point_l1_distances) + min_point_l1 = min(all_point_l1_distances) + max_point_l1 = max(all_point_l1_distances) + point_stats = { + "text_prediction_point_l1_avg": avg_point_l1, + "text_prediction_point_l1_min": min_point_l1, + "text_prediction_point_l1_max": max_point_l1, + "text_prediction_point_samples_count": len(all_point_l1_distances), + } + if logger is not None: + logger.log(point_stats, step=current_step) + if self.logger is not None: + self.logger.info( + f"Point L1 Statistics - Avg: {avg_point_l1:.4f}, " + f"Min: {min_point_l1:.4f}, Max: {max_point_l1:.4f}, " + f"Samples: {len(all_point_l1_distances)}" + ) + save_text_results_to_file( + current_step, + all_input_texts, + all_gt_texts, + all_pred_texts, + save_dir if save_dir is not None else "./", + ) + + +def _move_batch(batch, device): + """Recursively move dict/list/tensor to device (VLA-local copy to keep + VLAdapter self-contained; base_adapter has a similar helper used by + validation).""" + if isinstance(batch, dict): + return {k: _move_batch(v, device) for k, v in batch.items()} + if isinstance(batch, list): + return [_move_batch(v, device) for v in batch] + if torch.is_tensor(batch): + return batch.to(device) + return batch diff --git a/wall_x/trainer/fsdp_trainer/__init__.py b/wall_x/trainer/fsdp_trainer/__init__.py new file mode 100644 index 00000000..4bd20fc4 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/__init__.py @@ -0,0 +1,36 @@ +"""Distributed trainer package exports.""" + +from wall_x.trainer.fsdp_trainer.base_trainer import ( + BaseDistributedTrainer, + DistributedTrainer, + setup_distributed, + cleanup_distributed, + get_rank, + get_world_size, + get_local_rank, + is_main_process, + barrier, + all_reduce, + all_gather, +) + +from wall_x.trainer.fsdp_trainer.fsdp_trainer import FSDPTrainer + + +__all__ = [ + # Base classes and utilities + "BaseDistributedTrainer", + "DistributedTrainer", + "setup_distributed", + "cleanup_distributed", + "get_rank", + "get_world_size", + "get_local_rank", + "is_main_process", + "barrier", + "all_reduce", + "all_gather", + # Trainers + # "DDPTrainer", + "FSDPTrainer", +] diff --git a/wall_x/trainer/fsdp_trainer/base_trainer.py b/wall_x/trainer/fsdp_trainer/base_trainer.py new file mode 100644 index 00000000..8b8da59a --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/base_trainer.py @@ -0,0 +1,429 @@ +""" +Base Trainer for Pure PyTorch Distributed Training (DDP/FSDP) +No accelerate dependency - can be launched directly with torchrun +""" + +import os +import gc +import torch +import logging +import contextlib +import torch.distributed as dist +from abc import ABC, abstractmethod +from datetime import timedelta +from typing import Optional, Any + +from wall_x.config.train_config import TrainConfig +from wall_x.config.hyperparams_config import AdamWConfig +from wall_x.utils.logger import DistributedLogger +from wall_x.trainer.trainer_utils import seed_all +from wall_x.utils.timers import Timers + + +def setup_distributed(): + """Initialize distributed training environment. + + NCCL watchdog default is 10 min — some dataset backends (e.g. lerobot's + per-rank sequential LeRobotDataset indexing) can exceed that during the + first build, causing barrier #2 to time out on the waiting rank. Allow + the timeout to be raised via ``WALLX_DIST_TIMEOUT_MINUTES`` (default 30). + """ + if not dist.is_initialized(): + timeout_min = int(os.environ.get("WALLX_DIST_TIMEOUT_MINUTES", "30")) + dist.init_process_group( + backend="nccl", + timeout=timedelta(minutes=timeout_min), + ) + + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + + return local_rank + + +def cleanup_distributed(): + """Clean up distributed training environment""" + if dist.is_initialized(): + dist.destroy_process_group() + + +def get_rank() -> int: + """Get current process rank""" + if dist.is_initialized(): + return dist.get_rank() + return 0 + + +def get_world_size() -> int: + """Get total number of processes""" + if dist.is_initialized(): + return dist.get_world_size() + return 1 + + +def get_local_rank() -> int: + """Get local rank (GPU index on current node)""" + return int(os.environ.get("LOCAL_RANK", 0)) + + +def is_main_process() -> bool: + """Check if current process is the main process (rank 0)""" + return get_rank() == 0 + + +# Trainer-level collectives should use an explicit global process group. +# FSDP's mesh groups may be 1D or 2D and have sharding-specific semantics, +# while trainer barriers/metric reductions need all-rank semantics. +_trainer_process_group = None + + +def set_trainer_process_group(pg): + global _trainer_process_group + _trainer_process_group = pg + + +def barrier(): + """Synchronize all processes""" + if dist.is_initialized(): + dist.barrier(group=_trainer_process_group) + + +def all_reduce(tensor: torch.Tensor, op=dist.ReduceOp.SUM) -> torch.Tensor: + """All-reduce tensor across all processes""" + if dist.is_initialized(): + dist.all_reduce(tensor, op=op, group=_trainer_process_group) + return tensor + + +def all_gather(tensor: torch.Tensor) -> torch.Tensor: + """All-gather tensor from all processes""" + if not dist.is_initialized(): + return tensor + + group = _trainer_process_group + world_size = dist.get_world_size(group=group) + gathered = [torch.zeros_like(tensor) for _ in range(world_size)] + dist.all_gather(gathered, tensor, group=group) + return torch.stack(gathered) + + +class BaseDistributedTrainer(ABC): + """ + Abstract base class for distributed training without accelerate. + Supports both DDP and FSDP through subclasses. + + Launch with torchrun: + torchrun --nproc_per_node=8 train.py --config config.yaml + """ + + @abstractmethod + def train_loop(self, epoch: int): + """Training loop for one epoch""" + raise NotImplementedError("train_loop must be implemented") + + @abstractmethod + def val_loop(self): + """Validation loop""" + raise NotImplementedError("val_loop must be implemented") + + @abstractmethod + def save_checkpoint(self, epoch: int, step: int = 0): + """Save model checkpoint""" + raise NotImplementedError("save_checkpoint must be implemented") + + @abstractmethod + def load_model(self): + """Load model""" + raise NotImplementedError("load_model must be implemented") + + @abstractmethod + def load_dataset(self): + """Load dataset""" + raise NotImplementedError("load_dataset must be implemented") + + @abstractmethod + def backward(self, loss: torch.Tensor): + """Perform backward pass""" + raise NotImplementedError("backward must be implemented") + + @abstractmethod + def clip_grad_norm(self, max_norm: float) -> torch.Tensor: + """Clip gradient norm""" + raise NotImplementedError("clip_grad_norm must be implemented") + + +class DistributedTrainer(BaseDistributedTrainer): + """ + Concrete base trainer with common functionality for DDP/FSDP. + Subclasses should implement wrap_model, backward, clip_grad_norm, etc. + """ + + def __init__( + self, + train_config: TrainConfig, + wandb_run: Optional[Any] = None, + ): + self.cfg = train_config + + # wandb Run object (metric recorder) — None on non-main ranks and + # when use_wandb is disabled. Distinct from self.logger below. + self.wandb_run = wandb_run + + # Initialize distributed environment + self.local_rank = setup_distributed() + self.rank = get_rank() + self.world_size = get_world_size() + self.device = torch.device(f"cuda:{self.local_rank}") + + # Dataset config: optional backend-provided trainer config. + from wall_x.trainer.adapters.base_adapter import load_trainer_data_config + + self.data_config = load_trainer_data_config(self.cfg) + + # Text logger + self.logger = DistributedLogger( + name=self.__class__.__name__, + save_path=self.cfg.checkpoint.save_path, + ) + + # Training state + self.seed = self.cfg.hyperparams.seed + self.logger.info(f"seed {self.seed}") + + seed_all(self.seed) + self.start_epoch = 0 + self.global_step = 0 + self.micro_step = 0 + self.num_epoch = self.cfg.hyperparams.num_epoch + self.dataset_config_path = getattr(self.cfg.data, "dataset_config_path", None) + self.initial_step = 0 + + # Optimizer hyperparameters (betas / weight_decay / eps) + # used by build_optimizer for both AdamW and Muon paths. + opt = self.cfg.hyperparams.optimizer + if isinstance(opt, AdamWConfig): + self.adamw_betas = tuple(opt.betas) + self.adamw_weight_decay = opt.weight_decay + self.adamw_eps = opt.eps + elif getattr(opt, "optimizer_type", None) == "muon": + self.adamw_betas = (opt.beta_1, opt.beta_2) + self.adamw_weight_decay = opt.weight_decay + self.adamw_eps = opt.eps + else: + self.adamw_betas = (0.9, 0.98) + self.adamw_weight_decay = 1e-8 + self.adamw_eps = 1e-8 + + # Performance options + self.nvtx = self.cfg.debug.nvtx + self.timers = Timers(log_level=0, log_option="minmax") + + # Training hyperparameters + self.max_grad_norm = opt.max_grad_norm + self.grad_accum_steps = self.cfg.hyperparams.gradient_accumulation_steps + self.log_interval = self.cfg.logging.log_interval + self.log_stats_buffer = [] + + # Model and optimizer (to be set by subclass) + self.model = None + self.optimizer = None + self.lr_scheduler = None + + # Mixed precision + dist_cfg = self.cfg.distributed + self.use_amp = dist_cfg.use_amp + self.amp_dtype = torch.bfloat16 if dist_cfg.bf16 else torch.float16 + self.grad_scaler = None + if self.use_amp and self.amp_dtype == torch.float16: + self.grad_scaler = torch.amp.GradScaler("cuda") + + def log( + self, message: str, level: int = logging.INFO, main_process_only: bool = True + ): + """Log a text message via self.logger. + + Kept for backward compat with call sites that still use ``self.log(...)`` + (new code should call ``self.logger.info/warning/error(...)`` directly). + Under the new DistributedLogger the ``main_process_only`` flag is + redundant for stdout (stdout is rank-0-only by construction); it still + controls whether non-zero ranks write the message to their file log. + """ + if main_process_only and not is_main_process(): + return + self.logger.log(message, level=level, main_process_only=main_process_only) + + def fit(self): + """Main training loop""" + barrier() + + # Optional: validate before training + if self.cfg.checkpoint.validate_first: + self.val_loop() + barrier() + + if self.nvtx: + torch.cuda.cudart().cudaProfilerStart() + + # Profiler setup + enable_profiling = self.cfg.debug.profile + profiler = ( + self._setup_profiler() if enable_profiling else contextlib.nullcontext() + ) + + for epoch in range(self.start_epoch, self.num_epoch): + # Training + self.train_loop(epoch, profiler=profiler) + barrier() + num_training_steps = getattr(self, "num_training_steps", 0) + if num_training_steps > 0 and self.global_step >= num_training_steps: + break + + # Save checkpoint + if (epoch + 1) % self.cfg.logging.epoch_save_interval == 0: + self.save_checkpoint(epoch) + + # Validation + self.val_loop() + barrier() + + gc.collect() + torch.cuda.empty_cache() + + if self.nvtx: + torch.cuda.cudart().cudaProfilerStop() + + def _setup_profiler(self): + """Setup PyTorch profiler""" + return torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + schedule=torch.profiler.schedule( + wait=self.cfg.debug.profile_wait_iters, + warmup=self.cfg.debug.profile_warmup_iters, + active=self.cfg.debug.profile_active_iters, + ), + on_trace_ready=torch.profiler.tensorboard_trace_handler( + self.cfg.debug.profile_save_path, + worker_name=f"worker{self.rank}", + ), + record_shapes=True, + profile_memory=True, + with_stack=True, + ) + + def gather_tensor(self, tensor: torch.Tensor) -> torch.Tensor: + """Gather tensor from all processes""" + return all_gather(tensor) + + def reduce_tensor(self, tensor: torch.Tensor, average: bool = True) -> torch.Tensor: + """Reduce tensor across all processes""" + tensor = tensor.clone() + all_reduce(tensor) + if average: + tensor = tensor / self.world_size + return tensor + + def sync_gradients(self) -> bool: + """Check if gradients should be synchronized (for gradient accumulation)""" + return (self.micro_step + 1) % self.grad_accum_steps == 0 + + def optimizer_zero_grad(self): + """Zero optimizer gradients""" + self.optimizer.zero_grad(set_to_none=True) + + def optimizer_step(self): + """Perform optimizer step""" + if self.grad_scaler is not None: + self.grad_scaler.step(self.optimizer) + self.grad_scaler.update() + else: + self.optimizer.step() + + def lr_scheduler_step(self): + """Perform learning rate scheduler step""" + if self.lr_scheduler is not None: + self.lr_scheduler.step() + + def get_lr(self) -> float: + """Get current learning rate""" + if self.lr_scheduler is not None: + return self.lr_scheduler.get_last_lr()[0] + return self.optimizer.param_groups[0]["lr"] + + def autocast_context(self): + """Get autocast context for mixed precision""" + if self.use_amp: + return torch.amp.autocast("cuda", dtype=self.amp_dtype) + return contextlib.nullcontext() + + def training_log( + self, + current_epoch: int, + total_epoch: int, + current_iter: int, + total_iter: int, + loss: torch.Tensor, + lr: float, + time_per_step: float, + show_time_details: bool = False, + ): + """Log training progress""" + if not is_main_process(): + return + + # Prefer the cross-rank reduced loss (train_loss) from + # _current_step_stats so the console line matches what wall_x_2509's + # fsdp_trainer prints and what wandb sees. Falls back to the raw + # rank-0 local tensor when stats haven't been populated yet. + _smoothed = getattr(self, "_current_step_stats", None) or {} + _loss_to_print = _smoothed.get("train_loss", float(loss)) + + log_string = "" + log_string += f" epoch {current_epoch:3d}/{total_epoch:3d} |" + log_string += f" iter {current_iter:6d}/{total_iter:6d} |" + log_string += f" loss {_loss_to_print:.10f} |" + log_string += f" lr {lr:.6f} |" + log_string += f" time {time_per_step:.4f}s |" + + self.log(log_string) + + if show_time_details: + timers_to_log = [ + "interval-time", + "data-load", + "forward-compute", + "backward-compute", + "optimizer", + ] + self.timers.log(timers_to_log, normalizer=1) + + def true_gather(self, value): + """Gather values across all processes and compute mean over non-None values.""" + device = next(self.model.parameters()).device + + # CRITICAL: must clone() here, not just detach(). + # detach() only disconnects the computation graph but still shares the + # same underlying storage. dist.all_reduce() is **in-place** — it + # overwrites that storage with the sum across all ranks. If the caller + # passes a tensor that is also referenced elsewhere (e.g. outputs["loss"] + # and outputs["flow_loss"] pointing to the same scalar_loss), the + # in-place all_reduce silently mutates the original, causing downstream + # readers (like training_log) to see a value multiplied by world_size. + value_to_gather = ( + value.detach().clone() + if value is not None + else torch.tensor(0.0, device=device) + ) + count_to_gather = torch.tensor(1.0 if value is not None else 0.0, device=device) + + # all_reduce is cheaper than all_gather because only the sum is needed. + dist.all_reduce(value_to_gather, op=dist.ReduceOp.SUM) + dist.all_reduce(count_to_gather, op=dist.ReduceOp.SUM) + + total_processes = count_to_gather.item() + if total_processes > 0: + return value_to_gather.item() / total_processes + else: + return None diff --git a/wall_x/trainer/fsdp_trainer/checkpoint_io.py b/wall_x/trainer/fsdp_trainer/checkpoint_io.py new file mode 100644 index 00000000..bb21b6a3 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/checkpoint_io.py @@ -0,0 +1,1027 @@ +"""Checkpoint save/load helpers for distributed training.""" + +from __future__ import annotations + +import gc +import logging +import os +import random +import shutil +import time +from typing import Callable, Optional, Tuple + +import numpy as np +import torch +import torch.distributed as dist +import yaml +from safetensors.torch import load_file, save_file +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_state_dict, + set_model_state_dict, + set_optimizer_state_dict, +) +from torch.nn.parallel import DistributedDataParallel as DDP + +from wall_x.trainer.optimizer.dmuon import is_dmuon_model + + +def _noop_log(_msg: str, **_kw) -> None: + pass + + +def _dict_section(config: dict, key: str) -> dict: + section = config.get(key, {}) + return section if isinstance(section, dict) else {} + + +# ---------------------------------------------------------------------- +# Detectors +# ---------------------------------------------------------------------- + + +def _is_fsdp2_model(model: torch.nn.Module) -> bool: + """Detect FSDP2 models by the presence of DTensor parameters. + + FSDP2's ``fully_shard`` converts parameters in-place to DTensors without + wrapping the module in an outer class. DDP wraps in + ``DistributedDataParallel``; unwrapped models have plain + ``torch.Tensor`` parameters. + """ + if isinstance(model, DDP): + return False + try: + from torch.distributed.tensor import DTensor + except ImportError: + return False + for p in model.parameters(): + if isinstance(p, DTensor): + return True + return False + + +def _detect_legacy_fsdp1_format(checkpoint_path: str) -> bool: + """True when the ckpt directory looks like a pre-migration FSDP1 save. + + Legacy: ``model.safetensors`` is present (rank-0 full state), but + optimizer state lives in per-rank ``optimizer_rank{N}.pt`` files + instead of a single ``optimizer.pt``. After the FSDP1 → FSDP2 + migration we cannot reshard those flat_param-keyed optim files into + the new DTensor layout, so the legacy loader cold-starts the + optimizer and warns. + """ + if not os.path.isdir(checkpoint_path): + return False + if os.path.exists(os.path.join(checkpoint_path, "optimizer.pt")): + return False + try: + return any( + f.startswith("optimizer_rank") and f.endswith(".pt") + for f in os.listdir(checkpoint_path) + ) + except OSError: + return False + + +# ---------------------------------------------------------------------- +# Save +# ---------------------------------------------------------------------- + + +def save_checkpoint( + *, + ckpt_path: str, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + lr_scheduler, + config: dict, + rank: int, + is_main: bool, + epoch: int, + global_step: int, + seed: int, + normalizer_action, + normalizer_propri, + dataset=None, + grad_scaler=None, + log_fn: Optional[Callable] = None, + frozen_prefixes: Optional[Tuple[str, ...]] = None, +) -> None: + """Save model + optimizer + scheduler + metadata at *ckpt_path*. + + ``frozen_prefixes``: when set, model state-dict keys starting with any + of these prefixes are excluded from the saved file. This is used for + frozen-by-design submodules that should not be duplicated in every + checkpoint. Loaders already use ``strict=False``. + """ + log_fn = log_fn or _noop_log + os.makedirs(ckpt_path, exist_ok=True) + + if is_main: + _save_training_checkpoint_metadata( + ckpt_path=ckpt_path, + config=config, + epoch=epoch, + global_step=global_step, + seed=seed, + world_size=_world_size_for_metadata(), + normalizer_action=normalizer_action, + normalizer_propri=normalizer_propri, + grad_scaler=grad_scaler, + log_fn=log_fn, + ) + + if is_dmuon_model(model): + _save_dmuon_state_dict( + ckpt_path=ckpt_path, + model=model, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + is_main=is_main, + log_fn=log_fn, + frozen_prefixes=frozen_prefixes, + ) + elif _is_fsdp2_model(model): + _save_fsdp2_full_state_dict( + ckpt_path=ckpt_path, + model=model, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + is_main=is_main, + log_fn=log_fn, + frozen_prefixes=frozen_prefixes, + ) + else: + # DDP or unwrapped (fallback). + _save_ddp_state_dict( + ckpt_path=ckpt_path, + model=model, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + is_main=is_main, + log_fn=log_fn, + frozen_prefixes=frozen_prefixes, + ) + + # Per-rank dataset state (only when saving mid-epoch, step != 0). + if dataset is not None and global_step != 0: + _save_dataset_state( + ckpt_path=ckpt_path, dataset=dataset, rank=rank, log_fn=log_fn + ) + + +def _save_training_checkpoint_metadata( + *, + ckpt_path: str, + config: dict, + epoch: int, + global_step: int, + seed: int, + world_size: int, + normalizer_action, + normalizer_propri, + grad_scaler=None, + log_fn: Callable, +) -> None: + torch.save({"seed": seed}, os.path.join(ckpt_path, "seed.pth")) + torch.save({"global_step": global_step}, os.path.join(ckpt_path, "global_step.pth")) + torch.save({"current_epoch": epoch}, os.path.join(ckpt_path, "current_epoch.pth")) + + # world_size: public metadata, used at resume to detect reshard. + # Written unconditionally (previously sharded-mode only). + torch.save({"world_size": world_size}, os.path.join(ckpt_path, "world_size.pth")) + + # RNG state (rank-0 snapshot; every rank restores the same state on + # resume, matching the existing seed_all(seed) convention where every + # rank is seeded identically). + rng_state = { + "torch": torch.get_rng_state(), + "numpy": np.random.get_state(), + "python": random.getstate(), + } + if torch.cuda.is_available(): + rng_state["cuda"] = torch.cuda.get_rng_state() + torch.save(rng_state, os.path.join(ckpt_path, "rng_state.pt")) + + # GradScaler state (fp16 AMP only). + if grad_scaler is not None: + torch.save(grad_scaler.state_dict(), os.path.join(ckpt_path, "grad_scaler.pt")) + + with open(os.path.join(ckpt_path, "config.yml"), "w") as f: + yaml.dump(config, f, default_flow_style=False, allow_unicode=True) + + model_cfg = _dict_section(config, "model") + data_cfg = _dict_section(config, "data") + + # Copy processor files to checkpoint directory. + processor_dir = model_cfg.get("processor_path") or config.get("processor_path") + if processor_dir is None: + # Backward compatibility: fall back to pretrained_qwen_vl_path. + processor_dir = config.get("pretrained_qwen_vl_path") + if processor_dir is not None: + log_fn( + "WARNING: 'pretrained_qwen_vl_path' is deprecated for processor " + "file copying, please use 'processor_path' instead.", + level=logging.WARNING, + ) + + if processor_dir is not None: + for filename in ( + "preprocessor_config.json", + "tokenizer_config.json", + "tokenizer.json", + "tokenizer.model", + "special_tokens_map.json", + "vocab.json", + ): + src = os.path.join(processor_dir, filename) + if os.path.exists(src): + shutil.copy(src, os.path.join(ckpt_path, filename)) + + act_config_path = model_cfg.get("config_path") or config.get( + "qwen_vl_act_config_path" + ) + if act_config_path is not None: + if os.path.exists(act_config_path): + shutil.copy(act_config_path, os.path.join(ckpt_path, "config.json")) + log_fn(f"[Checkpoint] Copied act config to {ckpt_path}/config.json") + else: + log_fn(f"[Checkpoint] WARNING: {act_config_path} not found, skipping.") + + norm_stats_path = data_cfg.get("norm_stats_path") or config.get("norm_stats_path") + if norm_stats_path is not None: + if os.path.exists(norm_stats_path): + shutil.copy(norm_stats_path, os.path.join(ckpt_path, "norm_stats.json")) + log_fn(f"[Checkpoint] Copied norm stats to {ckpt_path}/norm_stats.json") + elif (data_cfg.get("dataset_type") or config.get("dataset_type")) == "lerobot": + log_fn(f"[Checkpoint] WARNING: {norm_stats_path} not found, skipping.") + + torch.save( + normalizer_action.state_dict(), + os.path.join(ckpt_path, "normalizer_action.pth"), + ) + torch.save( + normalizer_propri.state_dict(), + os.path.join(ckpt_path, "normalizer_propri.pth"), + ) + + +def _save_fsdp2_full_state_dict( + *, + ckpt_path: str, + model, + optimizer, + lr_scheduler, + is_main: bool, + log_fn: Callable, + frozen_prefixes: Optional[Tuple[str, ...]] = None, +) -> None: + """Save FSDP2 model + optimizer as rank-0 full state dict. + + ``get_state_dict`` with ``full_state_dict=True, cpu_offload=True`` + gathers DTensors to full CPU tensors on rank 0 (other ranks get empty + dicts). This is the cross-world-size-compatible format: on resume, + ``set_state_dict`` with ``broadcast_from_rank0=True`` re-shards from + rank 0's copy to whatever mesh the new run has. + """ + options = StateDictOptions(full_state_dict=True, cpu_offload=True) + model_sd, optim_sd = get_state_dict(model, optimizer, options=options) + + if is_main: + model_sd = _filter_frozen_prefixes(model_sd, frozen_prefixes, log_fn) + model_sd_out = _make_contiguous_and_clone_shared(model_sd) + save_file(model_sd_out, os.path.join(ckpt_path, "model.safetensors")) + torch.save(optim_sd, os.path.join(ckpt_path, "optimizer.pt")) + torch.save(lr_scheduler.state_dict(), os.path.join(ckpt_path, "scheduler.pt")) + log_fn("[Checkpoint] Saved FSDP2 full state dict (rank 0)") + + # Release the consolidated copy immediately; non-rank-0 already held {}. + del model_sd, optim_sd + gc.collect() + + +def _save_dmuon_state_dict( + *, + ckpt_path: str, + model, + optimizer, + lr_scheduler, + is_main: bool, + log_fn: Callable, + frozen_prefixes: Optional[Tuple[str, ...]] = None, +) -> None: + """Save via DMuon's state-dict helpers (full tensors, HF-compatible keys).""" + import dmuon + + model_sd = dmuon.get_model_state_dict(model, cpu_offload=True, rank0_only=True) + if is_main: + model_sd = _filter_frozen_prefixes(model_sd, frozen_prefixes, log_fn) + model_sd = _make_contiguous_and_clone_shared(model_sd) + save_file(model_sd, os.path.join(ckpt_path, "model.safetensors")) + log_fn("[Checkpoint] Saved DMuon model state dict (full, rank0)") + + optim_sd = dmuon.get_optimizer_state_dict( + model, optimizer, cpu_offload=True, rank0_only=True + ) + if is_main: + torch.save(optim_sd, os.path.join(ckpt_path, "optimizer.pt")) + torch.save(lr_scheduler.state_dict(), os.path.join(ckpt_path, "scheduler.pt")) + + del model_sd, optim_sd + gc.collect() + + +def _save_ddp_state_dict( + *, + ckpt_path: str, + model, + optimizer, + lr_scheduler, + is_main: bool, + log_fn: Callable = _noop_log, + frozen_prefixes: Optional[Tuple[str, ...]] = None, +) -> None: + if not is_main: + return + model_state = ( + model.module.state_dict() if isinstance(model, DDP) else model.state_dict() + ) + model_state = _filter_frozen_prefixes(model_state, frozen_prefixes, log_fn) + model_state_contiguous = _make_contiguous_and_clone_shared(model_state) + save_file(model_state_contiguous, os.path.join(ckpt_path, "model.safetensors")) + torch.save(optimizer.state_dict(), os.path.join(ckpt_path, "optimizer.pt")) + torch.save(lr_scheduler.state_dict(), os.path.join(ckpt_path, "scheduler.pt")) + + +def _save_dataset_state( + *, ckpt_path: str, dataset, rank: int, log_fn: Callable +) -> None: + """Save per-rank dataset resume state if the dataset supports it.""" + if not hasattr(dataset, "save_episode_containers"): + return + ec_path = os.path.join(ckpt_path, f"episode_containers_rank_{rank}.pkl") + dataset.save_episode_containers(ec_path) + + _t0 = time.time() + while not (os.path.exists(ec_path) and os.path.getsize(ec_path) > 0): + time.sleep(0.5) + if time.time() - _t0 > 120: + log_fn( + f"WARNING: episode container checkpoint save timeout: {ec_path}", + ) + break + + +def _filter_frozen_prefixes( + state_dict: dict, + frozen_prefixes: Optional[Tuple[str, ...]], + log_fn: Callable, +) -> dict: + """Drop entries whose key starts with any frozen prefix. + + Used to exclude frozen-by-design submodules from checkpoints. Returns + the input unchanged when no prefixes are configured or the dict is + already empty (non-rank-0 case). + """ + if not frozen_prefixes or not state_dict: + return state_dict + original = len(state_dict) + filtered = { + k: v for k, v in state_dict.items() if not k.startswith(frozen_prefixes) + } + log_fn( + f"[Checkpoint] Filtered state dict: {len(filtered)}/{original} entries " + f"(excluded {original - len(filtered)} frozen entries)" + ) + return filtered + + +def _make_contiguous_and_clone_shared(state_dict: dict) -> dict: + """Make tensors contiguous for safetensors, cloning any that share storage.""" + seen_data_ptrs = {} + out = {} + for k, v in state_dict.items(): + ptr = v.data_ptr() + if ptr in seen_data_ptrs: + v = v.clone() + else: + seen_data_ptrs[ptr] = k + out[k] = v.contiguous() if v.is_floating_point() or v.is_complex() else v + return out + + +def _world_size_for_metadata() -> int: + if dist.is_initialized(): + return dist.get_world_size() + return 1 + + +# ---------------------------------------------------------------------- +# Load +# ---------------------------------------------------------------------- + +_EMBED_WEIGHT_KEYS = ( + "model.embed_tokens.weight", + "model.language_model.embed_tokens.weight", +) + + +def _maybe_resize_token_embeddings_for_load( + model: torch.nn.Module, + state_dict: dict, + log_fn: Optional[Callable] = None, +) -> None: + """Resize model embeddings when checkpoint vocab size differs from the model.""" + log_fn = log_fn or _noop_log + for key in _EMBED_WEIGHT_KEYS: + if key not in state_dict: + continue + ckpt_vocab = state_dict[key].shape[0] + embed = model.get_input_embeddings() if hasattr(model, "get_input_embeddings") else None + if embed is None: + return + cur_vocab = embed.weight.shape[0] + if cur_vocab == ckpt_vocab: + return + log_fn( + f"resize_token_embeddings from {cur_vocab} to {ckpt_vocab} " + f"to match checkpoint ({key})" + ) + if hasattr(model, "resize_token_embeddings"): + model.resize_token_embeddings(ckpt_vocab) + elif hasattr(model, "model") and hasattr(model.model, "resize_token_embeddings"): + model.model.resize_token_embeddings(ckpt_vocab) + return + + +def load_weights( + *, + model: torch.nn.Module, + resume_config: dict, + model_class=None, + log_fn: Optional[Callable] = None, +): + """Load ``model`` weights from ``resume_config['ckpt']`` (file-level). + + Supports .safetensors / .pth sources and optional ``try_harder`` shape + matching for action-preprocessor weights when action dims changed. + """ + log_fn = log_fn or _noop_log + src = resume_config["ckpt"] + if src.endswith(".safetensors"): + log_fn(f"Loading model from safetensors: {src}") + state_dict = load_file(src) + elif src.endswith(".pth"): + checkpoint = torch.load(src, map_location="cpu", weights_only=False) + state_dict = checkpoint["model_state_dict"] + else: + raise ValueError(f"Unsupported checkpoint format: {src}") + + if model_class is not None and hasattr(model_class, "is_fused"): + if not model_class.is_fused(state_dict): + log_fn("Converting non-fused weights to fused format...") + state_dict = model_class.convert_to_fused(state_dict) + else: + log_fn("The weights is fused, skipping conversion.") + + filtered_state_dict = _drop_checkpoint_normalizer_state(state_dict, log_fn) + _maybe_resize_token_embeddings_for_load(model, filtered_state_dict, log_fn=log_fn) + + if resume_config.get("try_harder", False): + log_fn("### try harder to squeeze checkpoint weights into new model ###") + new_state_dict = reshape_compatible_state_dict( + filtered_state_dict, model.state_dict(), log_fn=log_fn + ) + err = model.load_state_dict(new_state_dict, strict=False) + else: + err = model.load_state_dict(filtered_state_dict, strict=False) + + log_fn(f"err in load model: {err}") + return model + + +def reshape_compatible_state_dict( + state_dict: dict, model_sd: dict, log_fn: Optional[Callable] = None +) -> dict: + """Pad / slice action-preprocessor weights to match target model shape.""" + log_fn = log_fn or _noop_log + out = {} + for name, param in state_dict.items(): + if name not in model_sd: + log_fn(f"Not used parameter: {name}") + continue + if "action_preprocessor" in name or "action_processor" in name: + if param.size() == model_sd[name].size(): + out[name] = param + continue + size_0 = param.size() + size_1 = model_sd[name].size() + if any(old_dim > new_dim for old_dim, new_dim in zip(size_0, size_1)): + raise ValueError( + f"Shape mismatch for '{name}': checkpoint shape {tuple(size_0)} is " + f"larger than model shape {tuple(size_1)} in at least one dimension. " + "Loading a larger checkpoint into a smaller model is not supported here. " + "If the action dimension has changed, please configure action padding " + "in the dataset processor so checkpoint actions match the new action " + "size." + ) + out[name] = model_sd[name].clone() + slices = [ + slice(0, min(old_dim, new_dim)) + for old_dim, new_dim in zip(size_0, size_1) + ] + out[name][slices] = param[slices] + log_fn( + f"Not match key: {name}, checkpoint shape: {tuple(size_0)}, " + f"model shape: {tuple(size_1)}. Filled checkpoint weights into the first " + f"{[s.stop for s in slices]} dims, remaining dims keep model init weights." + ) + else: + if param.size() == model_sd[name].size(): + out[name] = param + else: + log_fn( + f"Skipping '{name}': checkpoint shape {tuple(param.size())} " + f"!= model shape {tuple(model_sd[name].size())}" + ) + return out + + +def resume_from_checkpoint( + *, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + lr_scheduler, + resume_config: dict, + rank: int, + grad_scaler=None, + model_class=None, + log_fn: Optional[Callable] = None, +) -> None: + """Restore model / optimizer / scheduler / RNG / grad_scaler from ckpt.""" + log_fn = log_fn or _noop_log + checkpoint_path = resume_config["ckpt"] + + is_fsdp2 = _is_fsdp2_model(model) + is_dmuon = is_dmuon_model(model) + + # --- world_size diagnostic --------------------------------------- + ckpt_world_size = _read_ckpt_world_size(checkpoint_path, log_fn) + ws_mismatch = ( + ckpt_world_size is not None and ckpt_world_size != _world_size_for_metadata() + ) + if ws_mismatch: + log_fn( + f"Cross-world-size resume: ckpt_world_size={ckpt_world_size}, " + f"current={_world_size_for_metadata()}." + ) + + # --- single-file path (.safetensors / .pth) ---------------------- + if checkpoint_path.endswith(".safetensors") or checkpoint_path.endswith(".pth"): + _load_weights_into_model( + model=model, + resume_config=resume_config, + is_fsdp2=is_fsdp2, + is_dmuon=is_dmuon, + model_class=model_class, + log_fn=log_fn, + ) + log_fn(f"Resumed weights from single-file checkpoint: {checkpoint_path}") + return + + # --- directory path ---------------------------------------------- + if not os.path.isdir(checkpoint_path): + raise FileNotFoundError(f"Checkpoint path not found: {checkpoint_path}") + + safetensors_path = os.path.join(checkpoint_path, "model.safetensors") + rank_shard_path = os.path.join(checkpoint_path, f"model_rank{rank}.pt") + + if os.path.exists(safetensors_path): + if _detect_legacy_fsdp1_format(checkpoint_path): + # Pre-migration FSDP1 ckpt: per-rank optimizer files cannot be + # resharded into the new DTensor layout. Load model only and + # cold-start the optimizer. + _load_legacy_fsdp1_full( + checkpoint_path=checkpoint_path, + model=model, + is_fsdp2=is_fsdp2, + is_dmuon=is_dmuon, + model_class=model_class, + try_harder=resume_config.get("try_harder", False), + log_fn=log_fn, + ) + elif is_fsdp2 or is_dmuon: + # New-format path: model + optimizer via state_dict helpers, + # with automatic reshard on load via broadcast_from_rank0. + _load_fsdp2_or_dmuon_full( + checkpoint_path=checkpoint_path, + model=model, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + is_dmuon=is_dmuon, + ws_mismatch=ws_mismatch, + log_fn=log_fn, + ) + else: + # DDP / unwrapped: single-file model + single-file optimizer. + inner_resume = { + "ckpt": safetensors_path, + "try_harder": resume_config.get("try_harder", False), + } + _load_weights_into_model( + model=model, + resume_config=inner_resume, + is_fsdp2=False, + is_dmuon=False, + model_class=model_class, + log_fn=log_fn, + ) + _resume_ddp_optimizer_scheduler( + checkpoint_path=checkpoint_path, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + log_fn=log_fn, + ) + elif os.path.exists(rank_shard_path): + # FSDP1 sharded checkpoints are rank-layout-bound; dropping support. + raise RuntimeError( + f"Legacy FSDP1 sharded checkpoint at {checkpoint_path} is no " + f"longer supported. Convert to a single model.safetensors first." + ) + else: + raise FileNotFoundError( + f"No model.safetensors or model_rank*.pt found under {checkpoint_path}" + ) + + # --- auxiliary state (all paths) --------------------------------- + _resume_rng(checkpoint_path, log_fn) + if grad_scaler is not None: + _resume_grad_scaler(checkpoint_path, grad_scaler, log_fn) + log_fn(f"Resumed from checkpoint: {checkpoint_path}") + + +def _read_ckpt_world_size(checkpoint_path: str, log_fn: Callable): + if not os.path.isdir(checkpoint_path): + return None + ws_path = os.path.join(checkpoint_path, "world_size.pth") + if not os.path.exists(ws_path): + return None + try: + return int(torch.load(ws_path, map_location="cpu")["world_size"]) + except (KeyError, ValueError, RuntimeError, TypeError) as e: + log_fn( + f"world_size.pth unreadable ({e}); assuming same ws.", + level=logging.WARNING, + ) + return None + + +def _strip_fused_flags(osd: dict) -> None: + """In-place pop of fused / foreach flags from an optimizer state dict's + param_groups. Old checkpoints may have fused=True which breaks dtype + matching when the optimizer is reconstructed under a different precision.""" + for pg in osd.get("param_groups", []) or []: + pg.pop("fused", None) + pg.pop("foreach", None) + + +def _drop_checkpoint_normalizer_state(state_dict: dict, log_fn: Callable) -> dict: + """Drop saved normalizer buffers so current-run stats stay authoritative.""" + normalizer_state_prefixes = ( + "action_preprocessor.normalizer", + "action_processor.normalizer", + ) + filtered = { + k: v + for k, v in state_dict.items() + if not k.startswith(normalizer_state_prefixes) + } + dropped = len(state_dict) - len(filtered) + if dropped: + log_fn( + f"[Checkpoint] Dropped {dropped} checkpoint normalizer entries; " + "keeping current-run normalizers." + ) + return filtered + + +def _load_fsdp2_or_dmuon_full( + *, + checkpoint_path: str, + model, + optimizer, + lr_scheduler, + is_dmuon: bool, + ws_mismatch: bool, + log_fn: Callable, +) -> None: + """Load FSDP2 (or DMuon) model + optimizer from a full-state-dict ckpt. + + FSDP2 path: rank 0 deserializes, ``broadcast_from_rank0=True`` in + ``set_*_state_dict`` reshards to the current mesh. + + DMuon path: every rank MUST load the full state dict from disk. + ``dmuon.set_model_state_dict`` / ``set_optimizer_state_dict`` iterate + ``fqn_to_dp`` and skip any FQN not present in the provided dict — so + if only rank 0 has the data, new owners on other ranks (after a + cross-ws resume) silently miss their assigned params and + ``_owned_data`` stays at fresh-init (random weights), producing a + massive post-resume loss spike. The ckpt files live on CPFS and are + shared, so per-rank reads are essentially free. + """ + is_main = dist.get_rank() == 0 if dist.is_initialized() else True + safetensors_path = os.path.join(checkpoint_path, "model.safetensors") + optim_path = os.path.join(checkpoint_path, "optimizer.pt") + sched_path = os.path.join(checkpoint_path, "scheduler.pt") + + # --- model --- + if is_dmuon: + if not os.path.exists(safetensors_path): + raise FileNotFoundError( + f"DMuon resume requires {safetensors_path} on every rank." + ) + # Every rank loads the full state dict so new owners (after a + # cross-ws reshard, which re-runs dedicate_params and may assign + # FQNs to different ranks) each see their own FQN in the dict. + # dmuon.set_model_state_dict is FQN-keyed and only the new owner + # writes into _owned_data, so passing the full dict on every + # rank is correct (and required) for the cross-ws case. + model_sd = _drop_checkpoint_normalizer_state( + load_file(safetensors_path), log_fn + ) + import dmuon + + dmuon.set_model_state_dict(model, model_sd) + log_fn("[Checkpoint] DMuon model state loaded.") + else: + if is_main and os.path.exists(safetensors_path): + model_sd = _drop_checkpoint_normalizer_state( + load_file(safetensors_path), log_fn + ) + else: + model_sd = {} + options = StateDictOptions( + full_state_dict=True, + cpu_offload=True, + broadcast_from_rank0=True, + strict=False, + ) + set_model_state_dict( + model, + model_state_dict=model_sd, + options=options, + ) + log_fn("[Checkpoint] FSDP2 model state loaded (reshard via broadcast).") + del model_sd + gc.collect() + + # --- optimizer --- + if os.path.exists(optim_path): + if is_dmuon: + # Same rationale as model_sd: every rank needs full optim state. + optim_sd = torch.load(optim_path, map_location="cpu", weights_only=False) + _strip_fused_flags(optim_sd) + elif is_main: + optim_sd = torch.load(optim_path, map_location="cpu", weights_only=False) + _strip_fused_flags(optim_sd) + else: + optim_sd = {} + + try: + if is_dmuon: + import dmuon + + dmuon.set_optimizer_state_dict(model, optimizer, optim_sd) + log_fn("[Checkpoint] DMuon optimizer state loaded.") + else: + options = StateDictOptions( + full_state_dict=True, + cpu_offload=True, + broadcast_from_rank0=True, + ) + set_optimizer_state_dict( + model, + optimizers=optimizer, + optim_state_dict=optim_sd, + options=options, + ) + log_fn( + "[Checkpoint] FSDP2 optimizer state loaded " + "(reshard via broadcast)." + ) + except (ValueError, RuntimeError, TypeError, KeyError) as e: + # Most likely a legacy FSDP1 consolidated OSD (flat_param-specific). + log_fn( + f"[Checkpoint] Optimizer state incompatible with current " + f"layout ({e!r}); cold-starting optimizer. Model weights are " + f"loaded; momentum resets to zero.", + level=logging.WARNING, + ) + del optim_sd + gc.collect() + else: + log_fn( + "[Checkpoint] optimizer.pt not found; cold-starting optimizer.", + level=logging.WARNING, + ) + + # --- scheduler --- + if os.path.exists(sched_path): + try: + lr_scheduler.load_state_dict( + torch.load(sched_path, map_location="cpu", weights_only=False) + ) + log_fn("[Checkpoint] Scheduler state loaded.") + except (ValueError, RuntimeError, TypeError, KeyError) as e: + log_fn( + f"Scheduler load failed ({e}); keeping fresh state.", + level=logging.WARNING, + ) + + +def _resume_ddp_optimizer_scheduler( + *, + checkpoint_path: str, + optimizer, + lr_scheduler, + log_fn: Callable, +) -> None: + """Restore optimizer / scheduler for DDP layouts (single-file).""" + optimizer_path = os.path.join(checkpoint_path, "optimizer.pt") + if os.path.exists(optimizer_path): + try: + optim_sd = torch.load( + optimizer_path, map_location="cpu", weights_only=False + ) + _strip_fused_flags(optim_sd) + optimizer.load_state_dict(optim_sd) + log_fn("[Checkpoint] Optimizer state loaded (single file).") + except (ValueError, RuntimeError) as e: + log_fn( + f"Failed to load optimizer state dict, " + f"optimizer will be re-initialized. Error: {e}", + level=logging.WARNING, + ) + + sched_path = os.path.join(checkpoint_path, "scheduler.pt") + if os.path.exists(sched_path): + try: + lr_scheduler.load_state_dict( + torch.load(sched_path, map_location="cpu", weights_only=False) + ) + log_fn(f"[Checkpoint] Scheduler state loaded ({sched_path}).") + except (ValueError, RuntimeError, KeyError, TypeError) as e: + log_fn( + f"Scheduler load failed ({e}); keeping fresh state.", + level=logging.WARNING, + ) + + +def _load_legacy_fsdp1_full( + *, + checkpoint_path: str, + model, + is_fsdp2: bool, + is_dmuon: bool, + model_class, + try_harder: bool, + log_fn: Callable, +) -> None: + """Load a pre-migration FSDP1 ckpt: weights only, optimizer cold-starts. + + The per-rank ``optimizer_rank{N}.pt`` files use FSDP1's flat_param + layout, which cannot be resharded into FSDP2's DTensor layout. The + only safe action is to load the rank-0 ``model.safetensors`` (which + is layout-agnostic) and let the optimizer warm up from scratch. + """ + inner_resume = { + "ckpt": os.path.join(checkpoint_path, "model.safetensors"), + "try_harder": try_harder, + } + _load_weights_into_model( + model=model, + resume_config=inner_resume, + is_fsdp2=is_fsdp2, + is_dmuon=is_dmuon, + model_class=model_class, + log_fn=log_fn, + ) + log_fn( + f"Legacy FSDP1 ckpt detected at {checkpoint_path}: model weights " + f"loaded, but per-rank optimizer files cannot reshard into FSDP2. " + f"Optimizer state DROPPED — it will warm up from scratch. For long " + f"resume runs prefer a fresh FSDP2 ckpt; for short runs the loss " + f"bump is usually negligible.", + level=logging.WARNING, + ) + + +def _resume_rng(checkpoint_path: str, log_fn: Callable) -> None: + path = os.path.join(checkpoint_path, "rng_state.pt") + if not os.path.exists(path): + log_fn( + "rng_state.pt not found; RNG stays at seed_all() state.", + level=logging.WARNING, + ) + return + try: + sd = torch.load(path, map_location="cpu", weights_only=False) + if "torch" in sd: + torch.set_rng_state(sd["torch"]) + if "cuda" in sd and torch.cuda.is_available(): + torch.cuda.set_rng_state(sd["cuda"]) + if "numpy" in sd: + np.random.set_state(sd["numpy"]) + if "python" in sd: + random.setstate(sd["python"]) + log_fn("[Checkpoint] RNG state restored.") + except (RuntimeError, ValueError, TypeError, KeyError) as e: + log_fn( + f"RNG load failed ({e}); keeping current RNG state.", + level=logging.WARNING, + ) + + +def _resume_grad_scaler(checkpoint_path: str, grad_scaler, log_fn: Callable) -> None: + path = os.path.join(checkpoint_path, "grad_scaler.pt") + if not os.path.exists(path): + log_fn( + "grad_scaler.pt not found; GradScaler stays fresh.", + level=logging.WARNING, + ) + return + try: + grad_scaler.load_state_dict( + torch.load(path, map_location="cpu", weights_only=False) + ) + log_fn("[Checkpoint] GradScaler state restored.") + except (RuntimeError, ValueError, TypeError, KeyError) as e: + log_fn( + f"GradScaler load failed ({e}); keeping fresh state.", + level=logging.WARNING, + ) + + +def _load_weights_into_model( + *, + model: torch.nn.Module, + resume_config: dict, + is_fsdp2: bool = False, + is_dmuon: bool = False, + model_class, + log_fn: Callable, +) -> None: + """Weights-only load for single .safetensors / .pth sources.""" + if is_dmuon or is_fsdp2: + src = resume_config["ckpt"] + is_main = dist.get_rank() == 0 if dist.is_initialized() else True + # DMuon needs the full state dict on every rank (FQN-keyed lookup; + # see _load_fsdp2_or_dmuon_full above). FSDP2 with + # broadcast_from_rank0=True only needs rank 0 to read. + if is_dmuon or is_main: + if src.endswith(".safetensors"): + state_dict = load_file(src) + else: + state_dict = torch.load(src, map_location="cpu", weights_only=False) + if isinstance(state_dict, dict) and "model_state_dict" in state_dict: + state_dict = state_dict["model_state_dict"] + state_dict = _drop_checkpoint_normalizer_state(state_dict, log_fn) + else: + state_dict = {} + if is_dmuon: + import dmuon + + dmuon.set_model_state_dict(model, state_dict) + log_fn(f"[DMuon] Loaded model weights from {src}") + else: + options = StateDictOptions( + full_state_dict=True, + cpu_offload=True, + broadcast_from_rank0=True, + strict=False, + ) + set_model_state_dict( + model, + model_state_dict=state_dict, + options=options, + ) + log_fn(f"[FSDP2] Loaded model weights from {src}") + else: + # DDP / unwrapped. + unwrapped_model = model.module if isinstance(model, DDP) else model + load_weights( + model=unwrapped_model, + resume_config=resume_config, + model_class=model_class, + log_fn=log_fn, + ) + + +def finalize_save(log_fn: Optional[Callable] = None) -> None: + """Barrier + GC hygiene after a checkpoint write. Trainer calls this.""" + if dist.is_initialized(): + dist.barrier() + gc.collect() + torch.cuda.empty_cache() + if log_fn: + pass # trainer already logs the "Saved checkpoint to X" message diff --git a/wall_x/trainer/fsdp_trainer/distribution_strategy.py b/wall_x/trainer/fsdp_trainer/distribution_strategy.py new file mode 100644 index 00000000..57e33365 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/distribution_strategy.py @@ -0,0 +1,400 @@ +"""Distribution strategies for model wrapping and gradient coordination.""" + +from __future__ import annotations + +import logging +import os +from abc import ABC, abstractmethod +from contextlib import contextmanager +from importlib import import_module +from typing import NamedTuple, Optional + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh +from torch.distributed.fsdp import ( + MixedPrecisionPolicy, + OffloadPolicy, +) +from torch.nn.parallel import DistributedDataParallel as DDP + +_logger = logging.getLogger(__name__) + + +class FSDP2Layout(NamedTuple): + """Per-call FSDP2 wrap configuration produced by ``FSDPStrategy``.""" + + mesh: DeviceMesh + dp_process_group: Optional[dist.ProcessGroup] + shard_process_group: Optional[dist.ProcessGroup] + replicate_process_group: Optional[dist.ProcessGroup] + trainer_process_group: Optional[dist.ProcessGroup] + mp_policy: Optional[MixedPrecisionPolicy] + offload_policy: Optional[OffloadPolicy] + reshard_after_forward: bool + + +class DistributionStrategy(ABC): + """Common interface for FSDP2 / DDP wrapping + gradient operations.""" + + _cfg: dict + last_grad_clip_stats: Optional[dict] = None + + def _maybe_enable_grad_ckpt(self, model: torch.nn.Module) -> None: + """Call ``gradient_checkpointing_enable`` if the model exposes it.""" + if not self._cfg.get("use_gradient_checkpointing", False): + return + method = getattr(model, "gradient_checkpointing_enable", None) + if method is None: + _logger.warning( + "use_gradient_checkpointing=true but %s has no " + "gradient_checkpointing_enable method. If this model " + "routes recomputation through its config the request may " + "still be honored; otherwise " + "gradient checkpointing is effectively disabled for " + "this run.", + type(model).__name__, + ) + return + method() + + @abstractmethod + def wrap(self, model: torch.nn.Module) -> torch.nn.Module: + """Wrap *model* with the active distribution strategy.""" + + @abstractmethod + def clip_grad_norm( + self, + model: torch.nn.Module, + max_norm: float, + *, + optimizer=None, + ) -> torch.Tensor: + """Clip gradient L2 norm; return total norm before clipping.""" + + @abstractmethod + def no_sync(self, model: torch.nn.Module): + """Context manager that disables cross-rank gradient sync. + + Used by the trainer for gradient accumulation micro-batches. + FSDP2 toggles ``set_requires_gradient_sync``; DDP yields + ``model.no_sync()``. + """ + + @property + @abstractmethod + def params_sharded(self) -> bool: + """Whether parameters are actually sharded across ranks. + + FSDP2 always shards (every ``fully_shard`` unit is sharded across + the mesh). DDP never shards. Used by adapters to decide whether + per-rank norm contributions need an all-reduce. + """ + + +class FSDPStrategy(DistributionStrategy): + def __init__(self, config: dict): + self._cfg = config + self._layout: Optional[FSDP2Layout] = None + self.last_grad_clip_stats = None + + @property + def trainer_process_group(self) -> Optional[dist.ProcessGroup]: + """Process group for trainer-level collectives. + + Returns `dist.group.WORLD` (the implicit all-rank group created + by `init_process_group`) so that trainer barriers / metric + reductions behave consistently regardless of FSDP mesh shape + (1D for full_shard / shard_grad_op, 2D for hybrid_shard / + _hybrid_shard_zero2). No extra NCCL communicator is allocated. + """ + return self._layout.trainer_process_group if self._layout else None + + @staticmethod + def _build_trainer_process_group( + world_size: int, + ) -> Optional[dist.ProcessGroup]: + if not dist.is_initialized() or world_size <= 1: + return None + return dist.group.WORLD + + def _build_fsdp2_layout(self) -> FSDP2Layout: + """Construct the FSDP2 wrap layout from yaml config. + + Builds the device mesh (1D for full_shard / shard_grad_op, 2D + for hybrid_shard / _hybrid_shard_zero2), explicit FSDP mesh process + groups, a trainer-level global process group, the matching + mixed-precision policy, offload policy, and reshard_after_forward flag. + """ + name = self._cfg.get("fsdp_sharding_strategy", "full_shard") + world_size = dist.get_world_size() if dist.is_initialized() else 1 + dp_pg: Optional[dist.ProcessGroup] = None + shard_pg: Optional[dist.ProcessGroup] = None + replicate_pg: Optional[dist.ProcessGroup] = None + + # 2D mesh for HSDP variants; 1D for full_shard / shard_grad_op. + if name in ("hybrid_shard", "_hybrid_shard_zero2"): + replicate_size = self._cfg.get("fsdp_hsdp_replicate_size") or int( + os.environ.get("LOCAL_WORLD_SIZE", "0") + ) + if replicate_size <= 0: + # Fallback: assume one replica per node, equal-sized shard groups. + # Caller should set fsdp_hsdp_replicate_size explicitly when this + # default is wrong (e.g. uneven node sizes). + replicate_size = max(1, world_size // 8) if world_size >= 8 else 1 + if world_size % replicate_size != 0: + raise ValueError( + f"world_size ({world_size}) not divisible by HSDP replicate " + f"size ({replicate_size}); set fsdp_hsdp_replicate_size." + ) + shard_size = world_size // replicate_size + mesh = init_device_mesh( + "cuda", + (replicate_size, shard_size), + mesh_dim_names=("replicate", "shard"), + ) + replicate_pg = mesh.get_group("replicate") + shard_pg = mesh.get_group("shard") + else: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",)) + dp_pg = mesh.get_group("dp") + + trainer_pg = self._build_trainer_process_group(world_size) + + reshard_after_forward = name not in ( + "shard_grad_op", + "_hybrid_shard_zero2", + ) + + offload_policy: Optional[OffloadPolicy] = ( + OffloadPolicy(pin_memory=True) + if self._cfg.get("fsdp_cpu_offload", False) + else None + ) + + mp_policy: Optional[MixedPrecisionPolicy] + if self._cfg.get("use_mixed_precision", True): + dtype = torch.bfloat16 if self._cfg.get("bf16", True) else torch.float16 + reduce_dtype = ( + torch.float32 + if self._cfg.get("fsdp_reduce_dtype", "bf16") == "fp32" + else dtype + ) + mp_policy = MixedPrecisionPolicy( + param_dtype=dtype, + reduce_dtype=reduce_dtype, + cast_forward_inputs=False, + ) + else: + mp_policy = None + + assert mp_policy is not None, ( + "FSDP2 requires a MixedPrecisionPolicy in this torch version. " + "Set distributed.use_mixed_precision=true." + ) + + return FSDP2Layout( + mesh=mesh, + dp_process_group=dp_pg, + shard_process_group=shard_pg, + replicate_process_group=replicate_pg, + trainer_process_group=trainer_pg, + mp_policy=mp_policy, + offload_policy=offload_policy, + reshard_after_forward=reshard_after_forward, + ) + + def _wrap_fsdp( + self, + model: torch.nn.Module, + *, + use_dmuon: bool = False, + ) -> torch.nn.Module: + if not hasattr(model, "convert_to_fsdp"): + raise NotImplementedError( + f"Model {model.__class__.__name__} has no convert_to_fsdp method" + ) + self._maybe_enable_grad_ckpt(model) + layout = self._build_fsdp2_layout() + self._layout = layout + wrapped = model.convert_to_fsdp( + mesh=layout.mesh, + mp_policy=layout.mp_policy, + offload_policy=layout.offload_policy, + reshard_after_forward=layout.reshard_after_forward, + use_dmuon=use_dmuon, + ) + torch.cuda.empty_cache() + return wrapped + + def wrap(self, model: torch.nn.Module) -> torch.nn.Module: + return self._wrap_fsdp(model, use_dmuon=False) + + def clip_grad_norm( + self, + model: torch.nn.Module, + max_norm: float, + *, + optimizer=None, + ) -> torch.Tensor: + del optimizer + self.last_grad_clip_stats = None + # Separate DTensor (FSDP2) from regular tensor params — PyTorch's + # foreach_mul_ in clip_grad_norm_ can't mix them. + try: + from torch.distributed.tensor import DTensor + except ImportError: + DTensor = None + dtensor_params = [] + regular_params = [] + for p in model.parameters(): + if p.grad is None: + continue + if DTensor is not None and isinstance(p, DTensor): + dtensor_params.append(p) + else: + regular_params.append(p) + + device = next(iter(p for p in model.parameters())).device + total_sq = torch.tensor(0.0, device=device) + if dtensor_params: + n = torch.nn.utils.clip_grad_norm_(dtensor_params, max_norm) + total_sq = total_sq + n.to(device).pow(2) + if regular_params: + n = torch.nn.utils.clip_grad_norm_(regular_params, max_norm) + total_sq = total_sq + n.to(device).pow(2) + return total_sq.sqrt() + + @contextmanager + def no_sync(self, model: torch.nn.Module): + toggled = hasattr(model, "set_requires_gradient_sync") + if toggled: + model.set_requires_gradient_sync(False) + try: + yield + finally: + if toggled: + model.set_requires_gradient_sync(True) + + @property + def params_sharded(self) -> bool: + return True + + +class DMuonFSDPStrategy(FSDPStrategy): + """FSDP2 plus the DMuon dedicated-parameter runtime.""" + + @staticmethod + def _runtime(): + return import_module("dmuon") + + def wrap(self, model: torch.nn.Module) -> torch.nn.Module: + return self._wrap_fsdp(model, use_dmuon=True) + + def clip_grad_norm( + self, + model: torch.nn.Module, + max_norm: float, + *, + optimizer=None, + ) -> torch.Tensor: + total_norm = super().clip_grad_norm(model, max_norm, optimizer=optimizer) + self.last_grad_clip_stats = None + if optimizer is None: + return total_norm + + clip_stats = self._runtime().clip_grad_norm_(optimizer, max_norm) + self.last_grad_clip_stats = clip_stats.as_dict() + device = next(iter(p for p in model.parameters())).device + total_norm_t = ( + total_norm.to(device) + if torch.is_tensor(total_norm) + else torch.tensor(float(total_norm), device=device) + ) + dedicated_norm_t = torch.tensor(float(clip_stats.total_norm), device=device) + return (total_norm_t.pow(2) + dedicated_norm_t.pow(2)).sqrt() + + @contextmanager + def no_sync(self, model: torch.nn.Module): + with self._runtime().no_sync(model): + yield + + +class DDPStrategy(DistributionStrategy): + def __init__( + self, + config: dict, + *, + device: torch.device, + local_rank: int, + ): + self._cfg = config + self._device = device + self._local_rank = local_rank + self.find_unused_parameters = config.get("find_unused_parameters", False) + self.broadcast_buffers = config.get("broadcast_buffers", True) + self.bucket_cap_mb = config.get("bucket_cap_mb", 25) + self.last_grad_clip_stats = None + + def wrap(self, model: torch.nn.Module) -> torch.nn.Module: + # Adapter.convert_to_mix_precision_hint already placed the model on + # self._device; this .to() is a defensive no-op (same as the original + # wrap_ddp_model). + model = model.to(self._device) + self._maybe_enable_grad_ckpt(model) + return DDP( + model, + device_ids=[self._local_rank], + output_device=self._local_rank, + find_unused_parameters=self.find_unused_parameters, + broadcast_buffers=self.broadcast_buffers, + bucket_cap_mb=self.bucket_cap_mb, + ) + + def clip_grad_norm( + self, + model: torch.nn.Module, + max_norm: float, + *, + optimizer=None, + ) -> torch.Tensor: + del optimizer + self.last_grad_clip_stats = None + return torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) + + @contextmanager + def no_sync(self, model: torch.nn.Module): + with model.no_sync(): + yield + + @property + def params_sharded(self) -> bool: + return False + + +def build_strategy( + config: dict, + *, + device: torch.device, + local_rank: int, + use_dmuon: bool = False, +) -> DistributionStrategy: + """Pick FSDP vs DDP based on ``use_fsdp`` flag. + + ``use_gradient_checkpointing`` is read from the config dict by + each strategy's ``_maybe_enable_grad_ckpt``; not a separate kwarg here. + """ + if config.get("use_fsdp", False): + if use_dmuon: + return DMuonFSDPStrategy(config) + return FSDPStrategy(config) + if use_dmuon: + raise NotImplementedError( + "DMuon dedicated-parameter training requires " + "distributed.use_fsdp=true in this trainer." + ) + return DDPStrategy( + config, + device=device, + local_rank=local_rank, + ) diff --git a/wall_x/trainer/fsdp_trainer/fsdp_trainer.py b/wall_x/trainer/fsdp_trainer/fsdp_trainer.py new file mode 100644 index 00000000..11e7c8e7 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/fsdp_trainer.py @@ -0,0 +1,963 @@ +"""FSDP trainer implementation.""" + +import dataclasses +import gc +import os +import torch +import time +import logging +import contextlib +from torch.nn.parallel import DistributedDataParallel as DDP +from typing import Optional, Any, Tuple +from wall_x.trainer.optimizer import get_optimizer + +from wall_x.trainer.fsdp_trainer.base_trainer import ( + DistributedTrainer, + is_main_process, + barrier, +) +from wall_x.trainer.adapters import ADAPTER_REGISTRY, format_adapter_error +from wall_x.trainer.fsdp_trainer import checkpoint_io as _ckpt_io +from wall_x.trainer.fsdp_trainer.distribution_strategy import build_strategy +from wall_x.trainer.fsdp_trainer.metrics_logger import MetricsLogger +from wall_x.trainer.utils import move_batch_to_device + +from wall_x.model.core.action.normalizer import create_normalizers +from wall_x.trainer.scheduler.scheduler import get_scheduler + + +class FSDPTrainer(DistributedTrainer): + """ + FSDP-based Trainer using pure PyTorch FSDP2 (``fully_shard``). + + Features: + - No accelerate dependency + - Direct torchrun launch + - Mixed precision training support (bfloat16/float16) + - Gradient accumulation + - Full/Sharded state dict saving + - CPU offload support + - Activation checkpointing + + Launch: + torchrun --nproc_per_node=8 --master_port=29500 train_fsdp.py --config config.yaml + """ + + def __init__( + self, + train_config, + wandb_run: Optional[Any] = None, + ): + if train_config is None: + raise ValueError("train_config is required") + + self.model_type = train_config.model_type + super().__init__(train_config, wandb_run) + + if self.model_type not in ADAPTER_REGISTRY: + raise ValueError(format_adapter_error(self.model_type)) + adapter_cls = ADAPTER_REGISTRY[self.model_type] + self.adapter = adapter_cls( + cfg=self.cfg, + logger=self.logger, + model_type=self.model_type, + ) + + self.strategy = build_strategy( + dataclasses.asdict(self.cfg.distributed), + device=self.device, + local_rank=self.local_rank, + use_dmuon=self.cfg.hyperparams.optimizer.optimizer_type == "dmuon", + ) + + # Model config + self.action_dim = self.cfg.action_dim + self.use_selective_recompute = self.cfg.distributed.use_selective_recompute + self.show_time_details = self.cfg.debug.show_time_details + + # num_training_steps: read from any scheduler that exposes it + # (CosineSchedulerConfig currently; future schedulers may add it). + # Used in train_loop to trigger loss_guard_should_stop — works with + # constant scheduler too once the field is set. + sched = self.cfg.hyperparams.scheduler + self.num_training_steps = int(getattr(sched, "num_training_steps", 0) or 0) + self.metrics_logger = MetricsLogger( + wandb_run=self.wandb_run, + log_interval=self.log_interval, + smooth_window=self.cfg.logging.loss_log_smooth_window, + ) + + # Loss tracking + self.base_l1_loss = None + self.base_l1_loss_detail = {} + + self.load_normalizer() + self.load_processor() + self.load_model() + + # DDP needs to see correct requires_grad at wrap time, so freeze first. + self._freeze_params_if_needed(self.model) + + # Frozen-submodule prefixes for checkpoint filtering. Must be computed + # before FSDP wrapping, while child modules still have stable names. + self._frozen_prefixes = self._compute_frozen_prefixes() + + # Original shapes captured before FSDP flattening. + self.model._orig_param_shapes = { + 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() + + if self._resume_from_training_checkpoint(): + self.resume_from_checkpoint() + + self.load_dataset() + + self.adapter.init_validation(self.normalizer_action, self.normalizer_propri) + + def load_normalizer(self): + norm_cfg = self.cfg.data.normalizer_config or {} + custom_stats_path = norm_cfg.get("customized_action_statistic_dof") + if not custom_stats_path and self.cfg.data.dataset_type == "lerobot": + from wall_x.data.backends.lerobot.build import load_lerobot_normalizers + + loaded = load_lerobot_normalizers(self.cfg) + if loaded is not None: + self.normalizer_action, self.normalizer_propri = ( + loaded[0], + loaded[1], + ) + self._action_statistic_dof = None + self.logger.info( + "Loaded LeRobot normalizers from %s with dataset key %s", + loaded[2], + loaded[3], + ) + return + + merged = { + "dof_config": self.cfg.task.dof_config, + "agent_pos_config": self.cfg.task.agent_pos_config, + "customized_action_statistic_dof": custom_stats_path, + "min_key": norm_cfg.get("min_key", "min"), + "delta_key": norm_cfg.get("delta_key", "delta"), + } + self.normalizer_action, self.normalizer_propri, self._action_statistic_dof = ( + create_normalizers(merged) + ) + + def backward(self, loss: torch.Tensor): + """Perform backward pass.""" + if self.grad_scaler is not None: + self.grad_scaler.scale(loss).backward() + else: + loss.backward() + + def clip_grad_norm(self, max_norm: float) -> torch.Tensor: + """Unscale and clip gradient norm via the active distribution strategy.""" + if self.grad_scaler is not None: + self.grad_scaler.unscale_(self.optimizer) + + total_norm = self.strategy.clip_grad_norm( + self.model, + max_norm, + optimizer=self.optimizer, + ) + self._dedicated_param_grad_clip_stats = getattr( + self.strategy, + "last_grad_clip_stats", + None, + ) + return total_norm + + def load_model(self): + """Load and prepare model, optimizer, and scheduler""" + self.model_config = self.adapter.build_model_config() + self.model = self.adapter.create_model( + self.processor, self.tokenizer_mixin, self.model_config + ) + type(self.adapter).log_attention_implementation(self.logger, self.model) + self.adapter.load_weights( + self.model, + self.normalizer_action, + self.normalizer_propri, + processor=self.processor, + ) + + def load_processor(self): + """Load processor and tokenizers""" + self.adapter.normalizer_action = self.normalizer_action + self.adapter.normalizer_propri = self.normalizer_propri + processors_dict = self.adapter.load_processor(self._action_statistic_dof) + self.processor = processors_dict["processor"] + self.train_action_tokenizer = processors_dict["train_action_tokenizer"] + self.val_action_tokenizer = processors_dict["val_action_tokenizer"] + self.action_mapper = processors_dict["action_mapper"] + self.tokenizer_mixin = processors_dict.get("tokenizer_mixin") + + def _freeze_params_if_needed(self, model: torch.nn.Module): + """Freeze non-action parameters when train_action_expert_only is set. + + Must be called BEFORE wrapping with DDP/FSDP so the wrapper sees the + correct requires_grad flags and does not expect gradients for frozen params. + """ + from wall_x.trainer.optimizer.utils import resolve_lr_group_configs + + opt = self.cfg.hyperparams.optimizer + if not opt.train_action_expert_only: + return + lr_groups = resolve_lr_group_configs( + opt, self.adapter.default_action_lr_keywords + ) + if not lr_groups: + self.log( + "WARNING: train_action_expert_only is True but no optimizer LR " + "group is set. No parameters will be frozen.", + level=logging.WARNING, + ) + return + + frozen_count = 0 + grouped_count = 0 + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + matches_group = any( + any(keyword in name for keyword in group.include) for group in lr_groups + ) + if matches_group: + grouped_count += 1 + else: + param.requires_grad = False + frozen_count += 1 + + if grouped_count == 0: + raise ValueError( + "No grouped params found for train_action_expert_only. " + "Please check optimizer.lr_groups or legacy action_lr_keywords." + ) + + self.log( + f"*** train_action_expert_only: frozen {frozen_count} base params, " + f"keeping {grouped_count} grouped params trainable ***" + ) + + def _wrap_model(self, model: torch.nn.Module): + """Prepare model dtype/placement via adapter, then wrap via strategy. + + If the strategy created an explicit trainer process group, register + it so trainer-level barriers and metric collectives keep all-rank + semantics independent of the FSDP mesh topology. + """ + self.adapter.convert_to_mix_precision_hint( + model, + device=self.device, + use_fsdp=self.cfg.distributed.use_fsdp, + log_fn=self.log, + ) + self.model = self.strategy.wrap(model) + + trainer_pg = getattr(self.strategy, "trainer_process_group", None) + if trainer_pg is not None: + from wall_x.trainer.fsdp_trainer.base_trainer import ( + set_trainer_process_group, + ) + + set_trainer_process_group(trainer_pg) + self.log("[FSDP2] routing trainer collectives through trainer PG") + + def _create_optimizer(self): + """Create optimizer for FSDP wrapped model.""" + from wall_x.trainer.optimizer.utils import ( + build_lr_param_groups, + resolve_lr_group_configs, + uses_legacy_action_lr_groups, + ) + + opt_cfg = self.cfg.hyperparams.optimizer + + param_groups = None + lr_group_configs = resolve_lr_group_configs( + opt_cfg, self.adapter.default_action_lr_keywords + ) + if lr_group_configs: + if opt_cfg.optimizer_type not in ("adamw", "muon", "dmuon"): + raise ValueError( + "optimizer.lr_groups are only supported with adamw, muon, " + "or dmuon" + ) + base_group_name = ( + "base_lr_group" if uses_legacy_action_lr_groups(opt_cfg) else "base" + ) + param_groups = build_lr_param_groups( + self.model, + opt_cfg, + lr_group_configs, + base_group_name=base_group_name, + ) + summary = ", ".join( + f"{group.name}={group.lr} include={group.include}" + for group in lr_group_configs + ) + self.log( + f"setting optimizer LR groups: base={opt_cfg.learning_rate}; " + f"{summary} ({opt_cfg.optimizer_type})" + ) + + optimizer_kwargs = { + "opt_cfg": opt_cfg, + "param_groups": param_groups, + } + if opt_cfg.optimizer_type == "dmuon": + optimizer_kwargs["log_fn"] = self.log + + self.optimizer = get_optimizer( + opt_cfg.optimizer_type, + self.model, + **optimizer_kwargs, + ) + + def _create_scheduler(self): + """Create learning rate scheduler""" + from wall_x.config.hyperparams_config import ( + ConstantSchedulerConfig, + CosineSchedulerConfig, + ) + + sched = self.cfg.hyperparams.scheduler + lr = self.cfg.hyperparams.optimizer.learning_rate + + if isinstance(sched, ConstantSchedulerConfig): + self.lr_scheduler = torch.optim.lr_scheduler.ConstantLR( + self.optimizer, factor=1.0, total_iters=0 + ) + elif isinstance(sched, CosineSchedulerConfig): + if sched.num_training_steps <= 0: + raise ValueError( + "num_training_steps must be > 0 for cosine scheduler. " + "Please set it explicitly in the config." + ) + min_lr = sched.min_lr if sched.min_lr is not None else 0.1 * lr + self.lr_scheduler = get_scheduler( + optimizer=self.optimizer, + lr_scheduler_type="cosine", + num_warmup_steps=sched.num_warmup_steps, + num_training_steps=sched.num_training_steps, + peak_lr=lr, + end_lr=min_lr, + ) + else: + raise ValueError(f"Unsupported scheduler type: {type(sched).__name__}") + + def load_dataset(self): + """Load training and validation datasets""" + torch.cuda.empty_cache() + barrier() + # Dispatch to adapter with unified signature (supports model-specific kwargs) + extra_dataset_kwargs = dict( + model_config=self.model_config, + tokenizer_mixin=self.tokenizer_mixin, + normalizer_action=self.normalizer_action, + normalizer_propri=self.normalizer_propri, + ) + + resume_path = self.cfg.checkpoint.resume_from + if resume_path is not None and os.path.isdir(resume_path): + # Step and epoch restore are generic. Backend-owned episode + # metadata is optional and only loaded when the adapter exposes + # the hook. + is_incomplete_epoch = False + if hasattr(self.adapter, "_load_episode_indices"): + resume_indices = self.adapter._load_episode_indices( + resume_path, self.rank + ) + if resume_indices is not None: + is_incomplete_epoch = resume_indices["is_incomplete_epoch"] + extra_dataset_kwargs["resume_indices"] = resume_indices + + train_state = self.adapter.load_step_and_epoch( + resume_path, is_incomplete_epoch + ) + self.global_step = train_state["global_step"] + self.start_epoch = train_state["start_epoch"] + self.log( + f"global_step: {self.global_step}, start_epoch: " + f"{self.start_epoch} from {resume_path}" + ) + # global_step records completed steps, so resume from the next step + # to avoid triggering _should_save_checkpoint() immediately. + self.global_step += 1 + ( + self.dataset, + self.train_dataloader, + self.train_num, + ) = self.adapter.load_dataset( + self.data_config, + self.processor, + self.rank, + self.world_size, + **extra_dataset_kwargs, + ) + + def train_loop(self, epoch: int, profiler=contextlib.nullcontext()): + """Execute training for a single epoch""" + self.model.train() + + # VGDynamicRobotDataset (v1 path) doesn't expose set_epoch; other + # dataset backends (v2, lerobot) do. Dispatch conditionally. + if hasattr(self.dataset, "set_epoch"): + self.dataset.set_epoch(epoch) + total = len(self.train_dataloader) + # Drop last + stop_step = total - total % self.grad_accum_steps + + # Disable automatic GC to prevent random ~200ms stalls during compute + gc_interval = self.cfg.logging.gc_interval_steps + gc.disable() + + t0 = time.time() + with profiler: + self.timers("interval-time", log_level=0).start(barrier=True) + self.timers("data-load", log_level=0).start(barrier=True) + + for i, batch in enumerate(self.train_dataloader, self.initial_step): + self.timers("data-load").stop() + + # Save first batch for offline profiling (rank 0 only) + debug_batch_path = self.cfg.debug.save_debug_batch_path + if debug_batch_path and i == self.initial_step and self.rank == 0: + os.makedirs(os.path.dirname(debug_batch_path) or ".", exist_ok=True) + torch.save({"batch": batch}, debug_batch_path) + self.log(f"Saved debug batch to {debug_batch_path}") + + # Move batch to device + batch = move_batch_to_device(batch, self.device) + + if self.cfg.debug.enable_mfu_profile and self.global_step == 0: + self.log( + "[MFU] Forward FLOPs profiling is not available in the public package.", + level=logging.INFO, + ) + + # Forward pass with autocast + self.timers("forward-compute", log_level=0).start(barrier=False) + with self.autocast_context(): + outputs = self.adapter.forward( + self.model, batch, global_step=self.global_step + ) + self.timers("forward-compute").stop() + + loss = self.adapter.extract_loss(outputs) + + # Check for NaN loss -- replace with 0 instead of `continue` + # to avoid skipping collective ops (backward, all_reduce, etc.) + # which would cause NCCL deadlock across ranks. + nan_loss = torch.isnan(loss) + if nan_loss: + self.log( + f"Warning: nan in loss at epoch: {epoch}, step: {i}", + level=logging.WARNING, + ) + loss = torch.zeros_like(loss) + + # Backward pass + self.timers("backward-compute", log_level=0).start(barrier=True) + + context = ( + contextlib.nullcontext() + if self.sync_gradients() + else self.strategy.no_sync(self.model) + ) + with context: + scaled_loss = loss / self.grad_accum_steps + self.backward(scaled_loss) + self.timers("backward-compute").stop() + + # Gradient sync and optimizer step + if self.sync_gradients(): + self.timers("optimizer", log_level=0).start(barrier=True) + + # Per-component grad norms (before clipping) + self.timers("optimizer-grad-norms", log_level=0).start( + barrier=False + ) + self._component_grad_norms = {} + self.adapter.collect_grad_norms( + self.model, + self._component_grad_norms, + device=self.device, + reduce_tensor_fn=self.reduce_tensor, + params_sharded=self.strategy.params_sharded, + ) + self.timers("optimizer-grad-norms").stop() + + # Clip gradients + self.timers("optimizer-clip", log_level=0).start(barrier=False) + if self.cfg.hyperparams.optimizer.enable_grad_clip: + total_norm = self.clip_grad_norm(self.max_grad_norm) + else: + self._dedicated_param_grad_clip_stats = None + total_norm = 0 + self.timers("optimizer-clip").stop() + + self.timers("optimizer-step", log_level=0).start(barrier=False) + self.optimizer_step() + self.timers("optimizer-step").stop() + + self.timers("optimizer-zero-grad", log_level=0).start(barrier=False) + self.optimizer_zero_grad() + self.timers("optimizer-zero-grad").stop() + self.timers("optimizer").stop() + + # Scheduler step + self.lr_scheduler_step() + + # Logging + self.timers("logging", log_level=0).start(barrier=True) + _t_metrics_start = time.time() + self._log_training_metrics( + epoch, i, total, loss, total_norm, outputs + ) + _t_metrics_ms = (time.time() - _t_metrics_start) * 1000 + + t1 = time.time() + self.training_log( + epoch, + self.num_epoch, + i, + total, + loss, + self.get_lr(), + t1 - t0, + self.show_time_details, + ) + self.timers("logging").stop() + t0 = time.time() + + # Optional: log per-step breakdown to pinpoint spikes (e.g. param_norms every 100 steps) + if ( + self.show_time_details + and _t_metrics_ms > 5000 + and is_main_process() + ): + self.log( + f"[Step time breakdown] _log_training_metrics took {_t_metrics_ms:.0f} ms at global_step={self.global_step}", + level=logging.INFO, + ) + + # Checkpoint saving (FSDP full state_dict can take 10-20s every save_interval steps) + if self._should_save_checkpoint(): + _t_save_start = time.time() + self.save_checkpoint(epoch, self.global_step) + if is_main_process(): + self.log( + f"[Step time breakdown] save_checkpoint took {(time.time() - _t_save_start):.1f} s at global_step={self.global_step}", + level=logging.INFO, + ) + + # Validation + if self._should_validate(): + self.val_loop() + + self.global_step += 1 + self.micro_step = 0 + + # Manual GC outside timing window to avoid random stalls (can add 1-5s every gc_interval_steps) + if self.global_step % gc_interval == 0: + _t_gc_start = time.time() + gc.collect() + if is_main_process() and self.show_time_details: + self.log( + f"[Step time breakdown] gc.collect took {(time.time() - _t_gc_start):.1f} s at global_step={self.global_step}", + level=logging.INFO, + ) + else: + self.micro_step += 1 + + del batch + self.timers("interval-time").stop() + + # Drop last + if i == stop_step: + break + + if not isinstance(profiler, contextlib.nullcontext): + profiler.step() + + if ( + self.num_training_steps > 0 + and self.global_step >= self.num_training_steps + ): + break + + # Setup timers for next iteration + if i < total - 1: + self.timers("interval-time", log_level=0).start(barrier=True) + self.timers("data-load", log_level=0).start(barrier=True) + + # Re-enable automatic GC after training loop + gc.enable() + + # Reset dataloader for next epoch + self.train_dataloader = self.dataset.get_train_dataloader() + + def _log_training_metrics(self, epoch, step, total, loss, total_norm, outputs): + """Collect per-step stats, delegate buffering + wandb emission.""" + lr = self.get_lr() + train_loss = self.reduce_tensor(loss.detach()).item() + + step_stats = { + "lr": lr, + "train_loss": train_loss, + "grad_norm": ( + total_norm.item() if torch.is_tensor(total_norm) else float(total_norm) + ), + } + for idx, group in enumerate(self.optimizer.param_groups): + group_name = group.get("group_name", f"group_{idx}") + step_stats[f"lr_group/{group_name}"] = float(group["lr"]) + + # Model-family-specific auxiliary losses / accuracies. + self.adapter.collect_output_stats( + outputs, + step_stats, + reduce_tensor_fn=self.reduce_tensor, + true_gather_fn=self.true_gather, + tokenizer_mixin=self.tokenizer_mixin, + ) + + # Per-component grad norms (captured before clip & optimizer.zero_grad). + if hasattr(self, "_component_grad_norms"): + step_stats.update(self._component_grad_norms) + + dedicated_clip_stats = getattr(self, "_dedicated_param_grad_clip_stats", None) + if dedicated_clip_stats is not None: + step_stats.update( + { + "muon_grad_norm": dedicated_clip_stats["total_norm"], + "muon_grad_clip_coef": dedicated_clip_stats["clip_coef"], + "muon_grad_clipped": float(dedicated_clip_stats["clipped"]), + } + ) + + if self.global_step % 100 == 0: + self.adapter.collect_param_norms( + self.model, + step_stats, + device=self.device, + reduce_tensor_fn=self.reduce_tensor, + params_sharded=self.strategy.params_sharded, + ) + + self._current_step_raw_stats = step_stats + + # Display-smoothing rolling window (DZ-style) + self._current_step_stats = self.metrics_logger.smooth(step_stats) + + def training_log( + self, + current_epoch, + total_epoch, + current_train_iter, + total_train_iter, + loss, + lr, + time_per_step, + show_time_details=False, + ): + # timers.log() contains all_gather — must run on ALL ranks before the + # is_main_process() guard to avoid NCCL deadlock. + if show_time_details: + self.timers.log( + [ + "interval-time", + "data-load", + "forward-compute", + "backward-compute", + "optimizer", + "optimizer-grad-norms", + "optimizer-clip", + "optimizer-step", + "optimizer-zero-grad", + "logging", + ], + normalizer=1, + ) + + main = is_main_process() + stats = getattr(self, "_current_step_stats", None) or {} + + # MFU is computed after step time is known, then merged into both the + # console stats and the MetricsLogger buffer so wandb records it. + mfu_info = None + if self.cfg.debug.enable_mfu and hasattr(self.adapter, "compute_mfu"): + unwrapped = ( + self.model.module if hasattr(self.model, "module") else self.model + ) + mfu_info = self.adapter.compute_mfu(unwrapped, time_per_step) + if mfu_info is not None: + self.log( + "[MFU-rank] rank={} mfu={:.2f}% step={:.3f}s " + "flops_step={:.3f}T flops_fwd={:.3f}T " + "seq={} latent={} source={}".format( + getattr(self, "rank", 0), + mfu_info["mfu"] * 100.0, + time_per_step, + mfu_info["flops_per_step_T"], + mfu_info.get("flops_forward_T", 0.0) or 0.0, + mfu_info.get("seq_dims"), + mfu_info.get("latent_dims"), + mfu_info.get("latent_source"), + ), + main_process_only=False, + ) + if main and mfu_info is not None and stats is not None: + mfu_stats = { + "mfu": mfu_info["mfu"], + "flops_per_step_T": mfu_info["flops_per_step_T"], + "flops_forward_T": mfu_info.get("flops_forward_T"), + } + profile_fwd_flops = mfu_info.get("profile_fwd_flops_T") + if profile_fwd_flops is not None: + mfu_stats["profile_fwd_flops_T"] = profile_fwd_flops + stats.update(mfu_stats) + raw_stats = getattr(self, "_current_step_raw_stats", None) + if raw_stats is not None: + raw_stats.update(mfu_stats) + + if main: + raw_stats = getattr(self, "_current_step_raw_stats", None) + if raw_stats is not None: + self.metrics_logger.record_step(raw_stats, is_main=True) + avg_stats = self.metrics_logger.flush_if_due(self.global_step) + if avg_stats is not None: + self.logger.info( + f"[FSDP Train] Step {self.global_step}: {avg_stats}" + ) + + if not main: + return + + loss_to_print = stats.get( + "train_loss", loss.item() if torch.is_tensor(loss) else float(loss) + ) + fields = self.adapter.console_fields(tokenizer_mixin=self.tokenizer_mixin) + + line = self.metrics_logger.format_training_line( + epoch=current_epoch, + total_epoch=total_epoch, + current_iter=current_train_iter, + total_iter=total_train_iter, + loss=loss_to_print, + lr=lr, + time_per_step=time_per_step, + stats=stats, + fields=fields, + mfu_info=mfu_info, + ) + self.log(line) + + def _should_save_checkpoint(self) -> bool: + """Check if checkpoint should be saved""" + log = self.cfg.logging + return ( + self.global_step >= log.ignore_until_interval + and self.global_step % log.save_interval == 0 + and self.global_step != 0 + ) + + def _should_validate(self) -> bool: + """Check if validation should be performed""" + log = self.cfg.logging + return ( + self.global_step >= log.ignore_until_interval + and self.global_step % log.val_interval == 0 + and self.global_step > 0 + ) + + @torch.no_grad() + def val_loop(self): + """Delegate the full validation run to the adapter. + + Skips silently when the dataset has no val split — adapter's + run_validation iterates ``val_dataloader`` with ``tqdm(..., + total=len(...))`` and cannot accept None. v2 returns None here + when the YAML only declares a train split. + """ + self.val_dataloader = self.dataset.get_val_dataloader() + if self.val_dataloader is None: + if is_main_process(): + self.logger.info("No val split configured, skipping validation.") + return + save_path = self.cfg.checkpoint.save_path + self.adapter.run_validation( + model=self.model, + val_dataloader=self.val_dataloader, + rank=self.rank, + world_size=self.world_size, + device=self.device, + autocast_context=self.autocast_context, + reduce_fn=self.reduce_tensor, + gather_fn=self.true_gather, + logger=self.wandb_run if is_main_process() else None, + global_step=self.global_step, + output_path=os.path.join(save_path, f"val_rank_{self.rank}"), + tokenizer_mixin=self.tokenizer_mixin, + ) + barrier() + + def _compute_frozen_prefixes(self) -> Optional[Tuple[str, ...]]: + """Return state-dict key prefixes for submodules that are fully frozen. + + For models with a ``pipe`` container, any child not present in + ``cfg.model.trainable_models`` is treated as frozen. The matching key + prefixes let checkpoint_io drop those entries from the saved file. + Empty / missing ``trainable_models`` disables filtering and preserves + "save everything" behavior. + """ + if not hasattr(self.model, "pipe"): + return None + trainable_models = getattr(self.cfg.model, "trainable_models", None) + if trainable_models is None: + return None + if isinstance(trainable_models, str): + trainable_set = { + s.strip() for s in trainable_models.split(",") if s.strip() + } + else: + trainable_set = set(trainable_models) + if not trainable_set: + self.log( + "WARNING: trainable_models is empty; not filtering frozen " + "entries from checkpoint to avoid saving an empty file." + ) + return None + frozen_prefixes = tuple( + f"pipe.{name}." + for name, _ in self.model.pipe.named_children() + if name not in trainable_set + ) + if frozen_prefixes: + self.log( + f"Frozen submodule prefixes excluded from checkpoints: " + f"{list(frozen_prefixes)}" + ) + return frozen_prefixes + return None + + def save_checkpoint(self, epoch: int, step: int = 0): + """Save model checkpoint via checkpoint_io (dispatches on model wrapper type).""" + save_path = self.cfg.checkpoint.save_path + ckpt_path = f"{save_path}/{epoch}_{step}" if step else f"{save_path}/{epoch}" + _ckpt_io.save_checkpoint( + ckpt_path=ckpt_path, + model=self.model, + optimizer=self.optimizer, + lr_scheduler=self.lr_scheduler, + config=dataclasses.asdict(self.cfg), + rank=self.rank, + is_main=is_main_process(), + epoch=epoch, + global_step=self.global_step, + seed=self.seed, + normalizer_action=self.normalizer_action, + normalizer_propri=self.normalizer_propri, + dataset=getattr(self, "dataset", None) if step != 0 else None, + grad_scaler=self.grad_scaler, + log_fn=self.log, + frozen_prefixes=self._frozen_prefixes, + ) + _ckpt_io.finalize_save() + self.log(f"Saved checkpoint to {ckpt_path}") + + def load_state_dict(self, model, resume_config): + """Load state dict with fused-weight conversion + try_harder support.""" + return _ckpt_io.load_weights( + model=model, + resume_config=resume_config, + model_class=self._checkpoint_model_class(), + log_fn=self.log, + ) + + def _checkpoint_model_class(self): + model_class = getattr(self, "ModelClass", None) + if model_class is None and hasattr(type(self.adapter), "model_class"): + model_class = type(self.adapter).model_class() + return model_class + + def _resume_from_single_file(self) -> bool: + path = self.cfg.checkpoint.resume_from + return bool(path) and str(path).endswith((".safetensors", ".pth")) + + def _resume_from_training_checkpoint(self) -> bool: + path = self.cfg.checkpoint.resume_from + return bool(path) and not self._resume_from_single_file() + + def resume_from_checkpoint(self): + """Resume training from checkpoint via checkpoint_io.""" + _ckpt_io.resume_from_checkpoint( + model=self.model, + optimizer=self.optimizer, + lr_scheduler=self.lr_scheduler, + resume_config={"ckpt": self.cfg.checkpoint.resume_from}, + rank=self.rank, + grad_scaler=self.grad_scaler, + model_class=self._checkpoint_model_class(), + log_fn=self.log, + ) + barrier() + + def predict_action_loop( + self, + current_step=0, + max_iteration=None, + prediction_type="flow_action", + mode="acc", + ): + """Delegate action prediction to the adapter.""" + del mode # unused in current dispatch; kept for signature compatibility + return self.adapter.predict( + prediction_type, + model=self._unwrapped_model_for_inference(), + val_dataloader=self.dataset.get_val_dataloader(), + rank=self.rank, + world_size=self.world_size, + device=self.device, + processor=self.processor, + tokenizer_mixin=self.tokenizer_mixin, + logger=self.wandb_run, + current_step=current_step, + max_iteration=max_iteration, + ) + + def predict_text_loop(self, current_step=0, max_samples=None, save_dir="./"): + """Delegate text prediction to the adapter.""" + return self.adapter.predict( + "text", + model=self._unwrapped_model_for_inference(), + val_dataloader=self.dataset.get_val_dataloader(), + rank=self.rank, + world_size=self.world_size, + device=self.device, + processor=self.processor, + tokenizer_mixin=self.tokenizer_mixin, + logger=self.wandb_run, + current_step=current_step, + max_samples=max_samples, + save_dir=save_dir, + ) + + def _unwrapped_model_for_inference(self): + """Return the underlying module for inference: unwrap DDP, pass FSDP through.""" + if isinstance(self.model, DDP): + return self.model.module + return self.model diff --git a/wall_x/trainer/fsdp_trainer/metrics/__init__.py b/wall_x/trainer/fsdp_trainer/metrics/__init__.py new file mode 100644 index 00000000..31235735 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/metrics/__init__.py @@ -0,0 +1,3 @@ +from wall_x.trainer.fsdp_trainer.metrics.norms import DistributedNormAccumulator + +__all__ = ["DistributedNormAccumulator"] diff --git a/wall_x/trainer/fsdp_trainer/metrics/norms.py b/wall_x/trainer/fsdp_trainer/metrics/norms.py new file mode 100644 index 00000000..f5185341 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/metrics/norms.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +from collections import defaultdict +from typing import Callable, DefaultDict, Optional + +import torch +from torch.distributed.tensor import DTensor, Replicate + + +class DistributedNormAccumulator: + """Accumulate L2 norms without leaking parallel-layout branches to adapters.""" + + def __init__( + self, + *, + device: torch.device, + reduce_tensor_fn: Optional[Callable[[torch.Tensor, bool], torch.Tensor]], + ): + self.device = device + self.reduce_tensor_fn = reduce_tensor_fn + self._regular: DefaultDict[str, list[torch.Tensor]] = defaultdict(list) + self._dtensor: DefaultDict[str, list[torch.Tensor]] = defaultdict(list) + self._owner_only: DefaultDict[str, list[torch.Tensor]] = defaultdict(list) + + @staticmethod + def _is_dtensor(tensor: torch.Tensor) -> bool: + return isinstance(tensor, DTensor) + + def add(self, group: Optional[str], tensor: Optional[torch.Tensor]) -> None: + if group is None or tensor is None or tensor.numel() == 0: + return + tensor = tensor.detach() + if self._is_dtensor(tensor): + self._dtensor[group].append(tensor) + else: + self._regular[group].append(tensor) + + def add_owner_only( + self, group: Optional[str], tensor: Optional[torch.Tensor] + ) -> None: + if group is None or tensor is None or tensor.numel() == 0: + return + self._owner_only[group].append(tensor.detach()) + + def squared_norm(self, group: str) -> torch.Tensor: + with torch.no_grad(): + total = torch.zeros((), dtype=torch.float32, device=self.device) + total = total + self._regular_squared_norm(self._regular[group]) + total = total + self._owner_only_squared_norm(self._owner_only[group]) + total = total + self._dtensor_squared_norm(self._dtensor[group]) + return total + + def norm(self, group: str) -> torch.Tensor: + return self.squared_norm(group).sqrt() + + def _regular_squared_norm(self, tensors: list[torch.Tensor]) -> torch.Tensor: + if not tensors: + return torch.zeros((), dtype=torch.float32, device=self.device) + per = self._foreach_norm(tensors) + return self._sum_norms_squared(per) + + def _owner_only_squared_norm(self, tensors: list[torch.Tensor]) -> torch.Tensor: + local_sq = self._regular_squared_norm(tensors) + if self.reduce_tensor_fn is None: + return local_sq + return self.reduce_tensor_fn(local_sq, average=False) + + def _dtensor_squared_norm(self, tensors: list[torch.Tensor]) -> torch.Tensor: + if not tensors: + return torch.zeros((), dtype=torch.float32, device=self.device) + + per = self._foreach_norm(tensors) + + # Group by (mesh, placements) so each unique layout costs only one + # collective (redistribute). + groups: DefaultDict[object, list[torch.Tensor]] = defaultdict(list) + for norm in per: + if self._is_dtensor(norm): + key = (id(norm.device_mesh), tuple(norm.placements)) + else: + key = None + groups[key].append(norm) + + parts: list[torch.Tensor] = [] + for group_tensors in groups.values(): + try: + stacked = torch.stack(group_tensors) + local = self._to_local_replicated(stacked).to( + device=self.device, dtype=torch.float32 + ) + parts.append(local.pow(2).sum()) + except (RuntimeError, NotImplementedError): + # Stack/redistribute not supported for this layout in the + # current PyTorch version; fall back to per-tensor for safety. + for norm in group_tensors: + local = self._to_local_replicated(norm).to( + device=self.device, dtype=torch.float32 + ) + parts.append(local.pow(2).sum()) + + return torch.stack(parts).sum() + + def _foreach_norm(self, tensors: list[torch.Tensor]) -> list[torch.Tensor]: + try: + return torch._foreach_norm(tensors, 2.0) + except (RuntimeError, NotImplementedError): + return [torch.linalg.vector_norm(t, 2.0) for t in tensors] + + def _sum_norms_squared(self, norms: list[torch.Tensor]) -> torch.Tensor: + if not norms: + return torch.zeros((), dtype=torch.float32, device=self.device) + local_norms = [ + self._to_local_replicated(norm).to(device=self.device, dtype=torch.float32) + for norm in norms + ] + return torch.stack(local_norms).pow(2).sum() + + def _to_local_replicated(self, tensor: torch.Tensor) -> torch.Tensor: + if not self._is_dtensor(tensor): + return tensor.to(device=self.device) + mesh = tensor.device_mesh + tensor = tensor.redistribute(placements=[Replicate()] * mesh.ndim) + return tensor.to_local().to(device=self.device) diff --git a/wall_x/trainer/fsdp_trainer/metrics_logger.py b/wall_x/trainer/fsdp_trainer/metrics_logger.py new file mode 100644 index 00000000..ee00efe4 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/metrics_logger.py @@ -0,0 +1,111 @@ +"""Training metric buffering and wandb emission helpers.""" + +from __future__ import annotations + +from collections import deque +from typing import Any, Deque, Dict, List, Optional, Tuple + + +class MetricsLogger: + """Buffer per-step stats and emit averaged batches to wandb. + + Also formats the per-step console line via ``format_training_line``, + which consumes an ``(stat_key, pretty_label, fmt_spec)`` triple list + from ``adapter.console_fields()``. + + ``smooth_window`` controls a separate rolling-mean buffer used purely + for display smoothing (console + tqdm). It is independent of the + ``log_interval`` wandb-flush buffer; the smoothed dict is returned by + :meth:`smooth` for the caller to assign to ``_current_step_stats``. + 1 (default) preserves historical per-step values. + """ + + def __init__(self, *, wandb_run, log_interval: int, smooth_window: int = 1): + self._wandb_run = wandb_run + self._log_interval = log_interval + self._buffer: List[Dict[str, Any]] = [] + self._smooth_window = max(1, int(smooth_window)) + self._smooth_buffers: Dict[str, Deque[float]] = {} + + def record_step(self, step_stats: Dict[str, Any], *, is_main: bool) -> None: + """Push a per-step stats dict into the rank-0 buffer (no-op off rank-0).""" + if is_main: + self._buffer.append(step_stats) + + def flush_if_due(self, global_step: int) -> Optional[Dict[str, Any]]: + """At log_interval boundaries, average + emit; return avg dict or None.""" + if global_step % self._log_interval != 0 or not self._buffer: + return None + avg = self._average(self._buffer) + if self._wandb_run is not None and hasattr(self._wandb_run, "log"): + self._wandb_run.log(avg, step=global_step) + self._buffer = [] + return avg + + def smooth(self, stats: Dict[str, Any]) -> Dict[str, Any]: + """Return *stats* with numeric entries replaced by their rolling mean. + + Uses one ``deque(maxlen=smooth_window)`` per metric key. Non-numeric + entries pass through unchanged. When ``smooth_window <= 1`` this is + the identity (historical behavior). + """ + if self._smooth_window <= 1: + return stats + out: Dict[str, Any] = {} + for key, val in stats.items(): + if not isinstance(val, (int, float)): + out[key] = val + continue + buf = self._smooth_buffers.get(key) + if buf is None or buf.maxlen != self._smooth_window: + buf = deque(maxlen=self._smooth_window) + self._smooth_buffers[key] = buf + buf.append(float(val)) + out[key] = sum(buf) / len(buf) + return out + + @staticmethod + def _average(buf: List[Dict[str, Any]]) -> Dict[str, Any]: + all_keys: set = set() + for stats in buf: + all_keys.update(stats.keys()) + avg: Dict[str, Any] = {} + for key in all_keys: + values = [s[key] for s in buf if key in s] + if values: + avg[key] = sum(values) / len(values) + return avg + + def format_training_line( + self, + *, + epoch: int, + total_epoch: int, + current_iter: int, + total_iter: int, + loss: float, + lr: float, + time_per_step: float, + stats: Dict[str, Any], + fields: List[Tuple[str, str, str]], + mfu_info: Optional[Dict[str, Any]] = None, + ) -> str: + """Render a line matching pre-refactor training_log output exactly. + + Layout: " epoch E/T | iter I/T | loss L | [adapter fields] | lr L | + time_current_backward_step Ts | [MFU P% |]" + """ + parts = [ + " epoch {:3d}/{:3d} |".format(epoch, total_epoch), + " iter {:6d}/{:6d} |".format(current_iter, total_iter), + " loss {:.6f} |".format(loss), + ] + for key, label, fmt in fields: + value = stats.get(key) + if value is not None: + parts.append(" {} {:{fmt}} |".format(label, value, fmt=fmt)) + parts.append(" lr {:.6f} |".format(lr)) + parts.append(" time_current_backward_step {:.6f}s |".format(time_per_step)) + if mfu_info is not None: + parts.append(" MFU {:.2f}% |".format(mfu_info["mfu"] * 100)) + return "".join(parts) diff --git a/wall_x/trainer/fsdp_trainer/train_fsdp.py b/wall_x/trainer/fsdp_trainer/train_fsdp.py new file mode 100644 index 00000000..ea971239 --- /dev/null +++ b/wall_x/trainer/fsdp_trainer/train_fsdp.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 +""" +FSDP Training Entry Point +Launch with torchrun: + torchrun --nproc_per_node=8 --master_port=29500 train_fsdp.py --config config.yaml +""" +import argparse +import dataclasses +import logging +import os +import sys +import torch +from datetime import datetime +from wall_x.config.loader import load_config as load_typed_config +from wall_x.trainer.fsdp_trainer import ( + FSDPTrainer, + cleanup_distributed, + is_main_process, +) +import wandb + +logger = logging.getLogger(__name__) + + +class TeeOutput: + """Write output to both the terminal and a file.""" + + def __init__(self, file_path, mode="w"): + self.terminal = sys.stdout + self.log = open(file_path, mode) + + def write(self, message): + self.terminal.write(message) + self.log.write(message) + self.log.flush() + + def flush(self): + self.terminal.flush() + self.log.flush() + + +def parse_args(): + parser = argparse.ArgumentParser(description="FSDP Training Script") + parser.add_argument( + "--config", + type=str, + required=True, + help="Path to training config YAML file", + ) + # FSDP specific overrides + parser.add_argument( + "--fsdp_sharding_strategy", + type=str, + default=None, + choices=["full_shard", "shard_grad_op", "no_shard", "hybrid_shard"], + help="FSDP sharding strategy (overrides config)", + ) + parser.add_argument( + "--debug", + action="store_true", + default=False, + help="Enable debug mode: reduces buffer size, sets log/save path to debug.", + ) + parser.add_argument( + "--wandb_offline", + type=str, + default=None, + help="Whether to run wandb in offline mode (overrides config).", + ) + parser.add_argument( + "--visualize", + action="store_true", + default=False, + help="Whether to visualize samples during training.", + ) + parser.add_argument( + "--log_to_file", + action="store_true", + default=False, + help="Whether to redirect stdout and stderr to log files.", + ) + return parser.parse_args() + + +def load_config(config_path: str, cli_args=None): + """Load configuration from YAML file into TrainConfig.""" + return load_typed_config(config_path, cli_args=cli_args) + + +def setup_logger(cfg): + """Setup wandb logger if enabled""" + log = cfg.logging + if log.use_wandb and is_main_process(): + logger.info( + "rank %s is initializing wandb , is main process %s", + torch.distributed.get_rank(), + is_main_process(), + ) + wandb_run = wandb.init( + project=log.log_project, + name=log.log_name, + entity=log.log_entity, + config=dataclasses.asdict(cfg), + save_code=False, + force=False, + mode="offline" if log.wandb_offline else "online", + ) + logger.info("Wandb Initialized") + return wandb_run + return None + + +def print_fsdp_config(cfg): + """Print FSDP configuration""" + if is_main_process(): + dist = cfg.distributed + logger.info("%s", "=" * 60) + logger.info("FSDP Configuration:") + logger.info(" use_fsdp: %s", dist.use_fsdp) + logger.info("%s", "=" * 60) + + +def main(): + args = parse_args() + + # Redirect logs to files only on the main process to avoid write races. + if args.log_to_file: + import yaml as _yaml + + with open(args.config, "r") as f: + _tmp = _yaml.safe_load(f) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + # support both old (save_path) and new (checkpoint.save_path) schema + log_dir = _tmp.get("save_path") or _tmp.get("checkpoint", {}).get( + "save_path", "./ckpt" + ) + if args.debug: + log_dir = "./ckpt/debug" + os.makedirs(log_dir, exist_ok=True) + log_file = os.path.join(log_dir, f"training_log_{timestamp}.log") + sys.stdout = TeeOutput(log_file, mode="w") + sys.stderr = TeeOutput(log_file.replace(".log", "_stderr.log"), mode="w") + logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) + logger.info("\n%s", "=" * 80) + logger.info("LOG TO FILE MODE: All output will be saved to:") + logger.info(" STDOUT: %s", log_file) + logger.info(" STDERR: %s", log_file.replace(".log", "_stderr.log")) + logger.info("%s\n", "=" * 80) + else: + logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout) + + torch.cuda.init() + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) + torch.distributed.init_process_group("nccl", device_id=device) + cfg = load_config(args.config, cli_args=args) + wandb_run = setup_logger(cfg) + print_fsdp_config(cfg) + + try: + trainer = FSDPTrainer( + train_config=cfg, + wandb_run=wandb_run, + ) + trainer.fit() + + except Exception as e: + logger.exception("Training failed with error: %s", e) + raise + + finally: + cleanup_distributed() + if wandb_run is not None: + wandb_run.finish() + + +if __name__ == "__main__": + main() diff --git a/wall_x/trainer/optimizer/__init__.py b/wall_x/trainer/optimizer/__init__.py new file mode 100644 index 00000000..54e8a63a --- /dev/null +++ b/wall_x/trainer/optimizer/__init__.py @@ -0,0 +1,9 @@ +from .utils import get_optimizer, register_optimizer + +# DMuon is optional (external package). Skip registration if dmuon isn't +# installed so wall_x still imports; the registry will simply not have +# "dmuon" and get_optimizer("dmuon", ...) will raise a clear KeyError. +try: + from . import dmuon # noqa: F401 — side-effect registers "dmuon" +except ImportError: + pass diff --git a/wall_x/trainer/optimizer/dmuon/__init__.py b/wall_x/trainer/optimizer/dmuon/__init__.py new file mode 100644 index 00000000..b3d637c0 --- /dev/null +++ b/wall_x/trainer/optimizer/dmuon/__init__.py @@ -0,0 +1,16 @@ +import torch.nn as nn + +from .utils import get_dmuon_optimizer + + +def is_dmuon_model(model: nn.Module) -> bool: + """True if ``dmuon.dedicate_params()`` has been applied to this model. + + Checked via an attribute the external ``dmuon`` package attaches to + the root module, so this predicate works without importing ``dmuon`` + and returns ``False`` for ordinary (non-DMuon) models. + """ + return hasattr(model, "_dedicated_comm_ctx") + + +__all__ = ["get_dmuon_optimizer", "is_dmuon_model"] diff --git a/wall_x/trainer/optimizer/dmuon/utils.py b/wall_x/trainer/optimizer/dmuon/utils.py new file mode 100644 index 00000000..3d03b826 --- /dev/null +++ b/wall_x/trainer/optimizer/dmuon/utils.py @@ -0,0 +1,156 @@ +"""DMuon optimizer builder.""" + +import inspect +import logging + +from ..utils import register_optimizer + +_logger = logging.getLogger(__name__) + + +def _emit(log_fn, message, *args, level=logging.INFO): + if args: + message = message % args + if log_fn is not None: + log_fn(message, level=level) + else: + _logger.log(level, message) + + +def _is_rank0(): + try: + import torch.distributed as dist + + return ( + not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0 + ) + except Exception: + return True + + +def _build_ns_backend(dmuon, opt_cfg): + coefficients = getattr(opt_cfg, "ns_coefficients", "default") + if coefficients in (None, "default"): + return opt_cfg.ns_backend + + if coefficients != "wallx_muon": + raise ValueError( + "Unsupported DMuon ns_coefficients=" + f"{coefficients!r}. Supported: 'default', 'wallx_muon'." + ) + if opt_cfg.ns_backend != "direct": + raise ValueError( + "ns_coefficients='wallx_muon' is intended to match Wall-X's " + "direct-space Muon implementation. Set ns_backend='direct'." + ) + + wallx_coefficients = [[3.4445, -4.7750, 2.0315] for _ in range(opt_cfg.ns_steps)] + return dmuon.NewtonSchulz( + backend="direct", + coefficients=wallx_coefficients, + ) + + +def get_dmuon_optimizer(model, *, opt_cfg, param_groups=None, log_fn=None): + """Build dmuon.Muon from a DMuonConfig. + + When ``param_groups`` is provided, Wall-X expects DMuon to preserve the + PyTorch optimizer group semantics and then split each user group into + dedicated/Muon and non-dedicated/AdamW subgroups internally. + """ + from wall_x.config.hyperparams_config import DMuonConfig + + if not isinstance(opt_cfg, DMuonConfig): + raise TypeError( + f"get_dmuon_optimizer expects DMuonConfig, got {type(opt_cfg).__name__}" + ) + import dmuon + + muon_signature = inspect.signature(dmuon.Muon) + supports_param_groups = "param_groups" in muon_signature.parameters + if param_groups is not None and not supports_param_groups: + raise RuntimeError( + "Wall-X built optimizer param_groups for DMuon, but the installed " + "dmuon.Muon does not accept a param_groups= argument. Please update " + "DMuon to the param-group-aware implementation before enabling " + "action_expert_learning_rate with optimizer_type='dmuon'." + ) + + ns_backend = _build_ns_backend(dmuon, opt_cfg) + + _emit( + log_fn, + "DMuon: Muon lr=%s momentum=%s ns_steps=%s; " + "AdamW lr=%s betas=%s wd=%s; " + "ns_backend=%s ns_coefficients=%s nesterov=%s", + opt_cfg.muon_lr, + opt_cfg.momentum, + opt_cfg.ns_steps, + opt_cfg.adamw_lr, + opt_cfg.adamw_betas, + opt_cfg.adamw_weight_decay, + opt_cfg.ns_backend, + opt_cfg.ns_coefficients, + opt_cfg.nesterov, + ) + if param_groups is not None: + _emit( + log_fn, + "DMuon param_groups enabled: %s", + [ + { + "group_name": group.get("group_name", f"group_{idx}"), + "lr": group.get("lr"), + "num_params": len(group.get("params", [])), + } + for idx, group in enumerate(param_groups) + ], + ) + + kwargs = {} + if param_groups is not None: + kwargs["param_groups"] = param_groups + + optimizer = dmuon.Muon( + model, + lr=opt_cfg.muon_lr, + momentum=opt_cfg.momentum, + weight_decay=opt_cfg.muon_weight_decay, + ns_steps=opt_cfg.ns_steps, + adamw_lr=opt_cfg.adamw_lr, + adamw_betas=tuple(opt_cfg.adamw_betas), + adamw_weight_decay=opt_cfg.adamw_weight_decay, + adamw_eps=opt_cfg.adamw_eps, + ns_backend=ns_backend, + nesterov=opt_cfg.nesterov, + **kwargs, + ) + + if param_groups is not None and _is_rank0(): + summarize = getattr(dmuon, "summarize_param_groups", None) + format_summary = getattr(dmuon, "format_param_group_summary", None) + if summarize is None or format_summary is None: + _emit( + log_fn, + "DMuon param_groups are enabled, but the installed DMuon package " + "does not expose param-group diagnostics. Update DMuon if you need " + "startup verification of the Muon/AdamW subgroup split.", + level=logging.WARNING, + ) + else: + try: + summary = summarize(model, optimizer, max_rows=80) + _emit(log_fn, "%s", format_summary(summary)) + except Exception as exc: + _logger.exception("Failed to summarize DMuon param_groups") + _emit( + log_fn, + "Failed to summarize DMuon param_groups: %s", + exc, + level=logging.WARNING, + ) + + return optimizer + + +register_optimizer("dmuon", get_dmuon_optimizer) diff --git a/wall_x/trainer/optimizer/utils.py b/wall_x/trainer/optimizer/utils.py new file mode 100644 index 00000000..eb1c17c9 --- /dev/null +++ b/wall_x/trainer/optimizer/utils.py @@ -0,0 +1,230 @@ +import inspect + +from torch.optim import AdamW + +from wall_x.config.hyperparams_config import LRGroupConfig + +_OPTIMIZERS = {} + + +def register_optimizer(name, optimizer_fn): + _OPTIMIZERS[name] = optimizer_fn + + +def get_optimizer(name, *args, **kwargs): + if name not in _OPTIMIZERS: + raise KeyError(f"Unknown optimizer '{name}'. Registered: {sorted(_OPTIMIZERS)}") + return _OPTIMIZERS[name](*args, **kwargs) + + +def _group_weight_decay(opt_cfg): + return getattr(opt_cfg, "weight_decay", None) + + +def resolve_lr_group_configs(opt_cfg, default_action_lr_keywords): + """Return structured LR groups, with legacy action config fallback. + + ``optimizer.lr_groups`` is the preferred path. The legacy + ``action_expert_learning_rate`` fields are still converted into a single + action group so older configs keep working. + """ + if opt_cfg.lr_groups: + if ( + opt_cfg.action_expert_learning_rate is not None + or opt_cfg.action_lr_keywords is not None + ): + raise ValueError( + "Use either optimizer.lr_groups or legacy " + "action_expert_learning_rate/action_lr_keywords, not both." + ) + return opt_cfg.lr_groups + + if opt_cfg.action_expert_learning_rate is None: + return [] + + action_lr_keywords = ( + opt_cfg.action_lr_keywords + if opt_cfg.action_lr_keywords is not None + else default_action_lr_keywords + ) + return [ + LRGroupConfig( + name="action_lr_group", + lr=opt_cfg.action_expert_learning_rate, + include=action_lr_keywords, + fail_on_empty=True, + ) + ] + + +def uses_legacy_action_lr_groups(opt_cfg) -> bool: + return not opt_cfg.lr_groups and opt_cfg.action_expert_learning_rate is not None + + +def _make_param_group(name, params, lr, opt_cfg): + group = { + "params": params, + "lr": lr, + "group_name": name, + } + weight_decay = _group_weight_decay(opt_cfg) + if weight_decay is not None: + group["weight_decay"] = weight_decay + return group + + +def _validate_lr_group(group: LRGroupConfig, *, base_group_name: str): + if not group.name: + raise ValueError("optimizer.lr_groups entries must have a non-empty name") + if group.name == base_group_name: + raise ValueError( + f"optimizer.lr_groups name {group.name!r} is reserved for the base group" + ) + if "/" in group.name: + raise ValueError( + f"optimizer.lr_groups name {group.name!r} must not contain '/'. " + "DMuon appends '/muon' and '/adamw' to group names." + ) + if not group.include: + raise ValueError( + f"optimizer.lr_groups.{group.name} must define at least one include keyword" + ) + + +def build_lr_param_groups(model, opt_cfg, lr_groups, *, base_group_name="base"): + """Split trainable params into named LR groups plus a base group. + + ``lr_groups`` is a list of :class:`LRGroupConfig`. Each group matches + parameter names by substring. A parameter may match at most one explicit + group; unmatched trainable parameters remain in the ``base`` group using + ``opt_cfg.learning_rate``. + + Returns a list of torch.optim-compatible param_group dicts. The + ``group_name`` key is non-standard but preserved by torch.optim via + ``setdefault`` in ``add_param_group`` and is consumed downstream for + per-group lr logging. + + For AdamW / native Muon, each returned group includes ``weight_decay``. + For DMuon, weight decay is split between Muon and AdamW defaults, so the + groups only carry lr and metadata; DMuon applies its own per-route defaults. + The caller is expected to gate on ``opt_cfg.optimizer_type`` before calling + this. + """ + if not lr_groups: + return None + + names = [group.name for group in lr_groups] + duplicate_names = sorted({name for name in names if names.count(name) > 1}) + if duplicate_names: + raise ValueError( + f"optimizer.lr_groups contains duplicate names: {duplicate_names}" + ) + + for group in lr_groups: + _validate_lr_group(group, base_group_name=base_group_name) + + base_params = [] + grouped_params = {group.name: [] for group in lr_groups} + ambiguous = [] + + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + matches = [ + group.name + for group in lr_groups + if any(keyword in name for keyword in group.include) + ] + if len(matches) > 1: + ambiguous.append((name, matches)) + continue + if matches: + grouped_params[matches[0]].append(param) + else: + base_params.append(param) + + if ambiguous: + examples = ", ".join(f"{name} -> {matches}" for name, matches in ambiguous[:10]) + raise ValueError( + "Some parameters match multiple optimizer.lr_groups. Make group " + f"include patterns disjoint. Examples: {examples}" + ) + + if opt_cfg.train_action_expert_only: + assert len(base_params) == 0, ( + f"Expected 0 base_params after pre-wrap freeze, got {len(base_params)}. " + "Ensure base params are frozen before building the optimizer." + ) + + param_groups = [] + if len(base_params) > 0: + param_groups.append( + _make_param_group( + base_group_name, base_params, opt_cfg.learning_rate, opt_cfg + ) + ) + + for group in lr_groups: + params = grouped_params[group.name] + if len(params) == 0: + if group.fail_on_empty: + raise ValueError( + f"No params found for optimizer.lr_groups.{group.name}. " + f"Please check include={group.include!r}." + ) + continue + param_groups.append(_make_param_group(group.name, params, group.lr, opt_cfg)) + + return param_groups + + +def build_action_expert_param_groups(model, opt_cfg, action_lr_keywords): + """Compatibility wrapper for the legacy action-expert LR config.""" + return build_lr_param_groups( + model, + opt_cfg, + [ + LRGroupConfig( + name="action_lr_group", + lr=opt_cfg.action_expert_learning_rate, + include=action_lr_keywords, + fail_on_empty=True, + ) + ], + base_group_name="base_lr_group", + ) + + +def get_adamw_optimizer(model, *, opt_cfg, param_groups=None): + """Build AdamW from AdamWConfig.""" + from wall_x.config.hyperparams_config import AdamWConfig + + if not isinstance(opt_cfg, AdamWConfig): + raise TypeError( + f"get_adamw_optimizer expects AdamWConfig, got {type(opt_cfg).__name__}" + ) + + if param_groups is None: + params = [p for p in model.parameters() if p.requires_grad] + else: + # Per-group lr / weight_decay are preserved by torch.optim (via + # setdefault in add_param_group), so top-level values act only as + # defaults. Extra keys like ``group_name`` are kept in-place and used + # downstream for per-group lr logging. + params = param_groups + + kw = { + "lr": opt_cfg.learning_rate, + "weight_decay": opt_cfg.weight_decay, + "betas": tuple(opt_cfg.betas), + "eps": opt_cfg.eps, + } + sig_params = inspect.signature(AdamW.__init__).parameters + if "foreach" in sig_params and opt_cfg.foreach is not None: + kw["foreach"] = opt_cfg.foreach + if "fused" in sig_params: + kw["fused"] = opt_cfg.fused + return AdamW(params, **kw) + + +register_optimizer("adamw", get_adamw_optimizer) diff --git a/wall_x/trainer/qwen_vl_act_trainer.py b/wall_x/trainer/qwen_vl_act_trainer.py deleted file mode 100644 index c981e246..00000000 --- a/wall_x/trainer/qwen_vl_act_trainer.py +++ /dev/null @@ -1,1144 +0,0 @@ -import os -import gc -import time -import yaml -import shutil -import torch -import random -import numpy as np -import torch.nn as nn -import torch.distributed as dist -import json -from tqdm import tqdm -from functools import wraps -from datetime import datetime -from torch.optim import AdamW -from torch.distributed.tensor import distribute_tensor -from accelerate import Accelerator -from safetensors.torch import load_file -from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup -from transformers import AutoProcessor -from wall_x.model.action_head import Normalizer -from wall_x.utils.timers import Timers -from wall_x.model.qwen2_5_based import Qwen2_5_VLMoEForAction, Qwen2_5_VLConfig -from wall_x.utils.constant import action_statistic_dof as default_action_statistic_dof -from wall_x.data.config import ACTION_DATASET_NAMES, MULTIMODAL_DATASET_NAMES -from wall_x.data.load_lerobot_dataset import ( - PreprocessedDataset, - get_data_configs, - load_lerobot_data, -) -import copy - - -def timer(func): - """ - Decorator to measure function execution time. - - Args: - func: Function to be timed - - Returns: - Wrapped function with timing functionality - """ - - @wraps(func) - def wrapper(*args, **kwargs): - start_time = time.time() - result = func(*args, **kwargs) - end_time = time.time() - print( - f"\033[92m[current time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] Function {func.__name__} took {end_time - start_time:.2f} seconds to execute\033[0m" - ) - return result - - return wrapper - - -def print_rank_last(message): - """ - Print message only on the last rank in distributed training. - - Args: - message (str): Message to print - """ - if torch.distributed.is_initialized(): - if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1): - print(message, flush=True) - else: - print(message, flush=True) - - -def seed_all(seed): - """ - Set random seeds for reproducible training. - - Args: - seed (int): Random seed value - """ - np.random.seed(seed) - random.seed(seed) - torch.manual_seed(seed) - - -def update_model_config(train_config, model_config): - model_config.use_state_string_representation = train_config["data"].get( - "use_state_string_representation", False - ) - model_config.flow_loss_weight = train_config.get("flow_loss_weight", 1.0) - - model_config.dof_config = train_config["dof_config"] - model_config.agent_pos_config = train_config["agent_pos_config"] - - model_config.action_horizon_flow = train_config["data"].get( - "action_horizon_flow", 32 - ) - - if train_config.get("_attn_implementation", None) is not None: - model_config._attn_implementation = train_config["_attn_implementation"] - - if train_config.get("attn_deterministic", None) is not None: - model_config.attn_deterministic = train_config["attn_deterministic"] - model_config.vision_config.attn_deterministic = train_config[ - "attn_deterministic" - ] - print("[DEBUG] Attention is using deterministic kernel for this run!") - else: - model_config.attn_deterministic = False - model_config.vision_config.attn_deterministic = False - - return model_config - - -class QwenVlAct_Trainer: - """ - Vision-Language-Action trainer for Qwen-VL models with robotic action prediction. - - This trainer handles multi-modal learning combining vision, language, and action data - for robotic control applications. It supports distributed training, mixed precision, - gradient accumulation, and various optimization strategies including MoE (Mixture of Experts). - - Features: - - Multi-modal data processing (vision + language + actions) - - Distributed training with Accelerate - - Gradient accumulation and clipping - - Learning rate scheduling with warmup - - Checkpoint saving and resuming - - Comprehensive logging and monitoring - """ - - @timer - def __init__( - self, - config, - logger, - accelerator: Accelerator = None, - seed=42, - data_config_path=None, - ): - """ - Initialize the Vision-Language-Action trainer. - - Args: - config (dict): Training configuration dictionary containing: - - processor_path (str): Path to data preprocessing processor - - qwen_vl_act_config_path (str): Path to model configuration file - - learning_rate (float): Base learning rate for training - - num_epoch (int): Number of training epochs - - pretrained_wallx_path (str): Path to pretrained model - - And other training hyperparameters - logger: Logger instance for tracking metrics - accelerator (Accelerator, optional): Hugging Face Accelerate instance for distributed training - seed (int, optional): Random seed for reproducibility. Defaults to 42. - data_config_path (str, optional): Path to data configuration file - - Raises: - ValueError: If required configuration keys are missing - """ - # Validate required configuration keys - required_keys = ["learning_rate", "num_epoch"] - for key in required_keys: - if key not in config: - raise ValueError(f"Missing required configuration key: {key}") - - self.config = config - self.logger = logger - self.accelerator = accelerator - self.seed = seed - - # Initialize random seeds for reproducibility - seed_all(self.seed) - - # Training state variables - self.start_epoch = 0 - self.global_step = 0 - self.num_epoch = self.config["num_epoch"] - self.initial_step = 0 - - # Data and model configuration - self.dataload_config = get_data_configs(self.config["data"]) - self.data_config_path = data_config_path - self.use_fast_tokenizer = self.config.get("use_fast_tokenizer", False) - self.use_selective_recompute = self.config.get("use_selective_recompute", False) - - # Load model and initialize training components - self.load_normalizer() - self.load_model() - self.action_dim = sum(self.config["dof_config"].values()) - - # Distributed training setup - self.rank = self.accelerator.process_index - self.world_size = self.accelerator.num_processes - print( - f"rank {self.accelerator.process_index} after load model memory usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB", - flush=True, - ) - - # Load training data - self.load_qact_data() - print( - f"rank {self.accelerator.process_index} after load qact data usage: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB", - flush=True, - ) - - # Resume from checkpoint if specified - if "resume" in self.config: - self.resume_from_checkpoint() - - # Initialize special token IDs - self.propri_token_id = self.processor.tokenizer.convert_tokens_to_ids( - "<|propri|>" - ) - self.action_token_id = self.processor.tokenizer.convert_tokens_to_ids( - "<|action|>" - ) - - # Initialize evaluation metrics - self.base_l1_loss = None - self.base_l1_loss_detail = {} - - # Performance monitoring - self.timers = Timers(log_level=0, log_option="minmax") - - # Adjust global step if resuming from checkpoint - if self.initial_step != 0: - self.global_step = self.initial_step // self.config.get( - "gradient_accumulation_steps", 1 - ) - - def load_normalizer(self): - if self.config.get("norm_stats_path", None): - self.print_rank0( - f"loading customized action statistic dof from {self.config['norm_stats_path']}" - ) - action_statistic_dof = json.load(open(self.config["norm_stats_path"], "r")) - else: - self.print_rank0( - "loading default action statistic dof from default_action_statistic_dof" - ) - action_statistic_dof = default_action_statistic_dof - - self.normalizer_action = Normalizer( - action_statistic_dof, - self.config["dof_config"], - min_key=self.config.get("min_key", "min"), - delta_key=self.config.get("delta_key", "delta"), - ) - - print("self.normalizer_action.min: ", self.normalizer_action) - self.normalizer_propri = Normalizer( - action_statistic_dof, - self.config["agent_pos_config"], - min_key=self.config.get("min_key", "min"), - delta_key=self.config.get("delta_key", "delta"), - ) - - def print_rank0(self, msg, flush=True): - """ - Print message only on rank 0 (main process). - - Args: - msg: Message to print - flush (bool): Whether to flush output buffer - """ - if self.accelerator.is_main_process: - print(msg, flush=flush) - - def fit(self): - """ - Main training loop executing multiple epochs with validation. - - Handles the complete training process including: - - Training loop execution - - Validation after each epoch - - Process synchronization - - Memory cleanup - """ - self.accelerator.wait_for_everyone() - # Optional validation before training starts - if self.config.get("resume", None) is not None and self.config["resume"].get( - "validate_first", False - ): - self.val_loop() - self.accelerator.wait_for_everyone() - - # Main training loop - for epoch in range(self.start_epoch, self.num_epoch): - self.train_loop(epoch) - self.accelerator.wait_for_everyone() - - if (epoch + 1) % self.config.get("epoch_save_interval", 1) == 0: - self.save_checkpoint(epoch) - - # Validation after each epoch - # self.val_loop() - self.accelerator.wait_for_everyone() - - # Memory cleanup - gc.collect() - - def train_loop(self, epoch): - """ - Execute training for a single epoch. - - Args: - epoch (int): Current epoch number - - Handles: - - Data loading and batching - - Forward/backward passes - - Gradient accumulation and clipping - - Learning rate scheduling - - Loss logging and monitoring - - Performance profiling (optional) - """ - # Initialize training dataloader for current epoch - if isinstance(self.dataset, PreprocessedDataset): - if getattr(self, "train_dataloader", None) is not None: - self.dataset._train() - self.train_sampler.set_epoch(epoch) - else: - self.train_dataloader, self.train_sampler = ( - self.dataset.get_train_dataloader() - ) - self.train_sampler.set_epoch(epoch) - else: - self.train_dataloader = self.dataset.get_train_dataloader() - - self.model.train() - total = len(self.train_dataloader) - t0 = time.time() - enable_profiling = self.config["profile"] - - # Optional PyTorch profiler for performance analysis - if enable_profiling: - profiler = torch.profiler.profile( - activities=[ - torch.profiler.ProfilerActivity.CPU, - torch.profiler.ProfilerActivity.CUDA, - ], - schedule=torch.profiler.schedule( - wait=self.config["profile_wait_iters"], - warmup=self.config["profile_warmup_iters"], - active=self.config["profile_active_iters"], - ), - on_trace_ready=torch.profiler.tensorboard_trace_handler( - self.config["profile_save_path"], worker_name="worker0" - ), - record_shapes=True, - profile_memory=True, - with_stack=True, - ) - profiler.__enter__() - - try: - - # Setup timers for First iteration - self.timers("interval-time", log_level=0).start(barrier=False) - self.timers("data-load", log_level=0).start(barrier=False) - - for i, batch in enumerate(self.train_dataloader, self.initial_step): - # Move batch to device - if isinstance(self.dataset, PreprocessedDataset): - batch = { - k: ( - v.to(self.accelerator.device, non_blocking=True) - if isinstance(v, torch.Tensor) - else v - ) - for k, v in batch.items() - } - - self.timers("data-load").stop() - - with self.accelerator.accumulate(self.model): - # Forward pass - self.timers("forward-compute", log_level=0).start(barrier=False) - outputs = self.model(**batch, mode="train") - self.timers("forward-compute").stop() - - loss = outputs.loss - # Check for NaN loss - if torch.isnan(loss): - print( - f"Warning: NaN loss detected in epoch: {epoch}, step: {i}", - flush=True, - ) - continue - - # Backward pass - self.timers("backward-compute", log_level=0).start(barrier=False) - self.accelerator.backward(loss) - self.timers("backward-compute").stop() - - # Gradient clipping - total_norm = self.accelerator.clip_grad_norm_( - self.model.parameters(), self.config.get("max_grad_norm", 1.0) - ) - - # Optimizer step - self.timers("optimizer", log_level=0).start(barrier=False) - self.optimizer.step() - self.optimizer.zero_grad() - self.timers("optimizer").stop() - - # Update global step and learning rate after gradient accumulation - if self.accelerator.sync_gradients: - self.lr_scheduler.step() - self.global_step += 1 - lr = self.lr_scheduler.get_last_lr()[0] - - # Gather loss across all processes for logging - train_loss = ( - self.accelerator.gather(loss.detach()).mean().item() - ) - _log_dict = { - "lr": lr, - "train_loss": train_loss, - } - - # Log component losses - if ( - "cross_entropy_loss" in outputs - and outputs.cross_entropy_loss is not None - ): - _log_dict["cross_entropy_loss"] = ( - self.accelerator.gather( - outputs.cross_entropy_loss.detach() - ) - .mean() - .item() - ) - - if "flow_loss" in outputs and outputs.flow_loss is not None: - _log_dict["flow_loss"] = ( - self.accelerator.gather(outputs.flow_loss.detach()) - .mean() - .item() - ) - - # Log per-dataset channel losses - if ( - "channel_loss_dict" in outputs - and outputs.channel_loss_dict is not None - ): - for dataset_name_i in ( - ACTION_DATASET_NAMES + MULTIMODAL_DATASET_NAMES - ): - count_sum = ( - self.accelerator.gather( - outputs.channel_loss_count_dict[dataset_name_i] - ) - .sum() - .item() - ) - if count_sum > 0: - channel_loss = ( - self.accelerator.gather( - outputs.channel_loss_dict[ - dataset_name_i - ].detach() - ) - .sum() - .item() - / count_sum - ) - _log_dict[f"channel_loss_{dataset_name_i}"] = ( - channel_loss - ) - - # Log action accuracy for fast tokenizer - if ( - "action_accuracy" in outputs.channel_loss_dict - and self.use_fast_tokenizer - ): - _log_dict["action_accuracy"] = ( - self.accelerator.gather( - outputs.channel_loss_dict[ - "action_accuracy" - ].detach() - ) - .mean() - .item() - ) - - # Log metrics - if self.logger is not None: - self.logger.log(_log_dict, step=self.global_step) - - # Log gradient norm - if self.logger is not None and self.accelerator.sync_gradients: - self.logger.log( - {"total_norm": total_norm}, step=self.global_step - ) - - self.timers("interval-time").stop() - - # Setup timers for next iteration - if i < len(self.train_dataloader) - 1: - self.timers("interval-time", log_level=0).start(barrier=False) - self.timers("data-load", log_level=0).start(barrier=False) - - # Periodic logging - t1 = time.time() - if i % 1 == 0: - lr = self.lr_scheduler.get_last_lr()[0] - self.training_log( - epoch, self.num_epoch, i, total, loss, lr, t1 - t0 - ) - t0 = time.time() - - if enable_profiling: - profiler.step() - - finally: - if enable_profiling: - profiler.__exit__(None, None, None) - - @torch.no_grad() - def val_loop(self): - """ - Execute validation loop with gradient computation disabled. - - Evaluates model performance on validation set and logs validation loss. - """ - # Initialize validation dataloader - if getattr(self, "val_dataloader", None) is not None: - self.dataset._eval() - self.val_sampler.set_epoch(0) - else: - self.val_dataloader, self.val_sampler = self.dataset.get_val_dataloader() - self.val_sampler.set_epoch(0) - - self.model.eval() - self.val_loss = 0 - - # Validation loop - for i, batch in enumerate( - tqdm( - self.val_dataloader, - desc="Validating", - total=len(self.val_dataloader), - disable=not self.accelerator.is_main_process, - ) - ): - if isinstance(self.dataset, PreprocessedDataset): - batch = { - k: ( - v.to(self.accelerator.device, non_blocking=True) - if isinstance(v, torch.Tensor) - else v - ) - for k, v in batch.items() - } - - with torch.no_grad(): - outputs = self.model(**batch, mode="train") - loss = outputs.loss - self.val_loss += self.accelerator.gather(loss.detach()).mean().item() - - # Calculate average validation loss - self.val_loss /= len(self.val_dataloader) - - # Log validation metrics - if self.logger is not None: - self.logger.log({"val_loss": self.val_loss}, step=self.global_step) - - self.model.train() - - @timer - def load_model(self): - """ - Load and configure the Vision-Language-Action model. - - Handles: - - Model loading from pretrained weights - - Processor initialization - - Optimizer configuration (with support for different learning rates for different components) - - Learning rate scheduler setup - - Model preparation for distributed training - """ - # Load pretrained model - model_type = self.config.get("model_type", "qwen2_5") - assert model_type in ["wall-oss", "qwen2_5"] - if model_type == "wall-oss": - model = Qwen2_5_VLMoEForAction.from_pretrained( - self.config["pretrained_wallx_path"], - train_config=self.config, - action_tokenizer_path=( - self.config["action_tokenizer_path"] - if self.use_fast_tokenizer - else None - ), - ) - self.processor = model.processor - model = model.to(torch.bfloat16) - elif model_type == "qwen2_5": - - model_config = Qwen2_5_VLConfig.from_pretrained( - self.config["qwen_vl_act_config_path"] - ) - flow_loss_weight = self.config.get("flow_loss_weight", 1.0) - self.processor = AutoProcessor.from_pretrained( - self.config["pretrained_wallx_path"], use_fast=True - ) - new_tokens = ["<|propri|>", "<|action|>"] - self.processor.tokenizer.add_tokens(new_tokens) - if self.config.get("use_fast_tokenizer", False): - action_tokenizer_path = self.config["action_tokenizer_path"] - action_tokenizer = AutoProcessor.from_pretrained( - action_tokenizer_path, trust_remote_code=True - ) - # process for use fast - new_tokens = [ - f"<|action_token_{i}|>" for i in range(action_tokenizer.vocab_size) - ] - self.processor.tokenizer.add_tokens(new_tokens) - begin_idx_token = "<|action_token_0|>" - token_id = self.processor.tokenizer.convert_tokens_to_ids( - begin_idx_token - ) - self.processor.tokenizer.init_kwargs["action_token_start_index"] = ( - token_id - ) - self.processor.tokenizer.init_kwargs["action_token_vocab_size"] = ( - action_tokenizer.vocab_size - ) - self.processor.action_processor = action_tokenizer - - # Set the customized robot configuration to ensure consistency between cross-embodiment - # representations and the Wall-X action dimensionality. - # Qwen2_5_VLMoEForAction._set_customized_config(self.config) - customized_dof_config = self.config["customized_robot_config"][ - "customized_dof_config" - ] - customized_agent_pos_config = self.config["customized_robot_config"][ - "customized_agent_pos_config" - ] - setattr(model_config, "customized_dof_config", customized_dof_config) - setattr( - model_config, "customized_agent_pos_config", customized_agent_pos_config - ) - - model_config = update_model_config(self.config, model_config) - model = Qwen2_5_VLMoEForAction( - model_config, - self.use_fast_tokenizer, - self.processor, - flow_loss_weight=flow_loss_weight, - use_selective_recompute=self.use_selective_recompute, - ) - - model = model.to(torch.bfloat16) - model = self.load_qwen_pretrain_weight( - model, self.config["pretrained_wallx_path"] - ) - model.resize_token_embeddings(len(self.processor.tokenizer)) - model = model.to(torch.bfloat16) - else: - raise NotImplementedError(f"Invalid model type: {model_type}") - - # Configure optimizer based on training strategy - if "freeze_vlm" in self.config and self.config["freeze_vlm"]: - print("Freezing VLM parameters, training only MoE experts", flush=True) - moe_params = [] - for name, param in model.named_parameters(): - if "moe.experts.1." not in name: - param.requires_grad = False - else: - moe_params.append(param) - param_groups = [{"params": moe_params, "lr": self.config["learning_rate"]}] - self.optimizer = AdamW(param_groups, weight_decay=0.1) - - elif "action_expert_learning_rate" in self.config: - # Separate learning rates for VLM and action expert parameters - moe_params = [] - vlm_params = [] - for name, param in model.named_parameters(): - if "moe.experts.1." in name: - moe_params.append(param) - else: - vlm_params.append(param) - - # Configure parameter groups - if self.config.get("train_action_expert_only", False): - self.print_rank0("Training action expert only", flush=True) - param_groups = [ - { - "params": moe_params, - "lr": self.config["action_expert_learning_rate"], - } - ] - else: - param_groups = [ - {"params": vlm_params, "lr": self.config["learning_rate"]}, - { - "params": moe_params, - "lr": self.config["action_expert_learning_rate"], - }, - ] - - self.optimizer = AdamW(param_groups, weight_decay=0.1) - self.print_rank0( - f"Setting MoE learning rate to {self.config['action_expert_learning_rate']}, " - f"VLM learning rate to {self.config['learning_rate']}", - flush=True, - ) - else: - # Standard optimizer configuration - self.optimizer = AdamW( - model.parameters(), - lr=self.config["learning_rate"], - weight_decay=0.1, - ) - - # Configure learning rate scheduler - warmup_steps = self.config.get("num_warmup_steps", 0) - num_training_steps = self.config.get("num_training_steps", 1000000000) - min_lr = self.config.get("min_lr", 0.1 * self.config["learning_rate"]) - self.lr_scheduler = get_cosine_with_min_lr_schedule_with_warmup( - optimizer=self.optimizer, - num_warmup_steps=warmup_steps, - num_training_steps=num_training_steps, - min_lr=min_lr, - ) - - self.model = model - - # Enable gradient computation for embeddings - if hasattr(model, "enable_input_require_grads"): - self.model.enable_input_require_grads() - else: - - def make_inputs_require_grad(module, input, output): - output.requires_grad_(True) - - self.model.get_input_embeddings().register_forward_hook( - make_inputs_require_grad - ) - - # Prepare model, optimizer, and scheduler for distributed training - self.model, self.optimizer, self.lr_scheduler = self.accelerator.prepare( - self.model, self.optimizer, self.lr_scheduler - ) - - @timer - def load_qact_data(self): - """ - Load and configure training data for Vision-Language-Action learning. - - Supports LeRobot dataset format and handles distributed data loading - across multiple processes. - """ - print(f"Loading Vision-Language-Action data from {__file__}") - self.accelerator.wait_for_everyone() - - # Load LeRobot dataset - self.dataset, self.train_num = load_lerobot_data( - config=self.config, - lerobot_config=self.dataload_config.get("lerobot_config", {}), - normalizer_action=copy.deepcopy(self.normalizer_action), - normalizer_propri=copy.deepcopy(self.normalizer_propri), - rank=self.rank, - world_size=self.world_size, - ) - - @timer - def load_qwen_pretrain_weight(self, model, pretrain_weight_path): - """ - Load pretrained Qwen weights with MoE adaptation. - - Args: - model: Model instance to load weights into - pretrain_weight_path (str): Path to pretrained weight files - - Returns: - Model with loaded pretrained weights - - Handles weight key renaming for MoE architecture compatibility. - """ - # Load all safetensors files - weight_files = sorted( - [f for f in os.listdir(pretrain_weight_path) if f.endswith(".safetensors")] - ) - merged_weights = {} - - # Merge weights from all files - for weight_file in weight_files: - file_path = os.path.join(pretrain_weight_path, weight_file) - weights = load_file(file_path) - merged_weights.update(weights) - - # Rename weights for MoE compatibility - renamed_weights = {} - for key, value in merged_weights.items(): - if ( - key.startswith("model.layers") - and "mlp." in key - and model.config.mlp_moe - ): - # Rename MLP weights for MoE structure - layer_num = key.split(".layers.")[1].split(".mlp")[0] - new_key = key.replace( - f"layers.{layer_num}.mlp.", f"layers.{layer_num}.moe.experts.0." - ) - renamed_weights[new_key] = value - elif ( - key.startswith("model.layers") - and "self_attn." in key - and model.config.attention_moe - ): - # Rename attention weights for MoE structure - layer_num = key.split(".layers.")[1].split(".self_attn")[0] - proj_types = ["q_proj", "k_proj", "v_proj", "o_proj"] - for proj in proj_types: - if proj in key: - new_key = key.replace( - f"layers.{layer_num}.self_attn.{proj}", - f"layers.{layer_num}.self_attn.{proj}_experts.0", - ) - renamed_weights[new_key] = value - break - else: - renamed_weights[key] = value - - # Load weights into model - # err = model.load_state_dict(renamed_weights, strict=False) - # self.print_rank0(f"Weight loading report: {err}", flush=True) - if self.accelerator.is_main_process: - self.print_rank0(f"Loaded pretrained weights from: {pretrain_weight_path}") - - return model - - def training_log( - self, - current_epoch, - total_epoch, - current_train_iter, - total_train_iter, - loss, - lr, - time_per_step, - ): - """ - Log training progress and performance metrics. - - Args: - current_epoch (int): Current epoch number - total_epoch (int): Total number of epochs - current_train_iter (int): Current training iteration - total_train_iter (int): Total iterations in epoch - loss (torch.Tensor): Current loss value - lr (float): Current learning rate - time_per_step (float): Time taken for current step - """ - timers_to_log = [ - "interval-time", - "data-load", - "forward-compute", - "backward-compute", - "optimizer", - ] - - log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]" - log_string += " epoch {:3d}/{:3d} |".format(current_epoch, total_epoch) - log_string += " iter {:6d}/{:6d} |".format(current_train_iter, total_train_iter) - log_string += " loss {:.6f} |".format(loss) - log_string += " lr {:.6f} |".format(lr) - log_string += " time_per_step_avg {:.6f}s |".format(time_per_step) - - print_rank_last(log_string) - self.timers.log(timers_to_log, normalizer=1) - - def save_checkpoint(self, epoch, step=0): - save_path = self.config["save_path"] - if step == 0: - ckpt_path = f"{save_path}/{epoch}" - else: - ckpt_path = f"{save_path}/{epoch}_{step}" - self.accelerator.save_state(ckpt_path) - - # Save random seed - if self.accelerator.is_main_process: - # FIXME the dataset does not have random seed now. Should the dataset set the random seed? - torch.save( - {"seed": self.seed}, os.path.join(ckpt_path, "seed.pth") - ) # seed is shared by all ranks; seed follows dataset - torch.save( - {"global_step": self.global_step}, - os.path.join(ckpt_path, "global_step.pth"), - ) - torch.save( - {"current_epoch": epoch}, os.path.join(ckpt_path, "current_epoch.pth") - ) - - # Save configuration in YAML format - config_path = os.path.join(ckpt_path, "config.yml") - with open(config_path, "w", encoding="utf-8") as f: - yaml.dump( - self.config, - f, - default_flow_style=False, - allow_unicode=True, - indent=2, - sort_keys=False, - ) - - pretrained_dir = self.config.get("pretrained_qwen_vl_path", None) - if pretrained_dir is not None: - files_to_copy = [ - "preprocessor_config.json", - "tokenizer_config.json", - "tokenizer.json", - "vocab.json", - ] - - for filename in files_to_copy: - src = os.path.join(pretrained_dir, filename) - dst = os.path.join(ckpt_path, filename) - - if os.path.exists(src): - shutil.copy(src, dst) - print(f"[Checkpoint] Copied {filename} to {ckpt_path}") - else: - print(f"[Checkpoint] WARNING: {src} not found, skip copying.") - - act_config_path = self.config.get("qwen_vl_act_config_path", None) - if act_config_path is not None: - dst = os.path.join(ckpt_path, "config.json") - - if os.path.exists(act_config_path): - shutil.copy(act_config_path, dst) - print(f"[Checkpoint] Copied act config to {dst}") - else: - print( - f"[Checkpoint] WARNING: {act_config_path} not found, skipping." - ) - # Save normalizer - torch.save( - self.normalizer_action.state_dict(), - os.path.join(ckpt_path, "normalizer_action.pth"), - ) - torch.save( - self.normalizer_propri.state_dict(), - os.path.join(ckpt_path, "normalizer_propri.pth"), - ) - - # Save current iter steps - if step != 0: # step==0, no need for dataset resume - _rank = self.accelerator.process_index - if self.data_config["multimodal_data_ratio"] != 1: - torch.save( - { - "episode_start_index": self.dataset.primary_pool_start_index.value - }, - os.path.join(ckpt_path, f"episode_start_index_rank_{_rank}.pth"), - ) - torch.save( - { - "multimodal_episode_start_index": self.dataset.secondary_pool_start_index.value - }, - os.path.join( - ckpt_path, f"multimodal_episode_start_index_rank_{_rank}.pth" - ), - ) - else: - torch.save( - { - "episode_start_index": self.dataset.secondary_pool_start_index.value - }, - os.path.join(ckpt_path, f"episode_start_index_rank_{_rank}.pth"), - ) - torch.save( - { - "multimodal_episode_start_index": self.dataset.primary_pool_start_index.value - }, - os.path.join( - ckpt_path, f"multimodal_episode_start_index_rank_{_rank}.pth" - ), - ) - - def resume_from_checkpoint(self): - """ - Resume training from a saved checkpoint. - - Handles both full checkpoint loading and model-only loading based on configuration. - """ - - if self.config.get("resume", {}).get("load_ckpt_only", False): - if self.config.get("FSDP2", False): - self._load_fsdp_state_dict_with_distribute_tensor() - - else: - # Load only model weights - ckpt_path = self.config["resume"]["ckpt"] + "/model.safetensors" - state_dict = load_file(ckpt_path, device="cpu") - - # Add module prefix if needed for distributed training - new_state_dict = {} - for key in state_dict: - if not key.startswith("module."): - new_key = "module." + key - new_state_dict[new_key] = state_dict[key] - - self.model.load_state_dict(new_state_dict, strict=False) - else: - # Load full checkpoint including optimizer and scheduler states - - # self.accelerator.load_state(checkpoint_path) - state_dict = load_file( - self.config["resume"]["ckpt"] + "/model.safetensors", device="cpu" - ) - - filtered_state_dict = { - k: v - for k, v in state_dict.items() - if not k.startswith("action_preprocessor.normalizer") - } - - if self.config["resume"].get("try_harder", False): - new_state_dict = {} - for name, param in filtered_state_dict.items(): - if name in self.model.state_dict(): - if param.size() == self.model.state_dict()[name].size(): - new_state_dict[name] = param - else: - size_0 = param.size() - size_1 = self.model.state_dict()[name].size() - new_state_dict[name] = self.model.state_dict()[name] - slices = [ - slice(0, min(old_dim, new_dim)) - for old_dim, new_dim in zip(size_0, size_1) - ] - new_state_dict[name][slices] = param[slices] - self.print_rank0( - f"Not match key: {name}, required shape: {size_1}, loaded shape: {size_0}, new shape: {new_state_dict[name].size()}" - ) - elif "module." + name in self.model.state_dict(): - name = "module." + name - if param.size() == self.model.state_dict()[name].size(): - new_state_dict[name] = param - else: - size_0 = param.size() - size_1 = self.model.state_dict()[name].size() - new_state_dict[name] = self.model.state_dict()[name] - slices = [ - slice(0, min(old_dim, new_dim)) - for old_dim, new_dim in zip(size_0, size_1) - ] - new_state_dict[name][slices] = param[slices] - self.print_rank0( - f"Not match key: {name}, required shape: {size_1}, loaded shape: {size_0}, new shape: {new_state_dict[name].size()}" - ) - else: - self.print_rank0(f"Not used parameter: {name}") - err = self.model.load_state_dict(new_state_dict, strict=False) - else: - err = self.model.load_state_dict(filtered_state_dict, strict=False) - - self.print_rank0(f"err in load model: {err}", err) - - def _load_fsdp_state_dict_with_distribute_tensor(self): - - rank = dist.get_rank() if dist.is_initialized() else 0 - - full_sd = load_file( - self.config["resume"]["ckpt"] + "/model.safetensors", device="cpu" - ) - meta_sharded_sd = self.model.state_dict() - sharded_sd = {} - - def find_matching_key(target_key, available_keys): - if target_key in available_keys: - return target_key - prefixed_key = f"_orig_mod.{target_key}" - if prefixed_key in available_keys: - return prefixed_key - if target_key.startswith("_orig_mod."): - unprefixed_key = target_key[len("_orig_mod.") :] - if unprefixed_key in available_keys: - return unprefixed_key - return None - - for param_name, full_tensor in full_sd.items(): - matching_key = find_matching_key(param_name, meta_sharded_sd.keys()) - if matching_key is None: - if rank == 0: - print( - f"[Rank {rank}] Warning: Parameter not found:", - param_name, - flush=True, - ) - continue - sharded_meta_param = meta_sharded_sd[matching_key] - sharded_tensor = distribute_tensor( - full_tensor, - sharded_meta_param.device_mesh, - sharded_meta_param.placements, - ) - sharded_sd[matching_key] = nn.Parameter(sharded_tensor) - - self.model.load_state_dict(sharded_sd, assign=True, strict=False) - - def log_l1_details(self, all_label, all_pred, all_task, all_dof_mask): - """ - Log detailed L1 loss metrics by degrees of freedom. - - Args: - all_label (torch.Tensor): Ground truth action labels - all_pred (torch.Tensor): Predicted actions - all_task (list): Task identifiers - all_dof_mask (torch.Tensor): Degrees of freedom mask - - Computes and logs L1 loss for each DOF component separately for detailed analysis. - """ - all_task = all_task[: len(all_label)] - - # Apply DOF mask - all_label = all_label * all_dof_mask - all_pred = all_pred * all_dof_mask - - # Compute baseline L1 loss (predict mean action) - if self.base_l1_loss is None: - mean_action = all_label.mean(dim=0) - self.base_l1_loss = nn.functional.l1_loss(all_label, mean_action) - - self.logger.log( - {"base_l1_loss": self.base_l1_loss.item()}, step=self.global_step - ) - - # Log L1 loss for each DOF component - start_idx = 0 - dof_config = self.config.get("dof_config", {}) - for dof in dof_config: - end_idx = start_idx + dof_config[dof] - dof_label = all_label[:, :, start_idx:end_idx] - dof_pred = all_pred[:, :, start_idx:end_idx] - dof_l1 = nn.functional.l1_loss(dof_pred, dof_label) - - self.print_rank0(f"DOF {dof}, L1 loss: {dof_l1.item()}", flush=True) - self.logger.log( - {f"detail/l1_loss_{dof}": dof_l1.item()}, step=self.global_step - ) - - start_idx = end_idx diff --git a/wall_x/trainer/scheduler/scheduler.py b/wall_x/trainer/scheduler/scheduler.py new file mode 100644 index 00000000..8ad7eafc --- /dev/null +++ b/wall_x/trainer/scheduler/scheduler.py @@ -0,0 +1,70 @@ +from torch.optim.lr_scheduler import LambdaLR, ConstantLR +from torch.optim import Optimizer +import math + + +def create_cosine_scheduler( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + peak_lr: float | None = None, + end_lr: float | None = None, + last_epoch: int = -1, +): + if peak_lr is None: + peak_lr = float(optimizer.defaults["lr"]) + if end_lr is None: + end_lr = peak_lr * 0.1 + + def lr_lambda(current_step: int): + if current_step < num_warmup_steps: + # OpenPI warmup: start from peak_lr / (warmup_steps + 1) + init_lr = peak_lr / (num_warmup_steps + 1) + current_lr = init_lr + (peak_lr - init_lr) * current_step / num_warmup_steps + return current_lr / peak_lr # LambdaLR multiplies by base_lr + else: + # Cosine decay + decay_steps = num_training_steps - num_warmup_steps + progress = min(1.0, (current_step - num_warmup_steps) / max(1, decay_steps)) + cos = 0.5 * (1 + math.cos(math.pi * progress)) + current_lr = end_lr + (peak_lr - end_lr) * cos + return current_lr / peak_lr + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def create_step_scheduler( + optimizer: Optimizer, + lr_decay_steps: str, + lr_gamma: float = 0.1, +): + decay_steps = [int(s.strip()) for s in lr_decay_steps.split(",")] + + def lr_lambda(current_step): + factor = 1.0 + for step in decay_steps: + if current_step >= step: + factor *= lr_gamma + return factor + + return LambdaLR(optimizer, lr_lambda) + + +def create_constant_scheduler( + optimizer: Optimizer, + factor: float = 1 / 3, + total_iters: int = 5, + last_epoch: int = -1, +): + return ConstantLR(optimizer, factor, total_iters, last_epoch) + + +def get_scheduler(optimizer: Optimizer, lr_scheduler_type: str, **kwargs): + if lr_scheduler_type == "cosine": + return create_cosine_scheduler(optimizer, **kwargs) + elif lr_scheduler_type == "step": + return create_step_scheduler(optimizer, **kwargs) + elif lr_scheduler_type == "constant": + return create_constant_scheduler(optimizer, **kwargs) + else: + raise ValueError(f"Unsupported lr_scheduler: {lr_scheduler_type}") diff --git a/wall_x/trainer/trainer_utils.py b/wall_x/trainer/trainer_utils.py new file mode 100644 index 00000000..b6d2c974 --- /dev/null +++ b/wall_x/trainer/trainer_utils.py @@ -0,0 +1,574 @@ +import logging +import math +import os +import random +import threading + +import numpy as np +import psutil +import torch +import torch.nn.functional as F +from safetensors.torch import load_file +from torch.optim.lr_scheduler import LambdaLR +from transformers import AutoProcessor + +from wall_x.model.core.action.normalizer import Normalizer +from wall_x.model.qact.tokenizer_mixin import get_action_tokenizer_mixin +from wall_x.utils.metrics import get_action_accuracy, dtw_distance, frechet_distance + +_logger = logging.getLogger(__name__) + + +def load_wallx_processors( + config, + normalizer=None, + action_statistic_dof=None, + device: str = "cpu", +): + """ + Load Wall-X processors, including tokenizer and action mapper. + + Args: + config: Configuration dictionary. + device: Tokenizer device. Training usually uses "cpu"; inference + usually uses "cuda". + + Returns: + Dictionary containing: + - processor: HuggingFace processor + - train_action_tokenizer: action tokenizer for training + - val_action_tokenizer: action tokenizer for validation + - action_mapper: action mapper dictionary + - num_added_tokens: number of added tokens + - tokenizer_mixin: ActionTokenizerMixin instance + """ + processor = AutoProcessor.from_pretrained(config["processor_path"], use_fast=True) + # pad side = left + processor.tokenizer.padding_side = "left" + + new_tokens = ["<|propri|>", "<|action|>"] + if config.get("new_special_tokens", None) is not None: + new_tokens.extend(config.get("new_special_tokens")) + + action_tokenizer_type = config.get("action_tokenizer_type", None) + + train_action_tokenizer = None + val_action_tokenizer = None + action_mapper = None + tokenizer_mixin = None + + if action_tokenizer_type: + # Use tokenizer_mixin as the single action-tokenizer interface. + action_tokenizer_config = config.get("action_tokenizer", {}) + # Backward compatibility: read top-level keys as fallback values. + action_tokenizer_config.setdefault( + "action_tokenizer_path", config.get("action_tokenizer_path") + ) + action_tokenizer_config.setdefault( + "action_tokenizer_checkpoint_path", + config.get("action_tokenizer_checkpoint_path"), + ) + action_tokenizer_config.setdefault( + "action_tokenizer_config_dir", config.get("action_tokenizer_config_dir") + ) + # Pass action_horizon_ar to the tokenizer for DLLM. + data_config = config.get("data", {}) + action_tokenizer_config.setdefault( + "action_horizon_ar", data_config.get("action_horizon_ar", 32) + ) + # Fall back to dof_config when ar_dof_config is not provided. + ar_dof_config = config.get("ar_dof_config") or config.get("dof_config") + assert ar_dof_config is not None, "Missing ar_dof_config and dof_config" + if normalizer is None: + if action_statistic_dof is None: + raise ValueError( + "Action tokenizer setup requires an explicit normalizer or " + "action statistics. Public Wall-X builds do not bundle " + "default action statistics." + ) + ar_normalizer = Normalizer(action_statistic_dof, ar_dof_config) + else: + ar_normalizer = normalizer + tokenizer_mixin = get_action_tokenizer_mixin(action_tokenizer_type) + tokenizer_mixin.load_tokenizer( + action_tokenizer_config, ar_normalizer, device=device + ) + + # Collect special tokens. + _new_tokens, special_tokens = tokenizer_mixin.get_all_special_tokens() + new_tokens += _new_tokens + + # Add tokens to the vocabulary. + num_added_tokens = processor.tokenizer.add_tokens(new_tokens) + + # Set placeholder_seq for discrete diffusion. + if special_tokens and action_tokenizer_config.get( + "input_placeholder_flag", False + ): + processor.placeholder_seq = [ + processor.tokenizer.convert_tokens_to_ids(token) + for token in special_tokens + ] + + # Backward compatibility: use the first added token when + # <|action_token_0|> does not exist. + ar_first_token_id = processor.tokenizer.convert_tokens_to_ids( + "<|action_token_0|>" + ) + if ( + ar_first_token_id is None + or ar_first_token_id == processor.tokenizer.unk_token_id + ): + ar_first_token_id = processor.tokenizer.convert_tokens_to_ids( + _new_tokens[0] + ) + processor.ar_first_token = ar_first_token_id + + # Build action_mapper. + action_mapper = tokenizer_mixin.build_action_mapper(processor) + + # Fetch the underlying tokenizer. + train_action_tokenizer = tokenizer_mixin.tokenizer + # Fast tokenizers need a separate validation instance; others share one. + val_action_tokenizer = tokenizer_mixin.get_val_tokenizer(config) + else: + num_added_tokens = processor.tokenizer.add_tokens(new_tokens) + + return { + "processor": processor, + "train_action_tokenizer": train_action_tokenizer, + "val_action_tokenizer": val_action_tokenizer, + "action_mapper": action_mapper, + "num_added_tokens": num_added_tokens, + "tokenizer_mixin": tokenizer_mixin, + } + + +def load_wallx_processors_from_cfg( + cfg, + normalizer=None, + action_statistic_dof=None, + device: str = "cpu", +): + """Typed convenience wrapper around ``load_wallx_processors``. + + Builds the flat dict that the legacy function expects from typed + TrainConfig sub-configs, then delegates. Callers using TrainConfig + can use this directly instead of hand-flattening. + """ + import dataclasses + + flat = dataclasses.asdict(cfg.model) + flat["model_type"] = cfg.model_type + flat["data"] = dict(cfg._raw_data or {}) + 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: + flat["ar_dof_config"] = cfg.task.ar_dof_config + flat["batch_size_per_gpu"] = cfg.hyperparams.batch_size_per_gpu + return load_wallx_processors( + flat, + normalizer=normalizer, + action_statistic_dof=action_statistic_dof, + device=device, + ) + + +def load_qwen_pretrain_weight(model, pretrain_weight_path): + weight_files = sorted( + [f for f in os.listdir(pretrain_weight_path) if f.endswith(".safetensors")] + ) + # Initialize empty dictionary to store merged weights + merged_weights = {} + + # Load and merge each file sequentially + for weight_file in weight_files: + file_path = os.path.join(pretrain_weight_path, weight_file) + weights = load_file(file_path) + merged_weights.update(weights) + + renamed_weights = model.rename_vlm_weights_for_vla(merged_weights) + renamed_weights = { + k: v + for k, v in renamed_weights.items() + if "action_preprocessor.normalizer_" not in k + } # remove normalizer weights + if ( + model.config.model_type == "qwen2_5_vl" + and model.model.embed_tokens.weight.shape[0] + != renamed_weights["model.embed_tokens.weight"].shape[0] + ): + _logger.info( + "resize_token_embeddings from %d to %d", + model.model.embed_tokens.weight.shape[0], + renamed_weights["model.embed_tokens.weight"].shape[0], + ) + model.model.resize_token_embeddings( + renamed_weights["model.embed_tokens.weight"].shape[0] + ) + + err = model.load_state_dict(renamed_weights, strict=False) + + return model, err + + +def update_model_config(train_config, model_config): + model_config.use_state_string_representation = train_config["data"].get( + "use_state_string_representation", False + ) + model_config.ar_loss_weight = train_config.get("ar_loss_weight", 1.0) + + model_config.dof_config = train_config["dof_config"] + model_config.agent_pos_config = train_config["agent_pos_config"] + + model_config.action_horizon_flow = train_config["data"].get( + "action_horizon_flow", 32 + ) + + if train_config.get("_attn_implementation", None) is not None: + model_config._attn_implementation = train_config["_attn_implementation"] + + if train_config.get("attn_deterministic", None) is not None: + model_config.attn_deterministic = train_config["attn_deterministic"] + model_config.vision_config.attn_deterministic = train_config[ + "attn_deterministic" + ] + _logger.info("Attention is using deterministic kernel for this run") + else: + model_config.attn_deterministic = True + model_config.vision_config.attn_deterministic = True + + if train_config.get("noise_scheduler", None) is not None: + model_config.noise_scheduler = train_config["noise_scheduler"] + + return model_config + + +def update_data_config(config): + """Keep the top-level model type aligned with the nested data config.""" + config["data"]["model_type"] = config.get("model_type") + + if config.get("use_state_string_representation", None) is not None: + config["data"]["use_state_string_representation"] = config[ + "use_state_string_representation" + ] + + return config + + +def get_detailed_memory_usage(): + """Return process memory and thread-count diagnostics.""" + process = psutil.Process() + memory_info = process.memory_info() + current_threads = threading.active_count() + return { + "rss": f"{memory_info.rss / 1024 / 1024:.2f}MB ", + "vms": f"{memory_info.vms / 1024 / 1024:.2f}MB ", + "threads_count": current_threads, + } + + +def is_last_rank(): + return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1) + + +def print_rank_last(message): + """If distributed is initialized, log only on last rank.""" + if torch.distributed.is_initialized(): + if is_last_rank(): + _logger.info(message) + else: + _logger.info(message) + + +def seed_all(seed): + np.random.seed(seed) + random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + os.environ["PYTHONHASHSEED"] = str(seed) + + +def save_text_results_to_file( + current_step, all_input_texts, all_gt_texts, all_pred_texts, save_dir +): + """Save text-generation results to a local JSON file.""" + import json + + os.makedirs(save_dir, exist_ok=True) + + results = [] + for i, (input_text, gt_text, pred_text) in enumerate( + zip(all_input_texts, all_gt_texts, all_pred_texts) + ): + # Some callers pass one-item lists instead of plain strings. + input_clean = input_text[0] if isinstance(input_text, list) else input_text + gt_clean = gt_text[0] if isinstance(gt_text, list) else gt_text + pred_clean = pred_text[0] if isinstance(pred_text, list) else pred_text + + results.append( + { + "sample_id": i, + "input": input_clean, + "ground_truth": gt_clean, + "prediction": pred_clean, + "step": current_step, + } + ) + + filename = os.path.join(save_dir, f"text_predictions_step_{current_step}.json") + with open(filename, "w", encoding="utf-8") as f: + json.dump(results, f, ensure_ascii=False, indent=2) + + +# Add accuracy metrics. +# Add ade, fde, dtw, and frechet metrics from Ryan. +def compute_action_metrics(all_preds, all_actions, config, step_log={}): + metrics_settings = config.get("metrics_settings", "default") + metrics_available = ["l1", "mse", "accuracy", "ade", "fde", "dtw", "frechet"] + metrics_default = ["l1", "mse", "accuracy"] + if metrics_settings == "all": + metrics_enabled = metrics_available + elif metrics_settings == "default": + metrics_enabled = metrics_default + elif isinstance(metrics_settings, str): + metrics_enabled = [ + m.strip() + for m in metrics_settings.split(" ") + if m.strip() in metrics_available + ] + elif isinstance(metrics_settings, list): + metrics_enabled = [m for m in metrics_settings if m in metrics_available] + else: + metrics_enabled = metrics_default + _logger.warning( + 'Unknown metrics_settings "%s". Using default metrics: %s', + metrics_settings, + metrics_enabled, + ) + + overall_l1 = F.l1_loss(all_preds, all_actions) + overall_mse = F.mse_loss(all_preds, all_actions) + + step_log["val_action_l1"] = overall_l1.item() + step_log["val_action_mse"] = overall_mse.item() + + if "accuracy" in metrics_enabled: + accuracy_thresholds = [0.05, 0.1, 0.2, 0.4] + # The accuracy that all predicted action dimensions are within a certain range of the ground truth. + accuracies = get_action_accuracy( + all_preds, all_actions, thresholds=accuracy_thresholds + ) + for th_idx, threshold in enumerate(accuracy_thresholds): + step_log[f"val_action_acc_thr{threshold}"] = accuracies[th_idx].item() + + start_idx = 0 + dof_config = config["dof_config"] + for dof_key, dof_dim in dof_config.items(): + end_idx = start_idx + dof_dim + # all_preds.shape = (B, T, action_dim) + pred_dof = all_preds[..., start_idx:end_idx] + action_dof = all_actions[..., start_idx:end_idx] + dof_l1 = F.l1_loss(pred_dof, action_dof) + dof_mse = F.mse_loss(pred_dof, action_dof) + step_log[f"val_l1/{dof_key}"] = dof_l1.item() + step_log[f"val_mse/{dof_key}"] = dof_mse.item() + + if "accuracy" in metrics_enabled: + accuracies = get_action_accuracy( + action_dof, pred_dof, thresholds=accuracy_thresholds + ) + for th_idx, threshold in enumerate(accuracy_thresholds): + step_log[f"val_acc/{dof_key}_thr{threshold}"] = accuracies[ + th_idx + ].item() + + if "ee_cartesian_pos" in dof_key: + if "ade" in metrics_enabled: + displacement_error = torch.norm(pred_dof - action_dof, dim=-1) + ade = torch.mean(displacement_error) + step_log[f"val_ade/{dof_key}"] = ade.item() + + if "fde" in metrics_enabled: + final_pred = pred_dof[:, -1, :] + final_gt = action_dof[:, -1, :] + fde = torch.mean(torch.norm(final_pred - final_gt, dim=-1)) + step_log[f"val_fde/{dof_key}"] = fde.item() + + if "dtw" in metrics_enabled or "frechet" in metrics_enabled: + dtw_distances = [] + frechet_distances = [] + batch_size = pred_dof.shape[0] + for i in range(batch_size): + pred_seq = pred_dof[i] # shape: (T, D) + gt_seq = action_dof[i] # shape: (T, D) + dtw_dist = dtw_distance(pred_seq, gt_seq) + dtw_distances.append(dtw_dist) + + frechet_dist = frechet_distance(pred_seq, gt_seq) + frechet_distances.append(frechet_dist) + + avg_dtw = torch.mean(torch.stack(dtw_distances)) + if "dtw" in metrics_enabled: + step_log[f"val_dtw/{dof_key}"] = avg_dtw.item() + + avg_frechet = torch.mean(torch.stack(frechet_distances)) + if "frechet" in metrics_enabled: + step_log[f"val_frechet/{dof_key}"] = avg_frechet.item() + + start_idx = end_idx + + return step_log + + +def get_openpi_cosine_schedule_with_warmup( + optimizer, + num_warmup_steps: int, + num_training_steps: int, + peak_lr: float = None, + end_lr: float = None, + last_epoch: int = -1, +): + """ + Create a schedule with a learning rate that follows OpenPI style: + - Warmup: linearly increases from peak_lr/(warmup_steps+1) to peak_lr + - Decay: cosine decay from peak_lr to end_lr + + Args: + optimizer: The optimizer for which to schedule the learning rate. + num_warmup_steps: The number of steps for the warmup phase. + num_training_steps: The total number of training steps. + peak_lr: The peak learning rate. If None, uses optimizer's initial lr. + end_lr: The minimum learning rate at the end. If None, defaults to peak_lr * 0.1. + last_epoch: The index of the last epoch when resuming training. + + Return: + torch.optim.lr_scheduler.LambdaLR with the appropriate schedule. + """ + if peak_lr is None: + peak_lr = optimizer.defaults["lr"] + if end_lr is None: + end_lr = peak_lr * 0.1 + + def lr_lambda(current_step: int): + if current_step < num_warmup_steps: + # OpenPI warmup: start from peak_lr / (warmup_steps + 1) + init_lr = peak_lr / (num_warmup_steps + 1) + current_lr = init_lr + (peak_lr - init_lr) * current_step / num_warmup_steps + return current_lr / peak_lr # LambdaLR multiplies by base_lr + else: + # Cosine decay + decay_steps = num_training_steps - num_warmup_steps + progress = min(1.0, (current_step - num_warmup_steps) / max(1, decay_steps)) + cos = 0.5 * (1 + math.cos(math.pi * progress)) + current_lr = end_lr + (peak_lr - end_lr) * cos + return current_lr / peak_lr + + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def plot_openloop( + action_pred_list, + action_gt_list, + l1_loss, + episode_index, + save_path, + is_static_list=None, +): + """ + Plot openloop action comparison visualization. + + Args: + action_pred_list: List of predicted actions, each with shape (horizon, action_dim) + action_gt_list: List of ground truth actions, each with shape (horizon, action_dim) + l1_loss: L1 loss array with shape (total_frames, action_dim) + episode_index: Index of the episode being visualized + save_path: Directory path to save the plot + is_static_list: Optional list of booleans indicating static frames + """ + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + assert len(action_pred_list) == len( + action_gt_list + ), "Predicted action and ground truth action must have the same shape." + + dim = action_pred_list[0].shape[1] + plt.figure(figsize=(12, 4 * dim)) + + for i in range(dim): + plt.subplot(dim, 1, i + 1) + + has_labeled_static = False + for j in range(len(action_gt_list)): + gt_action = action_gt_list[j] + predict_action = action_pred_list[j] + + x_vals_gt = np.linspace(j, j + 1, len(gt_action)) + x_vals_pred = np.linspace(j, j + 1, len(predict_action)) + + if is_static_list is not None and is_static_list[j]: + label = None + if not has_labeled_static: + label = "Static GT" + has_labeled_static = True + plt.axvspan(j, j + 1, color="gray", alpha=0.2, label=label) + + if j == 0: + plt.plot( + x_vals_gt, + gt_action[:, i], + label="Ground Truth", + color="blue", + linewidth=2, + linestyle="-", + marker="o", + markersize=3, + ) + plt.plot( + x_vals_pred, + predict_action[:, i], + label="Model Output", + color="orange", + linewidth=2, + linestyle="--", + marker="x", + markersize=4, + ) + else: + plt.plot( + x_vals_gt, + gt_action[:, i], + color="blue", + linewidth=2, + linestyle="-", + marker="o", + markersize=3, + ) + plt.plot( + x_vals_pred, + predict_action[:, i], + color="orange", + linewidth=2, + linestyle="--", + marker="x", + markersize=4, + ) + + plt.title(f"Action Dimension {i + 1}, L1 Loss: {l1_loss[:, i].mean():.6f}") + plt.xlabel("Number of Chunk") + plt.ylabel("Action Value") + plt.legend() + + plt.suptitle( + f"Openloop Action Comparison for Episode {episode_index}, L1 Loss: {l1_loss.mean():.6f}" + ) + plt.tight_layout(rect=[0, 0, 1, 0.98]) + os.makedirs(save_path, exist_ok=True) + plt.savefig(f"{save_path}/{episode_index}.png") + plt.close() + _logger.info("Saved openloop plot to %s/%s.png", save_path, episode_index) diff --git a/wall_x/trainer/utils/__init__.py b/wall_x/trainer/utils/__init__.py new file mode 100644 index 00000000..d74f4750 --- /dev/null +++ b/wall_x/trainer/utils/__init__.py @@ -0,0 +1,6 @@ +"""Trainer utility helpers.""" + +from .data import move_batch_to_device +from .diagnostics import log_gpu_memory + +__all__ = ["move_batch_to_device", "log_gpu_memory"] diff --git a/wall_x/trainer/utils/data.py b/wall_x/trainer/utils/data.py new file mode 100644 index 00000000..8ca974cd --- /dev/null +++ b/wall_x/trainer/utils/data.py @@ -0,0 +1,46 @@ +"""Data-related utility functions used by the trainer main loop.""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def move_batch_to_device( + batch: Any, + device: torch.device, + *, + non_blocking: bool = True, +) -> Any: + """Move every tensor in ``batch`` to ``device``, recursing into dict/list. + + Returns a new structure with the same shape; the input ``batch`` is not + mutated. dict / list containers are rebuilt; tensors are moved via + ``.to(device, non_blocking=...)``; everything else is passed through + by reference. + + Parameters + ---------- + batch : Any + Typically a dict produced by the dataloader, but recursion accepts + dict / list / tensor / arbitrary leaves. + device : torch.device + Target device, typically ``self.device`` on the trainer. + non_blocking : bool + Whether to use pinned-memory async copy. Default True because the + trainer uses pinned loaders; pass False if the dataloader hasn't + pinned memory. + """ + if isinstance(batch, dict): + return { + k: move_batch_to_device(v, device, non_blocking=non_blocking) + for k, v in batch.items() + } + if isinstance(batch, list): + return [ + move_batch_to_device(v, device, non_blocking=non_blocking) for v in batch + ] + if isinstance(batch, torch.Tensor): + return batch.to(device, non_blocking=non_blocking) + return batch diff --git a/wall_x/trainer/utils/diagnostics.py b/wall_x/trainer/utils/diagnostics.py new file mode 100644 index 00000000..b8356fb9 --- /dev/null +++ b/wall_x/trainer/utils/diagnostics.py @@ -0,0 +1,40 @@ +"""Training-run diagnostic helpers (CUDA memory, etc.).""" + +from __future__ import annotations + +from typing import Callable, Optional + +import torch + + +def log_gpu_memory( + device: torch.device, + rank: int, + *, + stage: str = "", + log_fn: Optional[Callable] = None, +) -> None: + """Log per-rank GPU memory (allocated / reserved / total) via ``log_fn``. + + Calls ``torch.cuda.synchronize`` on ``device`` so the numbers reflect + the actual post-op usage, not pending work. If ``log_fn`` is None this + is a no-op. + """ + if log_fn is None: + return + torch.cuda.synchronize() + allocated = torch.cuda.memory_allocated(device) / 1024**3 + reserved = torch.cuda.memory_reserved(device) / 1024**3 + peak_allocated = torch.cuda.max_memory_allocated(device) / 1024**3 + peak_reserved = torch.cuda.max_memory_reserved(device) / 1024**3 + total = torch.cuda.get_device_properties(device).total_memory / 1024**3 + tag = f"[{stage}] " if stage else "" + log_fn( + f"{tag}GPU memory rank{rank} " + f"| allocated {allocated:.2f} GiB" + f" | reserved {reserved:.2f} GiB" + f" | peak_allocated {peak_allocated:.2f} GiB" + f" | peak_reserved {peak_reserved:.2f} GiB" + f" | total {total:.2f} GiB", + main_process_only=False, + ) diff --git a/wall_x/utils/constant.py b/wall_x/utils/constant.py index 3074591e..046f4b29 100644 --- a/wall_x/utils/constant.py +++ b/wall_x/utils/constant.py @@ -1,362 +1,59 @@ -action_statistic_dof = { - "x2_normal": { - # water flowers - "follow_left_arm_joint_cur": { - "min": [-3.7121], - "delta": [7.6008], - }, - "follow_right_arm_joint_cur": { - "min": [-3.6176], - "delta": [8.5015], - }, - "follow_left_ee_cartesian_pos": { - "min": [-0.036, -0.3241, -0.1245], - "delta": [0.4389, 0.557, 0.479], - }, - "follow_left_ee_rotation": { - "min": [-1.2373, -0.1929, -1.5182], - "delta": [2.2009, 1.5669, 2.0936], - }, - "follow_left_gripper": {"min": [-0.1196], "delta": [4.5226]}, - "follow_right_ee_cartesian_pos": { - "min": [-0.0326, -0.2273, -0.1377], - "delta": [0.4574, 0.5704, 0.4743], - }, - "follow_right_ee_rotation": { - "min": [-1.2201, -0.2611, -0.7427], - "delta": [2.6623, 1.6622, 2.4186], - }, - "follow_right_gripper": {"min": [-0.1208], "delta": [4.5261]}, - "height": {"min": [-0.0001], "delta": [0.5051]}, - "head_actions": {"min": [-1.5000, -1.4167], "delta": [2.5000, 1.8879]}, - "base_velocity": { - "min": [-0.0359, -0.084, -0.0162], - "delta": [0.1539, 0.1848, 0.0322], - }, - }, - "DobbE": { - "follow_right_ee_cartesian_pos": { - "min": [-0.6107, -0.3272, -0.4282], - "delta": [1.2629, 1.5297, 0.8349], - }, - "follow_right_ee_rotation": { - "min": [-1.7378, -1.4597, -1.8712], - "delta": [2.7031, 2.8182, 3.5921], - }, - "follow_right_gripper": {"min": [0.0], "delta": [0.9983]}, - }, - "RH20T": { - "follow_right_ee_cartesian_pos": { - "min": [0.3646, -0.2722, 0.0066], - "delta": [0.3813, 0.5973, 0.3277], - }, - "follow_right_ee_rotation": { - "min": [-1.8716, -0.4398, -3.1414], - "delta": [3.4145, 1.0225, 6.2828], - }, - "follow_right_gripper": {"min": [0.0], "delta": [95.0]}, - }, - "agibotworld_alpha": { - "follow_left_ee_cartesian_pos": { - "min": [0.4954, 0.0166, 0.1729], - "delta": [0.3336, 0.5123, 0.9189], - }, - "follow_left_ee_rotation": { - "min": [-3.1064, -1.2629, -3.1238], - "delta": [6.2127, 2.5923, 6.2496], - }, - "follow_left_gripper": {"min": [34.6222], "delta": [86.1921]}, - "follow_right_ee_cartesian_pos": { - "min": [0.4615, -0.5975, 0.1638], - "delta": [0.3823, 0.5577, 0.8873], - }, - "follow_right_ee_rotation": { - "min": [-3.0891, -1.0739, -2.5091], - "delta": [6.1707, 2.3074, 3.8533], - }, - "follow_right_gripper": {"min": [34.6222], "delta": [85.7635]}, - "height": {"min": [0.0], "delta": [0.4535]}, - "head_actions": {"min": [-0.1746, 0.0523], "delta": [0.2444, 0.4713]}, - }, - "austin_buds": { - "follow_right_ee_cartesian_pos": { - "min": [0.3496, -0.2855, 0.0105], - "delta": [0.3748, 0.492, 0.3116], - }, - "follow_right_ee_rotation": { - "min": [-3.1405, -0.151, -0.0737], - "delta": [6.2813, 0.3218, 0.1536], - }, - "follow_right_gripper": {"min": [0.0076], "delta": [0.0724]}, - }, - "austin_sailor": { - "follow_right_ee_cartesian_pos": { - "min": [0.387, -0.3165, 0.0244], - "delta": [0.2999, 0.5252, 0.2308], - }, - "follow_right_ee_rotation": { - "min": [-3.1402, -0.1618, -1.5918], - "delta": [6.2804, 0.337, 2.9478], - }, - "follow_right_gripper": {"min": [0.0005], "delta": [0.0773]}, - }, - "austin_sirius": { - "follow_right_ee_cartesian_pos": { - "min": [0.0, -0.1182, 0.0], - "delta": [0.5329, 0.3812, 0.2723], - }, - "follow_right_ee_rotation": { - "min": [-3.1407, -0.1243, -1.7434], - "delta": [6.2823, 0.1975, 1.8073], - }, - "follow_right_gripper": {"min": [0.0334], "delta": [0.046]}, - }, - "bc_z": { - "follow_right_ee_cartesian_pos": { - "min": [-0.3883, -0.1116, 0.6113], - "delta": [0.7199, 0.4288, 0.3709], - }, - "follow_right_ee_rotation": { - "min": [-1.056, -1.0587, -2.6295], - "delta": [1.9142, 1.9455, 4.8064], - }, - "follow_right_gripper": {"min": [0.2], "delta": [0.8]}, - }, - "berkeley_autolab_ur5": { - "follow_right_ee_cartesian_pos": { - "min": [0.3018, -0.2129, -0.1888], - "delta": [0.3121, 0.52, 0.3107], - }, - "follow_right_ee_rotation": { - "min": [-3.1396, -0.2278, 1.1413], - "delta": [6.279, 0.454, 0.9841], - }, - "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, - }, - "berkeley_cable_routing": { - "follow_right_ee_cartesian_pos": { - "min": [0.4617, -0.28, 0.03], - "delta": [0.1838, 0.5665, 0.1272], - }, - "follow_right_ee_rotation": { - "min": [-3.1413, -0.0299, -0.7665], - "delta": [6.2826, 0.0692, 3.322], - }, - }, - "berkeley_fanuc_manipulation": { - "follow_right_ee_cartesian_pos": { - "min": [0.3718, -0.4072, 0.0184], - "delta": [0.3483, 0.7201, 0.5229], - }, - "follow_right_ee_rotation": { - "min": [-3.1399, -1.0166, -1.6988], - "delta": [6.2802, 1.4498, 3.2074], - }, - "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, - }, - "bridge_data_v2": { - "follow_right_ee_cartesian_pos": { - "min": [0.1498, -0.2178, -0.0901], - "delta": [0.3012, 0.469, 0.298], - }, - "follow_right_ee_rotation": { - "min": [-0.3279, -0.6105, -1.0578], - "delta": [0.7378, 1.0353, 2.2552], - }, - "follow_right_gripper": {"min": [0.0692], "delta": [0.9426]}, - }, - "dlr_edan_shared_control": { - "follow_right_ee_cartesian_pos": { - "min": [-0.8387, 0.1473, -0.3934], - "delta": [0.6579, 0.6025, 1.1566], - }, - "follow_right_ee_rotation": { - "min": [-3.1217, -1.5197, -2.2516], - "delta": [6.2505, 1.5594, 4.2831], - }, - "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, - }, - "droid": { - "follow_right_ee_cartesian_pos": { - "min": [0.2667, -0.4396, -0.0472], - "delta": [0.5159, 0.8806, 0.8331], - }, - "follow_right_ee_rotation": { - "min": [-3.1374, -1.216, -2.1741], - "delta": [6.2749, 2.1075, 4.2259], - }, - "follow_right_gripper": {"min": [0.0], "delta": [0.9912]}, - }, - "fmb": { - "follow_right_ee_cartesian_pos": { - "min": [0.3554, -0.2844, 0.0354], - "delta": [0.336, 0.4961, 0.2943], - }, - "follow_right_ee_rotation": { - "min": [-3.1404, -0.9302, -0.0599], - "delta": [6.2807, 1.724, 1.8284], - }, - "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, - }, - "fractal": { - "follow_right_ee_cartesian_pos": { - "min": [0.3242, -0.2836, 0.1405], - "delta": [0.5518, 0.4963, 0.9328], - }, - "follow_right_ee_rotation": { - "min": [-3.1308, -0.2421, -2.9685], - "delta": [6.2609, 1.7343, 5.819], - }, - "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, - }, - "furniture_bench": { - "follow_right_ee_cartesian_pos": { - "min": [0.3691, -0.181, 0.0058], - "delta": [0.2962, 0.3582, 0.1775], - }, - "follow_right_ee_rotation": { - "min": [-3.1394, -0.6121, -1.9958], - "delta": [6.2786, 1.6114, 3.7748], - }, - "follow_right_gripper": {"min": [0.0035], "delta": [0.0762]}, - }, - "jaco_play": { - "follow_right_ee_cartesian_pos": { - "min": [-0.3787, -0.6294, 0.1682], - "delta": [0.5898, 0.3587, 0.2183], - }, - "follow_right_ee_rotation": { - "min": [0.9792, -0.0668, -0.0498], - "delta": [0.0175, 0.1277, 0.0686], - }, - "follow_right_gripper": {"min": [0.0791], "delta": [0.1033]}, - }, - "nyu_rot": { - "follow_right_ee_cartesian_pos": { - "min": [0.25, -1.0, -0.2], - "delta": [0.75, 2.0, 1.2], - }, - "follow_right_ee_rotation": { - "min": [-3.1416, -3.1416, 6.2831], - "delta": [9.4248, 4.1416, 0.0], - }, - "follow_right_gripper": {"min": [0.0], "delta": [1.0]}, - }, - "stanford_hydra": { - "follow_right_ee_cartesian_pos": { - "min": [0.2068, -0.274, 0.1317], - "delta": [0.4929, 0.4981, 0.4588], - }, - "follow_right_ee_rotation": { - "min": [-3.1321, -0.7496, -3.0269], - "delta": [6.2658, 1.5176, 5.8261], - }, - "follow_right_gripper": {"min": [0.0], "delta": [0.0811]}, - }, - "stanford_kuka_multimodal": { - "follow_right_ee_cartesian_pos": { - "min": [0.4781, -0.0659, 0.3424], - "delta": [0.0868, 0.0864, 0.1863], - }, - "follow_right_ee_rotation": { - "min": [-3.136, -0.0521, -3.1413], - "delta": [6.2727, 0.1109, 6.2825], - }, - "follow_right_gripper": {"min": [-0.4713], "delta": [0.9485]}, - }, - "taco_play": { - "follow_right_ee_cartesian_pos": { - "min": [0.1375, -0.4291, 0.2052], - "delta": [0.5327, 1.0237, 0.3913], - }, - "follow_right_ee_rotation": { - "min": [-3.1391, -0.6946, -1.2808], - "delta": [6.2784, 0.8196, 3.0856], - }, - "follow_right_gripper": {"min": [0.0001], "delta": [0.0806]}, - }, - "utaustin_mutex": { - "follow_right_ee_cartesian_pos": { - "min": [0.3213, -0.4734, 0.0141], - "delta": [0.2108, 0.8471, 0.5644], - }, - "follow_right_ee_rotation": { - "min": [-3.1404, -0.2202, -1.5489], - "delta": [6.2805, 0.582, 1.9282], - }, - "follow_right_gripper": {"min": [0.0019], "delta": [0.0738]}, - }, - "viola": { - "follow_right_ee_cartesian_pos": { - "min": [0.4011, -0.2521, 0.0103], - "delta": [0.2444, 0.4305, 0.4355], - }, - "follow_right_ee_rotation": { - "min": [-3.1403, -0.2737, -1.8626], - "delta": [6.2804, 0.4901, 2.0618], - }, - "follow_right_gripper": {"min": [0.0002], "delta": [0.0773]}, - }, - "kuka": { - "follow_right_ee_cartesian_pos": { - "min": [0.3914, -0.4901, 0.0175], - "delta": [0.3339, 0.8357, 0.9064], - }, - "follow_right_ee_rotation": { - "min": [-3.1416, -0.9903, -3.1416], - "delta": [6.2832, 2.2421, 6.2832], - }, - "follow_right_gripper": { - "min": [0.0000], - "delta": [1.0000], - }, - }, - "UMI-biarm": { - "follow_left_ee_cartesian_pos": { - "min": [-0.2917, -0.4926, 0.0063], - "delta": [0.9028, 0.8168, 0.3473], - }, - "follow_left_ee_rotation": { - "min": [-2.5309, -1.5706, -1.3309], - "delta": [0.8758, 2.3315, 1.7076], - }, - "follow_left_gripper": { - "min": [0.0029], - "delta": [0.0812], - }, - "follow_right_ee_cartesian_pos": { - "min": [-0.0023, -0.5191, -0.0358], - "delta": [0.7474, 0.8668, 0.351], - }, - "follow_right_ee_rotation": { - "min": [-2.4945, -2.0149, -0.8088], - "delta": [1.0941, 3.2628, 2.0018], - }, - "follow_right_gripper": { - "min": [0.0019], - "delta": [0.0814], - }, - }, - "agibotworld_beta": { - "follow_left_ee_cartesian_pos": { - "min": [0.4954, 0.0166, 0.1729], - "delta": [0.3336, 0.5123, 0.9189], - }, - "follow_left_ee_rotation": { - "min": [-3.1064, -1.2629, -3.1238], - "delta": [6.2127, 2.5923, 6.2496], - }, - "follow_left_gripper": {"min": [34.6222], "delta": [86.1921]}, - "follow_right_ee_cartesian_pos": { - "min": [0.4615, -0.5975, 0.1638], - "delta": [0.3823, 0.5577, 0.8873], - }, - "follow_right_ee_rotation": { - "min": [-3.0891, -1.0739, -2.5091], - "delta": [6.1707, 2.3074, 3.8533], - }, - "follow_right_gripper": {"min": [34.6222], "delta": [85.7635]}, - "height": {"min": [0.0], "delta": [0.4535]}, - "head_actions": {"min": [-0.1746, 0.0523], "delta": [0.2444, 0.4713]}, - }, -} +"""OSS-safe compatibility constants. + +The internal repository keeps canonical robot, dataset, action-key, and action +statistics tables in ``x2robot_utils.constants``. Those tables are intentionally +not bundled in the public export. This shim keeps the public VLA code importable +without exposing private dataset catalogues or default normalization stats. +""" + +from __future__ import annotations + +from collections.abc import Iterable + + +ACTION_KEY_FULL_MAPPING: dict[str, str] = {} +_ACTION_KEY_FULL_MAPPING = ACTION_KEY_FULL_MAPPING + +ACTION_DATASET_NAMES: tuple[str, ...] = () +MULTIMODAL_DATASET_NAMES: tuple[str, ...] = () +_ACTION_DATASET_NAMES = ACTION_DATASET_NAMES +_MULTIMODAL_DATASET_NAMES = MULTIMODAL_DATASET_NAMES + +VIEW_SLOT_KEYS = ["view1", "view2", "view3"] +PHYSICAL_VIEW_IDS: dict[str, int] = {} + + +def get_view_slot_keys(num_views: int) -> list[str]: + """Return canonical view slot names for ``num_views`` cameras.""" + return [f"view{i + 1}" for i in range(max(0, int(num_views)))] + + +def is_multimodal_dataset_name(dataset_name: str | None) -> bool: + """Return whether a dataset is marked multimodal in the public shim.""" + return dataset_name in MULTIMODAL_DATASET_NAMES + + +def is_action_dataset_name(dataset_name: str | None) -> bool: + """Return whether a dataset row should be treated as an action sample.""" + return dataset_name is not None and not is_multimodal_dataset_name(dataset_name) + + +def iter_action_dataset_names(dataset_names: Iterable[str]) -> list[str]: + """Filter a sequence down to action dataset names.""" + return [name for name in dataset_names if is_action_dataset_name(name)] + + +__all__ = [ + "ACTION_KEY_FULL_MAPPING", + "_ACTION_KEY_FULL_MAPPING", + "ACTION_DATASET_NAMES", + "MULTIMODAL_DATASET_NAMES", + "_ACTION_DATASET_NAMES", + "_MULTIMODAL_DATASET_NAMES", + "VIEW_SLOT_KEYS", + "PHYSICAL_VIEW_IDS", + "get_view_slot_keys", + "is_multimodal_dataset_name", + "is_action_dataset_name", + "iter_action_dataset_names", +] diff --git a/wall_x/utils/cudagraph_wrapper.py b/wall_x/utils/cudagraph_wrapper.py new file mode 100644 index 00000000..b0b21969 --- /dev/null +++ b/wall_x/utils/cudagraph_wrapper.py @@ -0,0 +1,277 @@ +import os +import torch + +ENABLE_CUDA_GRAPH = os.environ.get("ENABLE_CUDA_GRAPH", "True").lower() == "true" +# With a single master buffer + multi-bucket graph, max batch must be pre-allocated; +# resizing after capture invalidates addresses recorded by the graph. +CUDA_GRAPH_MAX_BS = int(os.environ.get("CUDA_GRAPH_MAX_BS", "128")) +_SHAPE_GUARD_FASTPATH = ( + os.environ.get("CUDAGRAPH_DISABLE_SHAPE_GUARD_FASTPATH", "0") != "1" +) + + +class CUDAGraph_Wrapper: + """ + CUDA Graph wrapper with bucketed batch sizes: + - Select/reuse a CUDAGraph per bucket_size when batch size changes + - For batch_size < bucket_size: copy inputs into static buffer [:B], return output [:B] after replay + """ + + def __init__( + self, + model, + warm_up_times: int = 3, + enable: bool | None = None, + batch_size_key: str | None = None, + ): + self.model = model + self.warm_up_times = warm_up_times + self._enable = ENABLE_CUDA_GRAPH if enable is None else bool(enable) + self._batch_size_key = batch_size_key or "suffix_inputs_embeds" + + # graph cache: bucket_bs -> CUDAGraph + self.graphs: dict[int, torch.cuda.CUDAGraph] = {} + # Per bucket: views into the master buffer (graph capture depends on address/shape/stride) + self.static_inputs_map: dict[int, dict[str, torch.Tensor]] = {} + self.static_output_tensor: dict[int, torch.Tensor] = {} + self.graph_pool = None + + # Master buffers (shared across all buckets) + self.graph_vars: dict[str, torch.Tensor] = {} + self._max_bs: int | None = None + + # Non-batch dimension shape signature (global); catches silent mismatches + self._shape_signature: dict[str, tuple] = {} + # Batch dimension index per input (default 0; e.g. suffix_position_ids batch at dim=1) + self._batch_dim: dict[str, int] = {} + self._shape_guard_verified: dict[int, bool] = {} + + def _get_batch_size(self, **kwargs) -> int: + key = self._batch_size_key + if key in kwargs: + v = kwargs[key] + if isinstance(v, torch.Tensor) and v.dim() >= 1: + return int(v.shape[0]) + raise ValueError(f"CUDAGraph_Wrapper.forward requires {key} argument") + + def _select_bucket_bs(self, bs: int) -> int: + """ + Bucket selection: + - 1,2,4,8 use power-of-two + - >8 align to 16 (16,32,48,...) + """ + if bs <= 1: + return 1 + if bs <= 2: + return 2 + if bs <= 4: + return 4 + if bs <= 8: + return 8 + return int(((bs + 15) // 16) * 16) + + def _batch_dim_for_key(self, key: str, tensor: torch.Tensor) -> int: + # Minimal special-case handling for known inputs + # suffix_position_ids: shape [3, B, T], batch at dim=1 + if key == "suffix_position_ids" and tensor.dim() >= 2: + return 1 + return 0 + + def _view_for_bucket(self, tensor: torch.Tensor, bucket_bs: int, batch_dim: int): + if tensor.dim() == 0: + return tensor + if batch_dim == 0: + return tensor[:bucket_bs] + if batch_dim == 1: + return tensor[:, :bucket_bs] + raise ValueError(f"Unsupported batch_dim={batch_dim} for cudagraph wrapper") + + def _copy_and_pad( + self, + master: torch.Tensor, + tensor: torch.Tensor, + bs: int, + bucket_bs: int, + batch_dim: int, + ): + if tensor.dim() == 0: + master.copy_(tensor) + return + if batch_dim == 0: + master[:bs].copy_(tensor) + if bs < bucket_bs: + master[bs:bucket_bs].zero_() + return + if batch_dim == 1: + master[:, :bs].copy_(tensor) + if bs < bucket_bs: + master[:, bs:bucket_bs].zero_() + return + raise ValueError(f"Unsupported batch_dim={batch_dim} for cudagraph wrapper") + + def _allocate_master_tensor( + self, key: str, tensor: torch.Tensor, max_bs: int, bs: int + ) -> torch.Tensor: + """ + Allocate a master buffer for one input: + - Extend only the batch dim to max_bs; other dims unchanged + - Copy the valid [:bs] region; zero the rest (avoid stale values on replay) + """ + if tensor.dim() == 0: + return tensor.clone() + batch_dim = self._batch_dim_for_key(key, tensor) + shape = list(tensor.shape) + if batch_dim >= len(shape): + raise ValueError( + f"Invalid batch_dim={batch_dim} for key={key}, tensor.shape={tuple(tensor.shape)}" + ) + shape[batch_dim] = max_bs + buf = torch.zeros(tuple(shape), device=tensor.device, dtype=tensor.dtype) + # Initial copy: pad to current bucket (at least cover valid [:bs]) + self._copy_and_pad(buf, tensor, bs=bs, bucket_bs=bs, batch_dim=batch_dim) + return buf + + def _ensure_master_buffers(self, bs: int, **kwargs): + """ + Ensure master buffers are allocated once. + - max_bs from CUDA_GRAPH_MAX_BS; if unset, first bucket_bs becomes max_bs + and later larger buckets are rejected (resize would invalidate captured graphs). + """ + if self._max_bs is not None: + return + + bucket_bs = self._select_bucket_bs(bs) + max_bs = CUDA_GRAPH_MAX_BS if CUDA_GRAPH_MAX_BS > 0 else bucket_bs + self._max_bs = int(max_bs) + + for k, v in kwargs.items(): + if not isinstance(v, torch.Tensor): + raise TypeError( + f"CUDAGraph_Wrapper only supports Tensor kwargs; got {k}={type(v)}" + ) + batch_dim = self._batch_dim_for_key(k, v) if v.dim() > 0 else 0 + self._batch_dim[k] = batch_dim + # Record non-batch dims (must stay consistent globally) + if v.dim() == 0: + self._shape_signature[k] = tuple() + else: + sig = list(v.shape) + sig.pop(batch_dim) + self._shape_signature[k] = tuple(sig) + self.graph_vars[k] = self._allocate_master_tensor( + k, v, max_bs=self._max_bs, bs=bs + ) + + def _initialize_bucket(self, bucket_bs: int, bs: int, **kwargs): + assert self._max_bs is not None + self._shape_guard_verified[bucket_bs] = False + if bucket_bs > self._max_bs: + raise ValueError( + f"CUDAGraph bucket_bs({bucket_bs}) exceeds master max_bs({self._max_bs}). " + f"Set CUDA_GRAPH_MAX_BS >= {bucket_bs}." + ) + + print(f"[CUDAGraph] Initializing bucket_bs={bucket_bs} (current bs={bs}) ...") + + # Build view map for this bucket (fixed address/shape/stride) + static_map: dict[str, torch.Tensor] = {} + for k, v in kwargs.items(): + # Shape guard: non-batch dims must match + expected = self._shape_signature.get(k) + if v.dim() == 0: + got = tuple() + else: + batch_dim = self._batch_dim.get(k, self._batch_dim_for_key(k, v)) + got_list = list(v.shape) + got_list.pop(batch_dim) + got = tuple(got_list) + if expected is not None and got != expected: + raise ValueError( + f"CUDAGraph bucket init: non-batch dim changed for key={k}, expected={expected}, got={got}" + ) + master = self.graph_vars[k] + batch_dim = self._batch_dim.get(k, 0) + static_map[k] = self._view_for_bucket( + master, bucket_bs=bucket_bs, batch_dim=batch_dim + ) + + # Warmup (stabilize kernels/caches) + out = None + for _ in range(self.warm_up_times): + out = self.forward_naive(**static_map) + assert isinstance(out, torch.Tensor) + + # Output master buffer (allocated once) + if "_outputs" not in self.graph_vars: + out_shape = (self._max_bs,) + tuple(out.shape[1:]) + self.graph_vars["_outputs"] = torch.empty( + out_shape, device=out.device, dtype=out.dtype + ) + + self.static_inputs_map[bucket_bs] = static_map + self.static_output_tensor[bucket_bs] = self.graph_vars["_outputs"][:bucket_bs] + + # capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=self.graph_pool): + self.static_output_tensor[bucket_bs].copy_(self.forward_naive(**static_map)) + if self.graph_pool is None: + self.graph_pool = graph.pool() + self.graphs[bucket_bs] = graph + torch.cuda.synchronize() + print(f"[CUDAGraph] Bucket {bucket_bs} captured.") + + def forward_naive(self, **kwargs): + return self.model(**kwargs) + + def forward(self, **kwargs): + if not self._enable: + return self.forward_naive(**kwargs) + bs = self._get_batch_size(**kwargs) + bucket_bs = self._select_bucket_bs(bs) + + # Initialize master buffers (first call) + self._ensure_master_buffers(bs=bs, **kwargs) + assert self._max_bs is not None + if bucket_bs > self._max_bs: + # Do not exceed master max_bs at runtime (would require resize and invalidate graphs) + return self.forward_naive(**kwargs) + + # lazy capture for this bucket + if bucket_bs not in self.graphs: + self._initialize_bucket(bucket_bs=bucket_bs, bs=bs, **kwargs) + + # Shape guard: non-batch dim change is unsafe for cudagraph; fallback to eager + # fast-path: skip guard if already verified for this bucket (stable-state optimization) + if not ( + _SHAPE_GUARD_FASTPATH and self._shape_guard_verified.get(bucket_bs, False) + ): + for k, v in kwargs.items(): + if not isinstance(v, torch.Tensor): + raise TypeError( + f"CUDAGraph_Wrapper only supports Tensor kwargs; got {k}={type(v)}" + ) + expected = self._shape_signature.get(k) + if v.dim() == 0: + got = tuple() + else: + batch_dim = self._batch_dim.get(k, self._batch_dim_for_key(k, v)) + got_list = list(v.shape) + got_list.pop(batch_dim) + got = tuple(got_list) + if expected is not None and got != expected: + return self.forward_naive(**kwargs) + # Mark verified for fast-path on subsequent calls + self._shape_guard_verified[bucket_bs] = True + + # Copy inputs into master buffer (per batch_dim, zero padding) + for k, tensor in kwargs.items(): + master = self.graph_vars[k] + batch_dim = self._batch_dim.get(k, 0) + self._copy_and_pad( + master, tensor, bs=bs, bucket_bs=bucket_bs, batch_dim=batch_dim + ) + + graph = self.graphs[bucket_bs] + graph.replay() + return self.graph_vars["_outputs"][:bs] diff --git a/wall_x/utils/logger.py b/wall_x/utils/logger.py new file mode 100644 index 00000000..eac25332 --- /dev/null +++ b/wall_x/utils/logger.py @@ -0,0 +1,95 @@ +"""Rank-aware text logger for distributed training.""" + +from __future__ import annotations + +import logging +import os +import sys +from pathlib import Path +from typing import Optional + +from torch.distributed import get_rank, is_initialized + + +class DistributedLogger: + def __init__( + self, + name: str = "wallx", + save_path: Optional[str] = None, + level: int = logging.INFO, + ): + if is_initialized(): + self._rank = get_rank() + else: + self._rank = int(os.environ.get("RANK", 0)) + logging.warning( + "DistributedLogger created before init_process_group; " + "falling back to RANK env var (rank=%d).", + self._rank, + ) + logger = logging.getLogger(f"{name}.rank{self._rank}") + logger.setLevel(level) + logger.propagate = False + for h in logger.handlers[:]: + h.close() + logger.removeHandler(h) + + fmt = logging.Formatter( + f"%(asctime)s - [rank{self._rank}] - %(levelname)s - %(message)s" + ) + + # All ranks write a file log (if a save_path was given). + if save_path: + log_dir = Path(save_path) / "logs" + log_dir.mkdir(parents=True, exist_ok=True) + file_handler = logging.FileHandler(log_dir / f"rank_{self._rank}.log") + file_handler.setFormatter(fmt) + logger.addHandler(file_handler) + + # Only rank 0 writes to stdout. + if self._rank == 0: + stream_handler = logging.StreamHandler(sys.stdout) + stream_handler.setFormatter(fmt) + logger.addHandler(stream_handler) + + self._logger = logger + + # Thin pass-throughs — callers use standard logging verbs. + def info(self, msg, *args, **kwargs): + self._logger.info(msg, *args, **kwargs) + + def warning(self, msg, *args, **kwargs): + self._logger.warning(msg, *args, **kwargs) + + def error(self, msg, *args, **kwargs): + self._logger.error(msg, *args, **kwargs) + + def debug(self, msg, *args, **kwargs): + self._logger.debug(msg, *args, **kwargs) + + @property + def rank(self) -> int: + return self._rank + + # --- backward-compatible shim ---------------------------------------- + # Some legacy call sites still invoke `.log(msg, level=..., + # main_process_only=...)`. Forward those to the standard logger so we + # don't break them during the migration; new code should use .info / + # .warning / .error / .debug directly. + def log( + self, + message, + level: int = logging.INFO, + main_process_only: bool = False, + ): + if main_process_only and self._rank != 0: + return + self._logger.log(level, message) + + # Pytorch's ``accelerate``-style fallback for code that still passes an + # accelerator object to the old constructor; accept and ignore it. + # Older codepaths can be migrated incrementally. + @classmethod + def legacy(cls, name: str, level: int = logging.INFO, accelerator=None): + del accelerator # ignored + return cls(name=name, save_path=None, level=level) diff --git a/wall_x/utils/metrics.py b/wall_x/utils/metrics.py new file mode 100644 index 00000000..6904f8b5 --- /dev/null +++ b/wall_x/utils/metrics.py @@ -0,0 +1,94 @@ +import torch +from typing import List + + +def get_action_accuracy( + gt: torch.FloatTensor, # [Batch_Size, Horizon, Action_Dim] + pred: torch.FloatTensor, + thresholds: List[float] = [0.1, 0.2], +) -> torch.FloatTensor: + device = gt.device + diff = torch.abs(gt - pred).reshape(-1, gt.shape[-1]) + + # get the percentage of diff lower than threshold for all action dimensions + accuracies = torch.zeros(len(thresholds), device=device) + for idx, threshold in enumerate(thresholds): + accuracy = torch.mean( + (torch.mean((diff < threshold).float(), dim=1) >= 1.0).float() + ) + accuracies[idx] = accuracy + return accuracies + + +def dtw_distance(seq1, seq2): + """ + Compute the Dynamic Time Warping distance between two sequences. + + ``seq1`` and ``seq2`` must have shape ``(T, D)``, where ``T`` is the + number of time steps and ``D`` is the feature dimension. + """ + n, m = seq1.shape[0], seq2.shape[0] + + seq1_d, seq2_d = seq1.double(), seq2.double() + dtw_matrix = torch.full( + (n + 1, m + 1), float("inf"), dtype=torch.float64, device=seq1.device + ) + dtw_matrix[0, 0] = 0 + + for i in range(1, n + 1): + for j in range(1, m + 1): + cost = torch.dist(seq1_d[i - 1], seq2_d[j - 1]) + dtw_matrix[i, j] = cost + torch.min( + torch.stack( + [ + dtw_matrix[i - 1, j], + dtw_matrix[i, j - 1], + dtw_matrix[i - 1, j - 1], + ] + ) + ) + + return dtw_matrix[n, m] + + +def frechet_distance(seq1, seq2): + """ + Compute the discrete Frechet distance between two sequences. + + ``seq1`` and ``seq2`` must have shape ``(T, D)``, where ``T`` is the + number of time steps and ``D`` is the feature dimension. + """ + n, m = seq1.shape[0], seq2.shape[0] + + seq1_d, seq2_d = seq1.double(), seq2.double() + frechet_matrix = torch.full( + (n, m), float("inf"), dtype=torch.float64, device=seq1.device + ) + frechet_matrix[0, 0] = torch.dist(seq1_d[0], seq2_d[0]) + + for j in range(1, m): + frechet_matrix[0, j] = torch.max( + frechet_matrix[0, j - 1], torch.dist(seq1_d[0], seq2_d[j]) + ) + + for i in range(1, n): + frechet_matrix[i, 0] = torch.max( + frechet_matrix[i - 1, 0], torch.dist(seq1_d[i], seq2_d[0]) + ) + + for i in range(1, n): + for j in range(1, m): + frechet_matrix[i, j] = torch.max( + torch.min( + torch.stack( + [ + frechet_matrix[i - 1, j], + frechet_matrix[i, j - 1], + frechet_matrix[i - 1, j - 1], + ] + ) + ), + torch.dist(seq1_d[i], seq2_d[j]), + ) + + return frechet_matrix[n - 1, m - 1] diff --git a/wall_x/utils/timers.py b/wall_x/utils/timers.py index 4233d19f..b1de4839 100644 --- a/wall_x/utils/timers.py +++ b/wall_x/utils/timers.py @@ -6,8 +6,11 @@ import torch from functools import wraps from contextlib import nullcontext +import logging import os +logger = logging.getLogger(__name__) + ENABLE_PERFORMANCE_TIMING = ( os.environ.get("ENABLE_PERFORMANCE_TIMING", "True").lower() == "true" ) @@ -32,23 +35,14 @@ def __exit__(self, exc_type, exc_value, traceback): torch.cuda.synchronize() end_time = time.perf_counter() cost_ms = (end_time - self.start_time) * 1e3 - print(f"\033[92m{self.msg} took {cost_ms:.3f} ms to execute\033[0m") + logger.info("%s took %.3f ms to execute", self.msg, cost_ms) ScopeTimer = ScopeTimerContext if ENABLE_PERFORMANCE_TIMING else nullcontext def timer(func, msg=None): - """ - Decorator to measure function execution time. - - Args: - func: Function to be timed - - Returns: - Wrapped function with timing functionality - """ - + """Decorator to measure function execution time.""" if msg is None: msg = func.__name__ else: @@ -58,44 +52,36 @@ def timer(func, msg=None): def wrapper(*args, **kwargs): with ScopeTimer(msg): result = func(*args, **kwargs) - return result return wrapper -# Helper functions to check for distributed environment def _is_distributed(): - """Checks if the current environment is set up for distributed training.""" return torch.distributed.is_available() and torch.distributed.is_initialized() def _get_world_size(): - """Safely retrieves the world size (number of processes).""" if _is_distributed(): return torch.distributed.get_world_size() return 1 def _get_rank(): - """Safely retrieves the rank of the current process.""" if _is_distributed(): return torch.distributed.get_rank() return 0 def _barrier(group=None): - """Safely executes a distributed barrier to synchronize processes.""" if _is_distributed(): torch.distributed.barrier(group=group) -# Dynamically set the all_gather function if torch.distributed.is_available(): try: dist_all_gather_func = torch.distributed.all_gather_into_tensor except AttributeError: - # Fallback to standard all_gather if all_gather_into_tensor is missing dist_all_gather_func = torch.distributed.all_gather else: dist_all_gather_func = None @@ -109,51 +95,35 @@ def __init__(self, name): @abstractmethod def start(self, barrier=False): - """Start the timer. - - Args: - barrier (bool, optional): Synchronizes ranks before starting. Defaults to False. - """ + """Start the timer, optionally syncing all ranks with a barrier first.""" pass @abstractmethod def stop(self, barrier=False): - """Stop the timer. - - Args: - barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False. - """ + """Stop the timer, optionally syncing all ranks with a barrier first.""" pass @abstractmethod def reset(self): - """Reset timer.""" + """Reset accumulated elapsed time to zero.""" pass @abstractmethod def elapsed(self, reset=True, barrier=False): - """Calculates the elapsed time and restarts timer. - - Args: - reset (bool, optional): Resets timer before restarting. Defaults to True. - barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False. - - Returns: - float: Elapsed time. - """ + """Return accumulated elapsed time in seconds; reset if reset=True.""" pass class DummyTimer(TimerBase): - """Dummy Timer.""" + """Dummy Timer — no-op placeholder used when log level exceeds threshold.""" def __init__(self): super().__init__("dummy timer") - def start(self, barrier=False, nvtx_push=False): + def start(self, barrier=False, nvtx_push=False, sync=False, **kwargs): return - def stop(self, barrier=False, nvtx_pop=False): + def stop(self, barrier=False, sync=False, **kwargs): return def reset(self): @@ -166,9 +136,6 @@ def elapsed(self, reset=True, barrier=False): ) def active_time(self): - """Returns the cumulative duration the timer has been active. - Note: Not supported for DummyTimer. - """ raise Exception( "active timer should not be used to calculate elapsed time, " "check if timer's log_level <= self._log_level." @@ -188,34 +155,18 @@ class Timer(TimerBase): """ def __init__(self, name): - """Initialize Timer. - - Args: - name (str): Name of the timer. - """ super().__init__(name) self._elapsed = 0.0 self._active_time = 0.0 self._started = False - # Note that None will default to the global process group self._barrier_group = None self._start_time = time.time() self.nvtx = False def set_barrier_group(self, barrier_group): - """Sets barrier group. - - Args: - barrier_group (ProcessGroup): Torch ProcessGroup for barrier. - """ self._barrier_group = barrier_group def start(self, barrier=False, nvtx_push=False, sync=False): - """Start the timer. - - Args: - barrier (bool, optional): Synchronizes ranks before starting. Defaults to False. - """ assert not self._started, "timer has already been started" if barrier: _barrier(group=self._barrier_group) @@ -228,11 +179,6 @@ def start(self, barrier=False, nvtx_push=False, sync=False): self.nvtx = True def stop(self, barrier=False, sync=False): - """Stop the timer. - - Args: - barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False. - """ if self.nvtx: nvtx.range_pop() assert self._started, "timer is not started" @@ -246,37 +192,21 @@ def stop(self, barrier=False, sync=False): self._started = False def reset(self): - """Reset timer.""" - # Don't reset _active_time self._elapsed = 0.0 self._started = False def elapsed(self, reset=True, barrier=False): - """Calculates the elapsed time and restarts timer. - - Args: - reset (bool, optional): Resets timer before restarting. Defaults to True. - barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False. - - Returns: - float: Elapsed time. - """ _started = self._started - # If the timing in progress, end it first. if self._started: self.stop(barrier=barrier) - # Get the elapsed time. _elapsed = self._elapsed - # Reset the elapsed time if reset: self.reset() - # If timing was in progress, set it back. if _started: self.start(barrier=barrier) return _elapsed def active_time(self): - """Calculates the cumulative duration for which the timer has been active""" return self._active_time @@ -284,13 +214,6 @@ class Timers: """Class for a group of Timers.""" def __init__(self, log_level, log_option): - """Initialize group of timers. - - Args: - log_level (int): Log level to control what timers are enabled. - log_option (str): Setting for logging statistics over ranks for all the timers. - Allowed: ['max', 'minmax', 'all']. - """ self._log_level = log_level allowed_log_options = set(["max", "minmax", "all"]) assert ( @@ -305,9 +228,6 @@ def __init__(self, log_level, log_option): self._max_log_level = 2 def __call__(self, name, log_level=None): - """Call timer with name and log level.""" - # If the timer has already been set, then check if the log-level - # is provided, it matches the one that the timer was created with. if name in self._timers: if log_level is not None: assert log_level == self._log_levels[name], ( @@ -317,8 +237,6 @@ def __call__(self, name, log_level=None): ) ) return self._timers[name] - # If timer does not exist and no log level is provided, - # set it to the max log level which is 2. if log_level is None: log_level = self._max_log_level assert ( @@ -326,38 +244,19 @@ def __call__(self, name, log_level=None): ), "log level {} is larger than max supported log level {}".format( log_level, self._max_log_level ) - # Now if the input log level is larger than the one set for - # the timers class, just ignore it and return a dummy timer. if log_level > self._log_level: return self._dummy_timer - # Otherwise, initalize the timer and set the level. self._timers[name] = Timer(name) self._log_levels[name] = log_level return self._timers[name] def _get_elapsed_time_all_ranks(self, names, reset, barrier): - """Returns elapsed times of timers in names. - - For single-node/single-GPU cases, directly returns the time for the current rank. - For distributed cases, maintains the existing all_gather logic. - - Args: - names (List[str]): list of timer names - reset (bool): reset the timer after recording the elapsed time - barrier (bool): if set, do a global barrier before time measurements - - Returns: - torch.tensor: Tensor of size [world_size, len(names)] with times in float. - """ - - # First make sure all the callers are in sync. if barrier: _barrier() world_size = _get_world_size() rank = _get_rank() - # Create device tensor if torch.cuda.is_available(): device = torch.cuda.current_device() else: @@ -367,33 +266,26 @@ def _get_elapsed_time_all_ranks(self, names, reset, barrier): (world_size, len(names)), dtype=torch.float, device=device ) - # Fill timing data for the current rank for i, name in enumerate(names): if name in self._timers: rank_name_to_time[rank, i] = self._timers[name].elapsed(reset=reset) - # Return directly for single-node; perform all_gather for distributed setup if world_size > 1 and _is_distributed() and dist_all_gather_func is not None: try: dist_all_gather_func( rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1) ) except Exception as e: - # If all_gather fails, print a warning and proceed with single rank timing - print(f"Warning: all_gather failed: {e}. Using single rank timing.") + logger.warning("all_gather failed: %s. Using single rank timing.", e) return rank_name_to_time def _get_global_min_max_time(self, names, reset, barrier, normalizer): - """Report only min and max times across all ranks.""" - rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier) name_to_min_max_time = {} for i, name in enumerate(names): rank_to_time = rank_name_to_time[:, i] - # filter out the ones we did not have any timings for rank_to_time = rank_to_time[rank_to_time > 0.0] - # If the timer exists: if rank_to_time.numel() > 0: name_to_min_max_time[name] = ( rank_to_time.min().item() / normalizer, @@ -404,7 +296,6 @@ def _get_global_min_max_time(self, names, reset, barrier, normalizer): def _get_global_min_max_time_string( self, names, reset, barrier, normalizer, max_only ): - """Report strings for max/minmax times across all ranks.""" name_to_min_max_time = self._get_global_min_max_time( names, reset, barrier, normalizer ) @@ -413,17 +304,13 @@ def _get_global_min_max_time_string( world_size = _get_world_size() if world_size == 1: - # Simplified output for single-node setup output_string = "time (ms):" for name in name_to_min_max_time: - _, max_time = name_to_min_max_time[ - name - ] # min and max are identical for a single rank + _, max_time = name_to_min_max_time[name] output_string += "\n {}: {:.2f}".format( (name + " ").ljust(48, "."), max_time ) else: - # Maintain original output format for multi-node setup if max_only: output_string = "max time across ranks (ms):" else: @@ -441,7 +328,6 @@ def _get_global_min_max_time_string( return output_string def _get_all_ranks_time_string(self, names, reset, barrier, normalizer): - """Report times across all ranks.""" rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier) world_size = _get_world_size() @@ -474,32 +360,20 @@ def get_all_timers_string( reset: bool = True, barrier: bool = False, ): - """Returns the output string with logged timer values according to configured options. + """Return a formatted timing string for the given timer names. Args: - names (List[str]): Names of the timers to log. If None, all registered timers are - fetched. Defaults to None. - normalizer (float, optional): Normalizes the timer values by the factor. - Defaults to 1.0. - reset (bool, optional): Whether to reset timer values after logging. Defaults to True. - barrier (bool, optional): Whether to do a global barrier before time measurments. - Defaults to False. - - Raises: - Exception: Raises if log option is invalid. - - Returns: - str: Formatted string with the timer values. + names: Timers to include; defaults to all registered timers. + normalizer: Divide raw seconds by this value (e.g. 1000 for ms output). + reset: Reset each timer after reading its elapsed time. + barrier: Synchronize across ranks before gathering times. """ - - if names is None: # get all registered timers + if names is None: names = list(self._timers.keys()) assert normalizer > 0.0 if self._log_option in ["max", "minmax"]: - max_only = False - if self._log_option == "max": - max_only = True + max_only = self._log_option == "max" output_string = self._get_global_min_max_time_string( names, reset, barrier, normalizer / 1000.0, max_only ) @@ -519,30 +393,23 @@ def log( reset: bool = True, barrier: bool = False, ): - """logs the timers passed in names to stdout. Example usage is to log average per step - value for timer 'foo', this function can be called with normalizer factor set to logging - interval. + """Print timing results for the given names to stdout on one rank. Args: - names (List[str]): Names of the timers to log. - rank (int, optional): logs the timers to a specific rank. If set to None, logs to the - last rank. Defaults to None. - normalizer (float, optional): Normalizes the timer values by the factor. - Defaults to 1.0. - reset (bool, optional): Whether to reset timer values after logging. Defaults to True. - barrier (bool, optional): Whether to do a global barrier before time measurments. - Defaults to False. + names: Timer names to log. + rank: Rank that prints; defaults to the last rank (world_size - 1). + normalizer: Divide raw seconds by this value before printing. + reset: Reset each timer after reading. + barrier: Synchronize across ranks first. """ - output_string = self.get_all_timers_string(names, normalizer, reset, barrier) - # If no input rank is provided, log on last rank. world_size = _get_world_size() current_rank = _get_rank() if rank is None: rank = world_size - 1 if rank == current_rank and output_string is not None: - print(output_string, flush=True) + logger.info("%s", output_string) def write( self, @@ -553,22 +420,16 @@ def write( reset: bool = True, barrier: bool = False, ): - """Write timers to a tensorboard writer. - Note that we only report maximum time across ranks to tensorboard. + """Write per-timer max times as TensorBoard scalars. Args: - names (List[str]): Names of the timers to log. - writer (SummaryWriter): Tensorboard SummaryWriter object - iteration (int): Current iteration. - normalizer (float, optional): Normalizes the timer values by the factor. - Defaults to 1.0. - reset (bool, optional): Whether to reset timer values after logging. Defaults to True. - barrier (bool, optional): Whether to do a global barrier before time measurments. - Defaults to False. + names: Timer names to write. + writer: TensorBoard SummaryWriter instance. + iteration: Global step value for the scalar. + normalizer: Divide raw seconds by this value. + reset: Reset each timer after reading. + barrier: Synchronize across ranks first. """ - # currently when using add_scalars, - # torch.utils.add_scalars makes each timer its own run, which - # polutes the runs list, so we just add each as a scalar assert normalizer > 0.0 name_to_min_max_time = self._get_global_min_max_time( names, reset, barrier, normalizer diff --git a/wall_x/utils/weight_convert/__init__.py b/wall_x/utils/weight_convert/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/wall_x/utils/weight_convert/convert.py b/wall_x/utils/weight_convert/convert.py new file mode 100644 index 00000000..2889ea6b --- /dev/null +++ b/wall_x/utils/weight_convert/convert.py @@ -0,0 +1,215 @@ +""" +Test script for weight loader - demonstrates usage with the ACT decode model. +""" + +from pathlib import Path + +import torch + +from weight_loader import WeightLoader + + +def test_weight_loader_basic(): + """Test basic weight loading and conversion.""" + print("=" * 80) + print("Test 1: Basic Weight Loading and Conversion") + print("=" * 80) + + source_path = "/x2robot_v2/share/liangyuxin/share_ckpts/multi_task_moe_1030/40/model.safetensors" + output_path = "/x2robot_v2/vincent/workspace/inference/tests/converted_weights_lucy.safetensors" + + # Load and convert weights + loader = WeightLoader( + source_path=source_path, + expert_index=1, # Use expert 1 as specified + verbose=True, + ) + + converted_weights = loader.convert() + loader.save(output_path) + + print("\n✓ Conversion successful!") + print(f"✓ Saved to: {output_path}") + + return converted_weights + + +def test_load_into_model(): + """Test loading converted weights into the ACT decode model.""" + print("\n" + "=" * 80) + print("Test 2: Loading Converted Weights into Model") + print("=" * 80) + + try: + from wall_x.infer.qwen2_5_based.modeling_qwen2_5_vl_act import ( + Qwen2_5_VL_ACT_Decode, + ) + from transformers import AutoConfig + + # Load config + config_path = "/x2robot_v2/vincent/workspace/inference/wall-x/workspace/models_config/moe_flash.json" + config = AutoConfig.from_pretrained(config_path) + + # Create model + print("\nCreating model...") + model = Qwen2_5_VL_ACT_Decode(config) + + converted_path = "/x2robot_v2/vincent/workspace/inference/tests/converted_weights_lucy.safetensors" + + if Path(converted_path).exists(): + from safetensors.torch import load_file + + state_dict = load_file(converted_path) + + # Load state dict + print(f"\nLoading weights from: {converted_path}") + missing_keys, unexpected_keys = model.load_state_dict( + state_dict, strict=False + ) + + print("\n✓ Weights loaded successfully!") + + if missing_keys: + print(f"\nMissing keys ({len(missing_keys)}):") + for key in missing_keys[:10]: # Show first 10 + print(f" - {key}") + if len(missing_keys) > 10: + print(f" ... and {len(missing_keys) - 10} more") + + if unexpected_keys: + print(f"\nUnexpected keys ({len(unexpected_keys)}):") + for key in unexpected_keys[:10]: # Show first 10 + print(f" - {key}") + if len(unexpected_keys) > 10: + print(f" ... and {len(unexpected_keys) - 10} more") + + return model, state_dict + else: + print(f"\n✗ Converted weights not found at: {converted_path}") + print(" Run test_weight_loader_basic() first to create converted weights") + return None, None + + except Exception as e: + print(f"\n✗ Error loading into model: {e}") + import traceback + + traceback.print_exc() + return None, None + + +def inspect_weight_shapes(): + """Inspect and compare weight shapes between source and target.""" + print("\n" + "=" * 80) + print("Test 3: Weight Shape Inspection") + print("=" * 80) + + source_path = "/x2robot_v2/share/liangyuxin/share_ckpts/multi_task_moe_1030/40/model.safetensors" + + if not Path(source_path).exists(): + print(f"✗ Source file not found: {source_path}") + return + + from safetensors.torch import load_file + + source_weights = load_file(source_path) + + print("\nSource weight shapes (first layer):") + print("-" * 80) + + layer_idx = 0 + + # QKV weights + q_key = f"model.layers.{layer_idx}.self_attn.q_proj.weight" + k_key = f"model.layers.{layer_idx}.self_attn.k_proj.weight" + v_key = f"model.layers.{layer_idx}.self_attn.v_proj.weight" + + if q_key in source_weights: + q_shape = source_weights[q_key].shape + k_shape = source_weights[k_key].shape + v_shape = source_weights[v_key].shape + + print(f"Q proj: {q_shape}") + print(f"K proj: {k_shape}") + print(f"V proj: {v_shape}") + print(f"→ Merged QKV: ({q_shape[0] + k_shape[0] + v_shape[0]}, {q_shape[1]})") + + # Gate/Up weights + gate_key = f"model.layers.{layer_idx}.moe.experts.1.gate_proj.weight" + up_key = f"model.layers.{layer_idx}.moe.experts.1.up_proj.weight" + + if gate_key in source_weights: + gate_shape = source_weights[gate_key].shape + up_shape = source_weights[up_key].shape + + print(f"\nGate proj (expert 1): {gate_shape}") + print(f"Up proj (expert 1): {up_shape}") + print(f"→ Merged gate_up: ({gate_shape[0] + up_shape[0]}, {gate_shape[1]})") + + converted_path = "/x2robot_v2/vincent/workspace/inference/tests/converted_weights_lucy.safetensors" + if Path(converted_path).exists(): + converted_weights = load_file(converted_path) + + print("\n" + "-" * 80) + print("Converted weight shapes (first layer):") + print("-" * 80) + + qkv_key = f"layers.{layer_idx}.self_attn.qkv_proj.weight" + if qkv_key in converted_weights: + print(f"QKV merged: {converted_weights[qkv_key].shape}") + + gate_up_key = f"layers.{layer_idx}.act_expert_mlp.gate_up_proj.weight" + if gate_up_key in converted_weights: + print(f"gate_up merged: {converted_weights[gate_up_key].shape}") + + +def compare_forward_pass(): + """Compare forward pass results (if possible).""" + print("\n" + "=" * 80) + print("Test 4: Forward Pass Verification") + print("=" * 80) + + print("\nCreating dummy input...") + + # Create a small dummy input + batch_size = 1 + seq_len = 10 + hidden_size = 4096 # Adjust based on your model + + hidden_states = torch.randn(batch_size, seq_len, hidden_size) + + print(f"Input shape: {hidden_states.shape}") + print( + "\nNote: Full forward pass test requires proper model initialization and KV cache." + ) + print("This is a placeholder for more detailed testing.") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Test weight loader functionality") + parser.add_argument( + "--mode", + type=str, + choices=["all", "convert", "load", "inspect", "forward"], + default="all", + help="Which test to run", + ) + + args = parser.parse_args() + + if args.mode == "all" or args.mode == "convert": + converted_weights = test_weight_loader_basic() + + if args.mode == "all" or args.mode == "load": + test_load_into_model() + + if args.mode == "all" or args.mode == "inspect": + inspect_weight_shapes() + + if args.mode == "all" or args.mode == "forward": + compare_forward_pass() + + print("\n" + "=" * 80) + print("All tests completed!") + print("=" * 80) diff --git a/wall_x/utils/weight_convert/weight_loader.py b/wall_x/utils/weight_convert/weight_loader.py new file mode 100644 index 00000000..2429da54 --- /dev/null +++ b/wall_x/utils/weight_convert/weight_loader.py @@ -0,0 +1,377 @@ +""" +Weight loader for converting HuggingFace/standard format weights to merged QKV and gate_up format. + +This loader supports: +1. Merging Q, K, V projections into a single qkv_proj weight +2. Merging gate and up projections into a single gate_up_proj weight +3. Remapping layer names from model.layers.X to layers.X +4. Using expert 1 for MoE models (configurable) +""" + +import torch +from safetensors.torch import save_file +from typing import Dict, Optional +from pathlib import Path + + +class WeightLoader: + def __init__( + self, + config: Dict, + source_weights: Dict[str, torch.Tensor], + expert_index: int = 1, + num_layers: Optional[int] = None, + verbose: bool = True, + ): + """ + Initialize the weight loader. + + Args: + source_path: Path to source weights (safetensors or pytorch checkpoint) + expert_index: Which expert to use for MoE models (default: 1) + num_layers: Number of layers in the model (auto-detected if None) + verbose: Print conversion progress + """ + self.config = config + self.source_weights = source_weights + self.expert_index = expert_index + self.verbose = verbose + + # Auto-detect number of layers if not provided + if num_layers is None: + self.num_layers = self._detect_num_layers() + else: + self.num_layers = num_layers + + if self.verbose: + print(f"Number of layers detected: {self.num_layers}") + print(f"Using expert index: {expert_index}") + + def _detect_num_layers(self) -> int: + """Auto-detect number of layers from weight keys.""" + layer_indices = set() + for key in self.source_weights.keys(): + if "layers." in key: + # Extract layer index + parts = key.split(".") + for i, part in enumerate(parts): + if part == "layers" and i + 1 < len(parts): + try: + layer_indices.add(int(parts[i + 1])) + except ValueError: + continue + return max(layer_indices) + 1 if layer_indices else 0 + + def _get_weight(self, key: str) -> Optional[torch.Tensor]: + """Safely get a weight tensor, returning None if not found.""" + return self.source_weights.get(key) + + def _merge_qkv(self, layer_idx: int) -> Dict[str, torch.Tensor]: + """ + Merge Q, K, V projection weights and biases for a layer. + + Expected source format: + model.layers.{layer_idx}.self_attn.q_proj.weight + model.layers.{layer_idx}.self_attn.k_proj.weight + model.layers.{layer_idx}.self_attn.v_proj.weight + (and corresponding .bias if present) + + Target format: + layers.{layer_idx}.self_attn.qkv_proj.qkv_proj.weight + layers.{layer_idx}.self_attn.qkv_proj.qkv_proj.bias + + Returns: + Dictionary with merged weights + """ + result = {} + + # Get Q, K, V weights + if self.config.attention_moe: + q_weight = self._get_weight( + f"model.layers.{layer_idx}.self_attn.q_proj_experts.{self.expert_index}.weight" + ) + k_weight = self._get_weight( + f"model.layers.{layer_idx}.self_attn.k_proj_experts.{self.expert_index}.weight" + ) + v_weight = self._get_weight( + f"model.layers.{layer_idx}.self_attn.v_proj_experts.{self.expert_index}.weight" + ) + else: + q_weight = self._get_weight( + f"model.layers.{layer_idx}.self_attn.q_proj.weight" + ) + k_weight = self._get_weight( + f"model.layers.{layer_idx}.self_attn.k_proj.weight" + ) + v_weight = self._get_weight( + f"model.layers.{layer_idx}.self_attn.v_proj.weight" + ) + + if q_weight is None or k_weight is None or v_weight is None: + merged_weight = self._get_weight( + f"model.layers.{layer_idx}.self_attn.qkv_proj.weight" + ) + if merged_weight is None: + merged_weight = self._get_weight( + f"model.layers.{layer_idx}.self_attn.qkv_proj_experts.{self.expert_index}.weight" + ) + + # Merge along output dimension (dim 0) + if q_weight is not None and k_weight is not None and v_weight is not None: + merged_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) + if merged_weight is None: + raise ValueError( + f"Missing QKV weights for layer {layer_idx} (expert {self.expert_index})" + ) + result[f"layers.{layer_idx}.self_attn.qkv_proj.qkv_proj.weight"] = merged_weight + + # Get Q, K, V biases (if they exist) + if self.config.attention_moe: + q_bias = self._get_weight( + f"model.layers.{layer_idx}.self_attn.q_proj_experts.{self.expert_index}.bias" + ) + k_bias = self._get_weight( + f"model.layers.{layer_idx}.self_attn.k_proj_experts.{self.expert_index}.bias" + ) + v_bias = self._get_weight( + f"model.layers.{layer_idx}.self_attn.v_proj_experts.{self.expert_index}.bias" + ) + else: + q_bias = self._get_weight(f"model.layers.{layer_idx}.self_attn.q_proj.bias") + k_bias = self._get_weight(f"model.layers.{layer_idx}.self_attn.k_proj.bias") + v_bias = self._get_weight(f"model.layers.{layer_idx}.self_attn.v_proj.bias") + + if q_bias is None or k_bias is None or v_bias is None: + merged_bias = self._get_weight( + f"model.layers.{layer_idx}.self_attn.qkv_proj.bias" + ) + if merged_bias is None: + merged_bias = self._get_weight( + f"model.layers.{layer_idx}.self_attn.qkv_proj_experts.{self.expert_index}.bias" + ) + else: + merged_bias = torch.cat([q_bias, k_bias, v_bias], dim=0) + result[f"layers.{layer_idx}.self_attn.qkv_proj.qkv_proj.bias"] = merged_bias + + if self.verbose: + print(f"Layer {layer_idx}: Merged QKV - shape {merged_weight.shape}") + + return result + + def _merge_gate_up(self, layer_idx: int) -> Dict[str, torch.Tensor]: + """ + Merge gate and up projection weights for a layer. + + Expected source format (MoE): + model.layers.{layer_idx}.moe.experts.{expert_idx}.gate_proj.weight + model.layers.{layer_idx}.moe.experts.{expert_idx}.up_proj.weight + + Alternative source format (standard MLP): + model.layers.{layer_idx}.mlp.gate_proj.weight + model.layers.{layer_idx}.mlp.up_proj.weight + + Target format: + layers.{layer_idx}.act_expert_mlp.gate_up_proj.weight + + Returns: + Dictionary with merged weights + """ + result = {} + + # Try MoE format first + gate_weight = self._get_weight( + f"model.layers.{layer_idx}.moe.experts.{self.expert_index}.gate_proj.weight" + ) + up_weight = self._get_weight( + f"model.layers.{layer_idx}.moe.experts.{self.expert_index}.up_proj.weight" + ) + + # If not found, try standard MLP format + if gate_weight is None or up_weight is None: + gate_weight = self._get_weight( + f"model.layers.{layer_idx}.mlp.gate_proj.weight" + ) + up_weight = self._get_weight(f"model.layers.{layer_idx}.mlp.up_proj.weight") + + if gate_weight is None or up_weight is None: + merged_weight = self._get_weight( + f"model.layers.{layer_idx}.moe.experts.{self.expert_index}.gate_up_proj.weight" + ) + else: + merged_weight = torch.cat([gate_weight, up_weight], dim=0) + + if merged_weight is None: + raise ValueError( + f"Missing gate/up weights for layer {layer_idx} (expert {self.expert_index})" + ) + result[f"layers.{layer_idx}.act_expert_mlp.gate_up_proj.weight"] = merged_weight + + if self.verbose: + print(f"Layer {layer_idx}: Merged gate_up - shape {merged_weight.shape}") + + return result + + def _copy_other_weights(self, layer_idx: int) -> Dict[str, torch.Tensor]: + """ + Copy other weights that don't need merging (o_proj, down_proj, layer norms, etc.) + + Mapping: + model.layers.X.self_attn.o_proj.weight -> layers.X.self_attn.o_proj.o_proj.weight + model.layers.X.moe.experts.{expert}.down_proj.weight -> layers.X.act_expert_mlp.down_proj.weight + model.layers.X.input_layernorm.weight -> layers.X.input_layernorm.weight + model.layers.X.post_attention_layernorm.weight -> layers.X.post_attention_layernorm.weight + """ + result = {} + + # O projection + if self.config.attention_moe: + o_weight = self._get_weight( + f"model.layers.{layer_idx}.self_attn.o_proj_experts.{self.expert_index}.weight" + ) + else: + o_weight = self._get_weight( + f"model.layers.{layer_idx}.self_attn.o_proj.weight" + ) + if o_weight is not None: + result[f"layers.{layer_idx}.self_attn.o_proj.o_proj.weight"] = o_weight + else: + raise ValueError( + f"Missing o_proj weight for layer {layer_idx} (expert {self.expert_index})" + ) + + # Down projection (MoE format) + down_weight = self._get_weight( + f"model.layers.{layer_idx}.moe.experts.{self.expert_index}.down_proj.weight" + ) + + # If not found, try standard MLP format + if down_weight is None: + down_weight = self._get_weight( + f"model.layers.{layer_idx}.mlp.down_proj.weight" + ) + if down_weight is None: + raise ValueError( + f"Missing down_proj weight for layer {layer_idx} (expert {self.expert_index})" + ) + + result[f"layers.{layer_idx}.act_expert_mlp.down_proj.weight"] = down_weight + + # Layer norms + if self.config.norm_moe: + input_ln = self._get_weight( + f"model.layers.{layer_idx}.input_layernorms.{self.expert_index}.weight" + ) + else: + input_ln = self._get_weight( + f"model.layers.{layer_idx}.input_layernorm.weight" + ) + if input_ln is not None: + result[f"layers.{layer_idx}.input_layernorm.weight"] = input_ln + + if self.config.norm_moe: + post_attn_ln = self._get_weight( + f"model.layers.{layer_idx}.post_attention_layernorms.{self.expert_index}.weight" + ) + else: + post_attn_ln = self._get_weight( + f"model.layers.{layer_idx}.post_attention_layernorm.weight" + ) + if post_attn_ln is not None: + result[f"layers.{layer_idx}.post_attention_layernorm.weight"] = post_attn_ln + + return result + + def _copy_embeddings_and_norms(self) -> Dict[str, torch.Tensor]: + """Copy embedding and final norm weights.""" + result = {} + + # Embeddings + embed = self._get_weight("model.embed_tokens.weight") + print("embed.shape: ", embed.shape) + if embed is not None: + result["embed_tokens.weight"] = embed + + # Final norm + if self.config.norm_moe: + norm = self._get_weight(f"model.norms.{self.expert_index}.weight") + else: + norm = self._get_weight("model.norm.weight") + print("norm.shape: ", norm.shape) + if norm is not None: + result["norm.weight"] = norm + + # LM head (if present) + lm_head = self._get_weight("lm_head.weight") + if lm_head is not None: + result["lm_head.weight"] = lm_head + + return result + + def convert(self) -> Dict[str, torch.Tensor]: + """ + Convert all weights to the target format. + + Returns: + Dictionary with converted weights in target format + """ + converted_weights = {} + if self.verbose: + print("\n" + "=" * 60) + print("Starting weight conversion...") + print("=" * 60) + + # Convert each layer + for layer_idx in range(self.num_layers): + if self.verbose: + print(f"\nProcessing layer {layer_idx}...") + + # Merge QKV + qkv_weights = self._merge_qkv(layer_idx) + converted_weights.update(qkv_weights) + + # Merge gate_up + gate_up_weights = self._merge_gate_up(layer_idx) + converted_weights.update(gate_up_weights) + + # Copy other weights + other_weights = self._copy_other_weights(layer_idx) + converted_weights.update(other_weights) + + # Copy embeddings and final norms + if self.verbose: + print("\nProcessing embeddings and final norms...") + embed_weights = self._copy_embeddings_and_norms() + converted_weights.update(embed_weights) + + if self.verbose: + print("\n" + "=" * 60) + print(f"Conversion complete! Total weights: {len(converted_weights)}") + print("=" * 60) + + return converted_weights + + def save( + self, + output_path: str, + converted_weights: Optional[Dict[str, torch.Tensor]] = None, + ): + """ + Save converted weights to file. + + Args: + output_path: Path to save converted weights + converted_weights: Pre-converted weights (if None, will run conversion) + """ + if converted_weights is None: + converted_weights = self.convert() + + output_path = Path(output_path) + + if output_path.suffix == ".safetensors": + save_file(converted_weights, str(output_path)) + else: + torch.save(converted_weights, str(output_path)) + + if self.verbose: + print(f"\nSaved converted weights to: {output_path}") + print(f"File size: {output_path.stat().st_size / (1024**3):.2f} GB") diff --git a/workspace/README.md b/workspace/README.md index b6a96774..ae2b8ff5 100644 --- a/workspace/README.md +++ b/workspace/README.md @@ -1,160 +1,355 @@ -# Training Guide +# Wall-X-OSS Usage Guide -This document explains the key configuration parameters and memory requirements for Wall-X training. +This guide explains how to fine-tune, evaluate in simulation, and deploy on real robots with the **Wall-OSS-0.5** pretrained model on LeRobot-format datasets. -## Quick Start Checklist +> All commands below assume you are in the **repository root** (`wall-x/`). -### 🚀 **Step 1: Prepare Model** -Choose one of our pretrained models: -- **WALL-OSS-FLOW**: https://huggingface.co/x-square-robot/wall-oss-flow -- **WALL-OSS-FAST**: https://huggingface.co/x-square-robot/wall-oss-fast -Or from Qwen-2.5-VL -- Download https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct, settings refer to `config_qact_from_vlm.yml` +> **Note:** This open-source release targets **Wall-OSS-0.5**. If you are using **Wall-OSS-FLOW** or **Wall-OSS-FAST** instead, switch back to the previous codebase version: +> +> ```bash +> git checkout 97406f2ab5de414c79b091873f946c112d105c72 +> ``` -### ⚙️ **Step 2: Configure Environment** -- Update `run.sh`: Set `code_dir` and `config_path` to your actual paths -- Set `CUDA_VISIBLE_DEVICES` for your available GPUs +--- -### 📝 **Step 3: Update Configuration Files** -- Replace all `/path/to/` placeholders in `config_qact.yml` with actual paths -- Configure robot settings: `dof_config` and `agent_pos_config` -- Set dataset: Choose appropriate `repo_id` -- Adjust `batch_size_per_gpu` based on your GPU memory +## Environment Setup -### ▶️ **Step 4: Start Training** ```bash -bash ./workspace/lerobot_example/run.sh +conda create --name wallx python=3.10 +conda activate wallx + +pip install torch torchvision transformers +pip install huggingface_hub lerobot pyyaml wandb websockets msgpack msgpack-numpy opencv-python + +# Install wall-x (from repo root) +pip install -e . ``` -## Enable FAST tokenizer -To fine-tune using the FAST tokenizer, please download the repository and update the `action_tokenizer_path`. Make sure to set `use_fast_tokenizer` to `true` and q01 and q99 to normalize the dataset, refer to `wall-x/scripts/compute_norm_stats.py`: +--- + +## Part 1: Download Wall-OSS-0.5 Weights + +Wall-OSS-0.5 is a VLA foundation model built on Qwen2.5-VL-3B. Fine-tuning requires both the **pretrained weights** and the **VLM processor**. + +### 1.1 Download Wall-OSS-0.5 + ```bash -git clone https://huggingface.co/physical-intelligence/fast +# Option A: huggingface-cli +huggingface-cli download X-Square-Robot/wall-oss-0.5 \ + --local-dir /path/to/wall-oss-0.5 + +# Option B: Python +python -c " +from huggingface_hub import snapshot_download +snapshot_download('X-Square-Robot/wall-oss-0.5', local_dir='/path/to/wall-oss-0.5') +" ``` -## Required Paths (Must Modify) -```yaml -pretrained_wallx_path: "/path/to/wallx_model/" # Path to pretrained wallx model -save_path: "/path/to/workspace/" # Path to save training outputs -use_fast_tokenizer: False # True: train FAST, False: train Flow -action_tokenizer_path: "/path/to/fast/" # Must set if use_fast_tokenizer is True -norm_stats_path: "/path/to/stats/" # Must set for normalize dataset -``` -## Customize your robot configuration -Ensure that the sum of the configuration dimensions corresponds to the values specified in norm_stats.json, and that each key is unique. The maximum dimensionality is set to 20, consistent with our robot configuration. +After download, the directory should contain at least: + +| File | Purpose | +|------|---------| +| `config.json` | Model architecture; maps to `model.config_path` in training YAML | +| `model.safetensors` | Pretrained weights; maps to `checkpoint.resume_from` | +| tokenizer / processor files | Used at inference time | + +HuggingFace: + +### 1.2 Download Qwen2.5-VL-3B-Instruct (processor) + +Set `model.processor_path` and `model.pretrained_path` to the Qwen2.5-VL processor directory: + +```bash +huggingface-cli download Qwen/Qwen2.5-VL-3B-Instruct \ + --local-dir /path/to/Qwen2.5-VL-3B-Instruct +``` + +--- + +## Part 2: Simulation Tasks (LIBERO) + +### 2.1 Download the LIBERO Dataset + +Download the LeRobot-format LIBERO dataset from HuggingFace: + +```bash +huggingface-cli download lerobot/libero \ + --repo-type dataset \ + --local-dir /path/to/libero_all +``` + +Dataset page: + +### 2.2 Edit Your Training Config (`config.yml`) + +Copy an example config and replace every `/path/to/*` placeholder: + +```bash +cp workspace/example/fintune_bus2602/libero.yml /path/to/my_libero_config.yml +``` + +**Required paths:** + ```yaml - customized_dof_config: - "action_eef": 6 - "action_gripper": 1 +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 # local LeRobot dataset root + norm_stats_path: /path/to/libero_all_norm_stats.json + key_mappings: # must match your dataset keys + 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 +``` + +See `workspace/example/fintune_bus2602/libero.yml` for the full example. LIBERO uses a 7-dim single-arm action; pad to 26 dims with `action_padding` to match the Wall-OSS-0.5 pretraining space (see comments in the YAML). + +### 2.3 Compute LIBERO Normalization Stats + +Generate `norm_stats.json` from the dataset before training: + +```bash +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 +``` + +Then set `data.norm_stats_path` to the generated JSON file. + +### 2.4 Start Training - customized_agent_pos_config: - "state_eef_with_gripper": 7 +```bash +# Single GPU +CUDA_VISIBLE_DEVICES=0 \ + python wall_x/trainer/fsdp_trainer/train_fsdp.py \ + --config /path/to/my_libero_config.yml + +# Multi-GPU (recommended) +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 ``` -## Using Lerobot Dataset -- Each dataset employs distinct keys; please specify the corresponding key mappings as described in `wall-x/wall_x/data/utils.py`. -```python -"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", - } +Logs and checkpoints are written to `checkpoint.save_path`. If training saved FSDP-sharded checkpoints, merge them before inference: + +```bash +python scripts/merge_sharded_weights.py \ + /path/to/sharded_checkpoint \ + /path/to/merged_checkpoint ``` -## Compute stats +Single-GPU training needs at least **48 GB** VRAM. For multi-GPU runs, enable `distributed.use_fsdp: true`. + +### 2.5 Run Inference (LIBERO Simulation) + +Use `scripts/run_libero.sh` for batch evaluation in the LIBERO simulator: + ```bash - python wall-x/scripts/compute_norm_stats.py +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 + +# Quick smoke test (1 trial per task) +SMOKE=1 CHECKPOINT_PATH=/path/to/checkpoint bash scripts/run_libero.sh ``` -## Configuration Explain -- `agent_pos_config` corresponds to `obs_action_keys` and subsequently to state, while `dof_config` corresponds to `predict_action_keys` and subsequently to action. Note that the state and action may not necessarily share the same set of DoF. +Useful environment variables: + +| Variable | Description | +|----------|-------------| +| `CHECKPOINT_PATH` | Checkpoint directory | +| `TRAIN_CONFIG_PATH` | Training YAML | +| `TASK_SUITE_NAME` | `libero_spatial` / `libero_object` / `libero_goal` / `libero_10` | +| `ALL_SUITES=1` | Run all 4 standard suites sequentially | +| `TASK_INDICES` | Task indices, e.g. `0,1,2` | +| `CUDA_ID` | GPU index | + +--- -## Training Parameters (Commonly Modified) +## Part 3: Real-Robot Tasks -### Learning Rate Settings -- `learning_rate`: Initial learning rate (default: 0.00009) -- `min_lr`: Minimum learning rate for scheduler (default: 0.00005) -- `num_warmup_steps`: Number of warmup steps (default: 100) +### 3.1 Prepare Your Real-Robot Training Dataset -### Batch Size and Memory -- `batch_size_per_gpu`: Batch size per GPU - adjust based on GPU memory -- `gradient_accumulation_steps`: Gradient accumulation steps -- `num_training_steps`: Total training steps -- `num_epoch`: Number of training epochs +Convert your robot data to **LeRobot v3** format and note the local dataset root path (used for `data.lerobot_config.repo_id`). -### Training Optimization Settings -- `FSDP2`: Enable FSDP2 for distributed training (default: True) - **Recommended for multi-GPU** -- `torch_compile`: Enable PyTorch compilation optimization (default: False) +Reference config: `workspace/example/fintune_bus2602/maniparena_example.yml` (dual-arm, 448px, 3 cameras). -**⚠️ Important Note on torch_compile:** -- **Benefits**: Enabling `torch_compile` can significantly improve training efficiency -- **Requirements**: Requires that the data input shape is always consistent throughout training -- **Caution**: If you don't have sufficient understanding of torch compile, please **DO NOT** enable it as it may cause unexpected issues with dynamic input shapes +### 3.2 Edit Your Training Config (`config.yml`) -## Robot Configuration (Modify for Your Robot) +```bash +cp workspace/example/fintune_bus2602/maniparena_example.yml /path/to/my_robot_config.yml +``` + +**Required paths:** + +```yaml +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/your_robot_dataset + norm_stats_path: /path/to/your_robot_norm_stats.json + key_mappings: # match your dataset camera / state / action keys + camera: + observation.images.faceImg: face_view + # ... + state: observation.state + action: action + +checkpoint: + save_path: /path/to/robot_training_output + resume_from: /path/to/wall-oss-0.5/model.safetensors +``` + +`task.dof_config` defines **predicted action dimensions**; `task.agent_pos_config` defines **observation state dimensions**. The sum of each config must match `norm_stats.json`. If your robot has fewer than 26 DOF, pad with `action_padding` (see comments in the example YAML). + +### 3.3 Compute Dataset Normalization Stats + +```bash +python scripts/compute_norm_stats.py \ + --train_config /path/to/my_robot_config.yml \ + --data_root /path/to/your_robot_dataset \ + --output_path /path/to/your_robot_norm_stats.json +``` -### DOF Configuration -Modify `dof_config` to match your robot's action space: -- Add/remove action keys based on your robot's capabilities -- Ensure DOF numbers match your robot's action dimensions +### 3.4 Start Training -### Agent Position Configuration -Keep `agent_pos_config` consistent with `dof_config`. +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ + torchrun --nproc_per_node=4 \ + wall_x/trainer/fsdp_trainer/train_fsdp.py \ + --config /path/to/my_robot_config.yml +``` -### Action Keys -- `obs_action_keys`: Actions used as observation context -- `predict_action_keys`: Actions to predict/control +To resume training, point `checkpoint.resume_from` to a checkpoint **directory** (not a single `.safetensors` file). -## Data Configuration +### 3.5 Start Inference (WebSocket Server) -### Dataset -- `repo_id`: LeRobot dataset identifier -- `train_test_split`: Training/validation split ratio (default: 0.95) -- `action_horizon`: Number of future actions to predict (default: 32) +Use `scripts/run_serving.sh` to launch the inference server for real-robot clients or open-loop evaluation: -### Image Settings -- `resolution`: Image resolution for different camera views -- `download_videos`: Whether to download video files (true/false) +```bash +# Edit CKPT_PATH in the script, then run: +bash scripts/run_serving.sh +``` -## Resume Training (Optional) -- `resume.ckpt`: Path to checkpoint for resuming training -- `resume.load_ckpt_only`: Only load model weights, not optimizer state +Or invoke directly (adjust parameters as needed): -## Merge checkpoint -- If FSDP SHARDED_STATE_DICT is used, please run command below to merge checkpoint into a single safetensors ```bash - # refer to accelerate/commands/merge.py - accelerate merge-weights /path/to/sharded_tensors /path/to/model.safetensors - # copy the saved processor files - cp /path/to/saved_processor_dir/* /path/to/model.safetensors +export ENABLE_CUDA_GRAPH=True +export ENABLE_EXPERIMENTAL_INFERENCE_ENGINE=True + +CKPT_PATH=/path/to/checkpoint + +python -m wall_x._vendor.harrix.serving.launch_serving \ + --env X2ROBOT \ + --port 32195 \ + model-config:server-model-config \ + --model-config.checkpoint-path "$CKPT_PATH" \ + --model-config.train-config-path /path/to/my_robot_config.yml \ + --model-config.action-horizon 32 \ + --model-config.robot-action-interpolate-multiplier 1 \ + --model-config.robot-action-end-ratio 1.0 \ + --model-config.robot-type desktop +``` + +Clients connect at `ws://127.0.0.1:32195` to send observations and receive predicted actions. + +### 3.6 Plot Open-Loop Results to Verify the Server + +Open-loop evaluation compares model-predicted action trajectories against dataset ground truth **without executing actions or feeding back state**. + +**Terminal 1:** Keep the inference server from §3.5 running. + +**Terminal 2:** Run the open-loop plotting script: - # In earlier versions of PyTorch, errors may occur. You can use our provided script to address this issue; refer to wall-x/scripts/merge_sharded_weights.py for details. +```bash +python scripts/draw_openloop_plot.py \ + --uri ws://127.0.0.1:32195 \ + --dataset-root /path/to/your_robot_dataset \ + --train-config /path/to/my_robot_config.yml \ + --episode-indices 0,1,2 \ + --save-dir ./openloop_plots ``` -## Memory Usage +The script loads episodes from the LeRobot dataset, queries the WebSocket server frame by frame, and saves **predicted vs. ground-truth** comparison plots (PNG) under `--save-dir`. + +Common options: + +| Option | Description | +|--------|-------------| +| `--episode-indices` | Comma-separated episode indices to evaluate | +| `--start-ratio` | Start position as a fraction of episode length (0.0 = from the beginning) | +| `--stride` | Frames between inference calls (default: `action_horizon`) | +| `--max-inferences` | Maximum inference requests per episode | + +--- + +## Appendix + +### Example Config Files + +| File | Scenario | +|------|----------| +| `workspace/example/fintune_bus2602/libero.yml` | LIBERO single-arm fine-tuning | +| `workspace/example/fintune_bus2602/maniparena_example.yml` | Real-robot dual-arm fine-tuning | + +### Helper Scripts + +| Script | Purpose | +|--------|---------| +| `scripts/compute_norm_stats.py` | Compute dataset normalization statistics | +| `scripts/fake_inference.py` | Inference smoke test | +| `scripts/run_libero.sh` | Batch LIBERO simulation evaluation | +| `scripts/run_serving.sh` | Launch WebSocket inference server | +| `scripts/draw_openloop_plot.py` | Open-loop evaluation and plotting | +| `scripts/merge_sharded_weights.py` | Merge FSDP sharded checkpoints | -Below are the memory consumption benchmarks for different training configurations using the `lerobot/aloha_mobile_cabinet` dataset: +See [`scripts/README.md`](../scripts/README.md) for more script details. -| Dataset | Batch Size | FSDP2 | Torch Compile | Num GPUs | Max Allocated Memory | -|---------|------------|--------|---------------|----------|---------------------| -| lerobot/aloha_mobile_cabinet | 1 | ❌ | ❌ | 1 | 40.11G | -| lerobot/aloha_mobile_cabinet | 1 | ❌ | ❌ | 8 | 48.02G | -| lerobot/aloha_mobile_cabinet | 1 | ✅ | ❌ | 2 | 43.70G | -| lerobot/aloha_mobile_cabinet | 1 | ✅ | ❌ | 8 | 24.96G | -| lerobot/aloha_mobile_cabinet | 1 | ✅ | ✅ | 8 | 24.21G | +### Training Hyperparameters +| Parameter | Description | Default (libero example) | +|-----------|-------------|--------------------------| +| `hyperparams.batch_size_per_gpu` | Batch size per GPU | 4 | +| `hyperparams.gradient_accumulation_steps` | Gradient accumulation steps | 4 | +| `hyperparams.optimizer.learning_rate` | Learning rate | 5e-5 | +| `hyperparams.num_epoch` | Number of training epochs | 100 | +| `distributed.use_fsdp` | Multi-GPU FSDP training | true | +| `logging.save_interval` | Checkpoint save interval (steps) | 2000 | -**Hardware Recommendations:** +### Quick Start Checklist -- For single GPU training: Ensure at least 48GB VRAM (e.g., RTX 6000 Ada, A6000) -- For multi-GPU training: Enable FSDP2 for optimal memory distribution +**Simulation (LIBERO)** -## Reproduce +- [ ] Downloaded `wall-oss-0.5` and `Qwen2.5-VL-3B-Instruct` +- [ ] Downloaded the `lerobot/libero` dataset +- [ ] Wrote `config.yml` and replaced all `/path/to/*` placeholders +- [ ] Ran `compute_norm_stats.py` to generate norm stats +- [ ] Launched `train_fsdp.py` and completed fine-tuning +- [ ] Ran simulation evaluation with `run_libero.sh` -Openloop plot `wall-x/workspace/lerobot_example/evaluation/lerobot_openloop.png` +**Real Robot** -To reproduce the results, use the config file wall-x/workspace/lerobot_example/config_qact_from_vlm.yml with a global batch size of 128, adjusted via `gradient_accumulation_steps` and numbers of gpu. +- [ ] Prepared a LeRobot v3 real-robot dataset +- [ ] Wrote `config.yml` and configured `key_mappings` / DOF +- [ ] Ran `compute_norm_stats.py` to generate norm stats +- [ ] Launched `train_fsdp.py` and completed fine-tuning +- [ ] Started the inference server with `run_serving.sh` +- [ ] Verified server output with `draw_openloop_plot.py` diff --git a/workspace/example/libero.yml b/workspace/example/libero.yml new file mode 100644 index 00000000..e7629a8a --- /dev/null +++ b/workspace/example/libero.yml @@ -0,0 +1,125 @@ +# LIBERO single-arm finetune example (Euler delta action, 256px, 2 cameras). +# +# Replace every /path/to/* placeholder before training: +# +# model.config_path -> model architecture JSON (mot_flash_mask_causal_xloss.json) +# model.processor_path -> Qwen2.5-VL-3B-Instruct directory +# model.pretrained_path -> same as processor_path, or HuggingFace cache path +# data.lerobot_config.repo_id -> local LeRobot dataset root (libero_all) +# data.norm_stats_path -> q01/q99 normalization JSON +# checkpoint.save_path -> writable directory for training checkpoints +# checkpoint.resume_from -> bus2602 pretrain .safetensors or checkpoint directory +# +# Compute norm stats first: +# +# python scripts/compute_norm_stats.py \ +# --train_config workspace/example/fintune_bus2602/libero.yml \ +# --data_root /path/to/libero_all \ +# --output_path /path/to/libero_all_norm_stats.json +# +# Launch training (from repo root): +# +# torchrun --nproc_per_node= wall_x/trainer/fsdp_trainer/train_fsdp.py \ +# --config workspace/example/fintune_bus2602/libero.yml +# +# Strategy: keep dof / agent_pos totals at 26 (matching the bus2602 pretrain ckpt) +# via ``action_padding``. The lerobot collator right-pads libero's 7-dim action / +# 8-dim state with zeros; loss does not flow through the padded tail. + +model_type: qwen2_5 + +task: + # 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 + gripper1 = 7, plus action_padding(19) = 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 + use_state_string_representation: true + +model: + backbone: qwen2_5 + 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 + attn_deterministic: true + use_ema: false + flow_loss_weight: 1.0 + ar_loss_weight: 0.01 + +hyperparams: + num_epoch: 100 + batch_size_per_gpu: 4 + gradient_accumulation_steps: 4 + seed: 10222 + optimizer: + optimizer_type: adamw + learning_rate: 5.0e-05 + max_grad_norm: 1.0 + enable_grad_clip: true + betas: [0.9, 0.95] + weight_decay: 1.0e-8 + eps: 1.0e-8 + scheduler: + scheduler_type: cosine + num_warmup_steps: 1000 + num_training_steps: 200000 + min_lr: 1.0e-6 + +distributed: + use_fsdp: true + use_mixed_precision: true + bf16: true + +data: + dataset_type: lerobot + lerobot_config: + repo_id: /path/to/libero_all + root: null + key_mappings: + # libero_all v3.0 only has faceImg + rightImg (no leftImg). + camera: + observation.images.faceImg: face_view + observation.images.rightImg: right_wrist_view + state: observation.state + action: action + norm_stats_path: /path/to/libero_all_norm_stats.json + train_test_split: 0.95 + num_workers: 4 + max_length: 1024 + resolution: + face_view: 256 + right_wrist_view: 256 + +logging: + log_name: libero_ft + log_project: lerobot_libero_ft + log_entity: your_wandb_entity + use_wandb: true + log_interval: 10 + save_interval: 2000 + val_interval: 1000000 + epoch_save_interval: 1 + +checkpoint: + save_path: /path/to/libero + # Single-file .safetensors loads as pretrain weights before FSDP wrapping. + # Use a checkpoint directory for full resume (optimizer / scheduler / RNG). + resume_from: /path/to/wall-oss-0.5/model.safetensors + +debug: + profile: false + nvtx: false diff --git a/workspace/example/maniparena_example.yml b/workspace/example/maniparena_example.yml new file mode 100644 index 00000000..4034d748 --- /dev/null +++ b/workspace/example/maniparena_example.yml @@ -0,0 +1,134 @@ +# ManipArena / CVPR dual-arm finetune example (6D relative action, 448px, 3 cameras). +# +# Replace every /path/to/* placeholder before training: +# +# model.config_path -> model architecture JSON (mot_flash_mask_causal_xloss.json) +# model.processor_path -> Qwen2.5-VL-3B-Instruct directory +# model.pretrained_path -> same as processor_path, or HuggingFace cache path +# data.lerobot_config.repo_id -> local LeRobot dataset root (cvpr_4tasks) +# data.norm_stats_path -> q01/q99 normalization JSON +# checkpoint.save_path -> writable directory for training checkpoints +# checkpoint.resume_from -> bus2602 pretrain .safetensors or checkpoint directory +# +# Compute norm stats first: +# +# python scripts/compute_norm_stats.py \ +# --train_config workspace/example/fintune_bus2602/maniparena_example.yml \ +# --data_root /path/to/cvpr_4tasks \ +# --output_path /path/to/cvpr_4tasks_norm_stats.json +# +# Launch training (from repo root): +# +# torchrun --nproc_per_node= wall_x/trainer/fsdp_trainer/train_fsdp.py \ +# --config workspace/example/fintune_bus2602/maniparena_example.yml +# +# Strategy: keep dof / agent_pos totals at 26 (matching the bus2602 pretrain ckpt) +# via ``action_padding``. The lerobot collator right-pads real action / state dims; +# loss does not flow through the padded tail. + +model_type: qwen2_5 + +task: + # Dual-arm 6D relative action: 10 + 10 = 20, plus action_padding(6) = 26. + dof_config: + follow_left_ee_cartesian_pos_relative: 3 + follow_left_ee_rotation_6D_relative: 6 + follow_left_gripper: 1 + follow_right_ee_cartesian_pos_relative: 3 + follow_right_ee_rotation_6D_relative: 6 + follow_right_gripper: 1 + action_padding: 6 + ar_dof_config: + follow_left_ee_cartesian_pos_relative: 3 + follow_left_ee_rotation_6D_relative: 6 + follow_left_gripper: 1 + follow_right_ee_cartesian_pos_relative: 3 + follow_right_ee_rotation_6D_relative: 6 + follow_right_gripper: 1 + action_padding: 6 + agent_pos_config: + follow_left_ee_cartesian_pos: 3 + follow_left_ee_rotation_6D: 6 + follow_left_gripper: 1 + follow_right_ee_cartesian_pos: 3 + follow_right_ee_rotation_6D: 6 + follow_right_gripper: 1 + action_padding: 6 + action_horizon: 32 + action_horizon_flow: 32 + use_state_string_representation: true + +model: + backbone: qwen2_5 + 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 + attn_deterministic: true + use_ema: false + flow_loss_weight: 1.0 + ar_loss_weight: 0.01 + +hyperparams: + num_epoch: 100 + batch_size_per_gpu: 4 + gradient_accumulation_steps: 4 + seed: 10222 + optimizer: + optimizer_type: adamw + learning_rate: 5.0e-05 + max_grad_norm: 1.0 + enable_grad_clip: true + betas: [0.9, 0.95] + weight_decay: 1.0e-8 + eps: 1.0e-8 + scheduler: + scheduler_type: cosine + num_warmup_steps: 1000 + num_training_steps: 200000 + min_lr: 1.0e-6 + +distributed: + use_fsdp: true + use_mixed_precision: true + bf16: true + +data: + dataset_type: lerobot + lerobot_config: + repo_id: /path/to/cvpr_4tasks + root: null + key_mappings: + camera: + observation.images.faceImg: face_view + observation.images.leftImg: left_wrist_view + observation.images.rightImg: right_wrist_view + state: observation.state + action: action + norm_stats_path: /path/to/cvpr_4tasks_norm_stats.json + train_test_split: 0.95 + num_workers: 4 + max_length: 1024 + resolution: + face_view: 448 + left_wrist_view: 448 + right_wrist_view: 448 + +logging: + log_name: maniparena_ft + log_project: lerobot_maniparena_ft + log_entity: your_wandb_entity + use_wandb: true + log_interval: 10 + save_interval: 2000 + val_interval: 1000000 + epoch_save_interval: 1 + +checkpoint: + save_path: /path/to/cvpr_4tasks + # Single-file .safetensors loads as pretrain weights before FSDP wrapping. + # Use a checkpoint directory for full resume (optimizer / scheduler / RNG). + resume_from: /path/to/wall-oss-0.5/model.safetensors + +debug: + profile: false + nvtx: false diff --git a/workspace/lerobot_example/config_qact.yml b/workspace/lerobot_example/config_qact.yml deleted file mode 100644 index 1248fa27..00000000 --- a/workspace/lerobot_example/config_qact.yml +++ /dev/null @@ -1,151 +0,0 @@ -# Training Configuration for Wall-X Robotic Multi-Modal Learning -# This configuration supports multi-modal learning with vision, language, and action data - -# Model and paths configuration -log_name: "robotic_training" -log_project: "vla_training" -model_type: wall-oss -pretrained_wallx_path: "/path/to/wallx_model/" # Must set -save_path: "/path/to/workspace/" # Must set -use_fast_tokenizer: False # True: train FAST, False: train Flow -action_tokenizer_path: "/path/to/fast/" # Must set if use_fast_tokenizer is true - -# Torch Profile -profile: False -profile_save_path: /path/to/profile/ -profile_wait_iters: 10 -profile_warmup_iters: 5 -profile_active_iters: 2 - -# Training hyperparameters -num_warmup_steps: 100 -num_training_steps: 64000000 -learning_rate: 0.00005 -min_lr: 0.00005 -num_epoch: 100 -gradient_accumulation_steps: 32 -batch_size_per_gpu: 8 -padding_side: left -epoch_save_interval: 10 - -# Training optimization settings -FSDP2: True -torch_compile: False - -# Robot configuration - Define degrees of freedom for each component -dof_config: - follow_left_ee_cartesian_pos: 3 # Left end-effector Cartesian position - follow_left_ee_rotation: 3 # Left end-effector rotation - follow_left_gripper: 1 # Left gripper control - follow_right_ee_cartesian_pos: 3 # Right end-effector Cartesian position - follow_right_ee_rotation: 3 # Right end-effector rotation - follow_right_gripper: 1 # Right gripper control - head_actions: 2 # Head/camera movement - height: 1 # Mobile base height control - car_pose: 3 # Mobile base pose (x, y, theta) - -# Agent proprioception configuration (typically matches DOF config) -agent_pos_config: - follow_left_ee_cartesian_pos: 3 - follow_left_ee_rotation: 3 - follow_left_gripper: 1 - follow_right_ee_cartesian_pos: 3 - follow_right_ee_rotation: 3 - follow_right_gripper: 1 - head_actions: 2 - height: 1 - car_pose: 3 - -# # Checkpoint resuming configuration -# resume: -# ckpt: "/path/to/resume_model/" -# load_ckpt_only: true - -norm_stats_path: "/path/to/norm_stats.json" - -enable_customized_robot_config: true -customized_robot_config: - name: "lerobot/aloha_mobile_cabinet" - customized_dof_config: - "action_left_shoulder" : 1 - "action_left_elbow" : 1 - "action_left_forearm_roll" : 1 - "action_left_wrist_angle" : 1 - "action_left_wrist_rotate" : 1 - "action_left_gripper" : 1 - "action_right_waist" : 1 - "action_right_shoulder" : 1 - "action_right_elbow" : 1 - "action_right_forearm_roll" : 1 - "action_right_wrist_angle" : 1 - "action_right_wrist_rotate" : 1 - "action_right_gripper" : 1 - - customized_agent_pos_config: - "state_left_shoulder" : 1 - "state_left_elbow" : 1 - "state_left_forearm_roll" : 1 - "state_left_wrist_angle" : 1 - "state_left_wrist_rotate" : 1 - "state_left_gripper" : 1 - "state_right_waist" : 1 - "state_right_shoulder" : 1 - "state_right_elbow" : 1 - "state_right_forearm_roll" : 1 - "state_right_wrist_angle" : 1 - "state_right_wrist_rotate" : 1 - "state_right_gripper" : 1 - -# Data configuration -data: - use_lerobot: true - - # LeRobot dataset configuration - lerobot_config: - repo_id: "lerobot/aloha_mobile_cabinet" - root: null - episodes: null - image_transforms: null - delta_timestamps: null - tolerance_s: 1e-4 - revision: null - force_cache_sync: false - download_videos: true - video_backend: null - - action_horizon: 32 - train_test_split: 0.95 - - # Action keys for observation and prediction - obs_action_keys: - - follow_left_ee_cartesian_pos - - follow_left_ee_rotation - - follow_left_gripper - - follow_right_ee_cartesian_pos - - follow_right_ee_rotation - - follow_right_gripper - - head_actions - - height - - car_pose - - predict_action_keys: - - follow_left_ee_cartesian_pos - - follow_left_ee_rotation - - follow_left_gripper - - follow_right_ee_cartesian_pos - - follow_right_ee_rotation - - follow_right_gripper - - head_actions - - height - - car_pose - - # Image resolution configuration for different camera views - resolution: - face_view: 256 - left_wrist_view: 256 - right_wrist_view: 256 - move1_view: 256 - move2_view: 256 - top_view: 256 - wall_view: 256 - multi_modal: 256 diff --git a/workspace/lerobot_example/config_qact_from_vlm.yml b/workspace/lerobot_example/config_qact_from_vlm.yml deleted file mode 100644 index e8723110..00000000 --- a/workspace/lerobot_example/config_qact_from_vlm.yml +++ /dev/null @@ -1,152 +0,0 @@ -# Train from Qwen-2.5-VL - -# Model and paths configuration -log_name: "robotic_training" -log_project: "vla_training" -model_type: qwen2_5 -pretrained_wallx_path: "/path/to/wallx_model/" # Must set -save_path: "/path/to/workspace/" # Must set -use_fast_tokenizer: True # True: train FAST, False: train Flow -action_tokenizer_path: "/path/to/fast/" # Must set if use_fast_tokenizer is true -qwen_vl_act_config_path: "wall-x/workspace/lerobot_example/qwen25_config.json" - - -# Torch Profile -profile: False -profile_save_path: /path/to/profile/ -profile_wait_iters: 10 -profile_warmup_iters: 5 -profile_active_iters: 2 - -# Training hyperparameters -num_warmup_steps: 100 -num_training_steps: 64000000 -learning_rate: 0.00009 -min_lr: 0.00005 -num_epoch: 100 -gradient_accumulation_steps: 1 -batch_size_per_gpu: 8 -padding_side: left -epoch_save_interval: 10 - -# Training optimization settings -FSDP2: True -torch_compile: False - -# Robot configuration - Define degrees of freedom for each component -dof_config: - follow_left_ee_cartesian_pos: 3 # Left end-effector Cartesian position - follow_left_ee_rotation: 3 # Left end-effector rotation - follow_left_gripper: 1 # Left gripper control - follow_right_ee_cartesian_pos: 3 # Right end-effector Cartesian position - follow_right_ee_rotation: 3 # Right end-effector rotation - follow_right_gripper: 1 # Right gripper control - head_actions: 2 # Head/camera movement - height: 1 # Mobile base height control - car_pose: 3 # Mobile base pose (x, y, theta) - -# Agent proprioception configuration (typically matches DOF config) -agent_pos_config: - follow_left_ee_cartesian_pos: 3 - follow_left_ee_rotation: 3 - follow_left_gripper: 1 - follow_right_ee_cartesian_pos: 3 - follow_right_ee_rotation: 3 - follow_right_gripper: 1 - head_actions: 2 - height: 1 - car_pose: 3 - -# # Checkpoint resuming configuration -# resume: -# ckpt: "/path/to/resume_model/" -# load_ckpt_only: true - -norm_stats_path: "/path/to/norm_stats.json" - -enable_customized_robot_config: true -customized_robot_config: - name: "physical-intelligence/libero" - customized_dof_config: - "action_left_shoulder" : 1 - "action_left_elbow" : 1 - "action_left_forearm_roll" : 1 - "action_left_wrist_angle" : 1 - "action_left_wrist_rotate" : 1 - "action_left_gripper" : 1 - "action_right_waist" : 1 - "action_right_shoulder" : 1 - "action_right_elbow" : 1 - "action_right_forearm_roll" : 1 - "action_right_wrist_angle" : 1 - "action_right_wrist_rotate" : 1 - "action_right_gripper" : 1 - - customized_agent_pos_config: - "state_left_shoulder" : 1 - "state_left_elbow" : 1 - "state_left_forearm_roll" : 1 - "state_left_wrist_angle" : 1 - "state_left_wrist_rotate" : 1 - "state_left_gripper" : 1 - "state_right_waist" : 1 - "state_right_shoulder" : 1 - "state_right_elbow" : 1 - "state_right_forearm_roll" : 1 - "state_right_wrist_angle" : 1 - "state_right_wrist_rotate" : 1 - "state_right_gripper" : 1 - -# Data configuration -data: - use_lerobot: true - - # LeRobot dataset configuration - lerobot_config: - repo_id: "lerobot/aloha_mobile_cabinet" - root: null - episodes: null - image_transforms: null - delta_timestamps: null - tolerance_s: 1e-4 - revision: null - force_cache_sync: false - download_videos: true - video_backend: null - - action_horizon: 32 - train_test_split: 0.95 - - # Action keys for observation and prediction - obs_action_keys: - - follow_left_ee_cartesian_pos - - follow_left_ee_rotation - - follow_left_gripper - - follow_right_ee_cartesian_pos - - follow_right_ee_rotation - - follow_right_gripper - - head_actions - - height - - car_pose - - predict_action_keys: - - follow_left_ee_cartesian_pos - - follow_left_ee_rotation - - follow_left_gripper - - follow_right_ee_cartesian_pos - - follow_right_ee_rotation - - follow_right_gripper - - head_actions - - height - - car_pose - - # Image resolution configuration for different camera views - resolution: - face_view: 256 - left_wrist_view: 256 - right_wrist_view: 256 - move1_view: 256 - move2_view: 256 - top_view: 256 - wall_view: 256 - multi_modal: 256 diff --git a/workspace/lerobot_example/evaluation/lerobot_openloop.png b/workspace/lerobot_example/evaluation/lerobot_openloop.png deleted file mode 100644 index dfc72e74..00000000 --- a/workspace/lerobot_example/evaluation/lerobot_openloop.png +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:aae8646e566b79b64a669956d6d2c779d72809010c802f75b2623ee371444b47 -size 985979 diff --git a/workspace/lerobot_example/libero/config_qact_libero_from_vlm.yml b/workspace/lerobot_example/libero/config_qact_libero_from_vlm.yml deleted file mode 100644 index 21198e48..00000000 --- a/workspace/lerobot_example/libero/config_qact_libero_from_vlm.yml +++ /dev/null @@ -1,100 +0,0 @@ -# Training Configuration for Wall-X Robotic Multi-Modal Learning -# This configuration supports multi-modal learning with vision, language, and action data - -# Model and paths configuration -log_name: "opensource_training" -log_project: "libero" -model_type: qwen2_5 -use_fast_tokenizer: true -pretrained_wallx_path: "/path/to/qwen/" -action_tokenizer_path: "/path/to/fast/" -qwen_vl_act_config_path: "/path/to/qwen25_config.json" - -save_path: "/path/to/save" -# Torch Profile -profile: False -profile_save_path: /path/to/profile/ -profile_wait_iters: 10 -profile_warmup_iters: 5 -profile_active_iters: 2 - -# Training hyperparameters -num_warmup_steps: 100 -num_training_steps: 64000000 -learning_rate: 0.00005 -min_lr: 0.00005 -num_epoch: 100 -gradient_accumulation_steps: 1 -batch_size_per_gpu: 8 -padding_side: left -epoch_save_interval: 1 - -# Robot configuration - Define degrees of freedom for each component -dof_config: - master_right_ee_cartesian_pos: 3 # Right end-effector Cartesian position - master_right_ee_rotation: 3 # Right end-effector rotation - master_right_gripper: 1 # Right gripper control - -# Agent proprioception configuration (typically matches DOF config) -agent_pos_config: - follow_right_ee_cartesian_pos: 3 - follow_right_ee_rotation: 3 - follow_right_gripper: 1 - -norm_stats_path: "/path/to/libero_norm_stats.json" - -enable_customized_robot_config: true -customized_robot_config: - name: "libero_all" - customized_dof_config: - "panda_action_eef_with_gripper": 7 - - customized_agent_pos_config: - "panda_state_eef_with_gripper": 7 - -# Checkpoint resuming configuration -# resume: -# ckpt: "/path/to/ckpt" -# load_ckpt_only: false - -# Data configuration -data: - use_lerobot: true - - # LeRobot dataset configuration - lerobot_config: - repo_id: "libero_all" - root: null - episodes: null - image_transforms: null - delta_timestamps: null - tolerance_s: 1e-4 - revision: null - force_cache_sync: false - download_videos: true - video_backend: null - - action_horizon: 10 - train_test_split: 0.95 - - # Action keys for observation and prediction - obs_action_keys: - - follow_right_ee_cartesian_pos - - follow_right_ee_rotation - - follow_right_gripper - - predict_action_keys: - - master_right_ee_cartesian_pos - - master_right_ee_rotation - - master_right_gripper - - # Image resolution configuration for different camera views - resolution: - face_view: 256 - left_wrist_view: 256 - right_wrist_view: 256 - move1_view: 256 - move2_view: 256 - top_view: 256 - wall_view: 256 - multi_modal: 256 diff --git a/workspace/lerobot_example/qwen25_config.json b/workspace/lerobot_example/qwen25_config.json deleted file mode 100644 index c92cc207..00000000 --- a/workspace/lerobot_example/qwen25_config.json +++ /dev/null @@ -1,106 +0,0 @@ -{ - "architectures": [ - "Qwen2_5_VLForConditionalGeneration" - ], - "attention_dropout": 0.0, - "bos_token_id": 151643, - "eos_token_id": 151645, - "vision_start_token_id": 151652, - "vision_end_token_id": 151653, - "vision_token_id": 151654, - "image_token_id": 151655, - "video_token_id": 151656, - "hidden_act": "silu", - "hidden_size": 2048, - "initializer_range": 0.02, - "intermediate_size": 11008, - "max_position_embeddings": 128000, - "max_window_layers": 70, - "model_type": "qwen2_5_vl", - "num_attention_heads": 16, - "num_hidden_layers": 36, - "num_key_value_heads": 2, - "rms_norm_eps": 1e-06, - "rope_theta": 1000000.0, - "sliding_window": 32768, - "tie_word_embeddings": true, - "torch_dtype": "bfloat16", - "transformers_version": "4.41.2", - "_attn_implementation": "flash_attention_2", - "use_cache": true, - "use_sliding_window": false, - "vision_config": { - "depth": 32, - "hidden_act": "silu", - "hidden_size": 1280, - "intermediate_size": 3420, - "num_heads": 16, - "in_chans": 3, - "out_hidden_size": 2048, - "patch_size": 14, - "spatial_merge_size": 2, - "spatial_patch_size": 14, - "window_size": 112, - "fullatt_block_indexes": [ - 7, - 15, - 23, - 31 - ], - "tokens_per_second": 2, - "temporal_patch_size": 2 - }, - "rope_scaling": { - "type": "mrope", - "mrope_section": [ - 16, - 24, - 24 - ] - }, - "vocab_size": 151936, - "num_experts": 2, - "experts":[ - { - "hidden_size": 2048, - "intermediate_size": 11008, - "hidden_act": "silu" - }, - { - "hidden_size": 2048, - "intermediate_size": 2048, - "hidden_act": "silu" - } - ], - "dof_config": { - "follow_left_ee_cartesian_pos": 3, - "follow_left_ee_rotation": 3, - "follow_left_gripper": 1, - "follow_right_ee_cartesian_pos": 3, - "follow_right_ee_rotation": 3, - "follow_right_gripper": 1, - "head_actions": 2, - "height": 1, - "car_pose": 3 - }, - "agent_pos_config": { - "follow_left_ee_cartesian_pos": 3, - "follow_left_ee_rotation": 3, - "follow_left_gripper": 1, - "follow_right_ee_cartesian_pos": 3, - "follow_right_ee_rotation": 3, - "follow_right_gripper": 1, - "head_actions": 2, - "height": 1, - "car_pose": 3 - }, - "noise_scheduler": { - "beta_alpha": 1.5, - "beta_beta": 1.0, - "s": 0.999, - "num_inference_timesteps": 5 - }, - "dim_inputs": [2048,2048], - "attention_moe": false, - "mlp_moe": true - } diff --git a/workspace/lerobot_example/run.sh b/workspace/lerobot_example/run.sh deleted file mode 100644 index 7d181068..00000000 --- a/workspace/lerobot_example/run.sh +++ /dev/null @@ -1,23 +0,0 @@ -#!/bin/bash -export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 -NUM_GPUS=$(echo $CUDA_VISIBLE_DEVICES | tr ',' '\n' | wc -l) - -# print current time -echo "[current time: $(date +'%Y-%m-%d %H:%M:%S')]" - -code_dir="/path/to/wall-x" -config_path="/path/to/wall-x/workspace/lerobot_example" - -# Use a fixed port instead of a random one -export PORT=$((21000 + $RANDOM % 30000)) - -MASTER_PORT=10239 # use 5 digits ports - -export LAUNCHER="accelerate launch --num_processes=$NUM_GPUS --main_process_port=$PORT" - -export SCRIPT="${code_dir}/train_qact.py" -export SCRIPT_ARGS="--config ${config_path}/config_qact.yml --seed $MASTER_PORT" - -echo "Running command: $LAUNCHER $SCRIPT $SCRIPT_ARGS" - -$LAUNCHER $SCRIPT $SCRIPT_ARGS From aadc0cc996539ec8c41440e6c3846cc19a844167 Mon Sep 17 00:00:00 2001 From: yangping Date: Wed, 10 Jun 2026 11:07:46 +0800 Subject: [PATCH 2/3] fix: add trailing newline to run_serving.sh for pre-commit Co-authored-by: Cursor --- scripts/run_serving.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/run_serving.sh b/scripts/run_serving.sh index d05f2695..4ff0f08b 100755 --- a/scripts/run_serving.sh +++ b/scripts/run_serving.sh @@ -13,4 +13,4 @@ python -m wall_x._vendor.harrix.serving.launch_serving \ --model-config.action-horizon 32 \ --model-config.robot-action-interpolate-multiplier 1 \ --model-config.robot-action-end-ratio 1.0 \ - --model-config.robot-type desktop \ No newline at end of file + --model-config.robot-type desktop From 84968afbceca3b8aaa3a7f78aad9d0136438c0ec Mon Sep 17 00:00:00 2001 From: yangping Date: Wed, 10 Jun 2026 11:13:48 +0800 Subject: [PATCH 3/3] style: apply ruff and black formatting for pre-commit CI Co-authored-by: Cursor --- scripts/draw_openloop_plot.py | 33 ++++++---- scripts/infer_libero.py | 5 +- scripts/vqa_inference.py | 4 +- wall_x/_vendor/harrix/adapters/qwen_vlact.py | 10 +-- .../harrix/serving-old/_wallx_infer/robot.py | 9 ++- .../_wallx_infer/socket_controller.py | 8 ++- .../harrix/serving-old/_wallx_infer/utils.py | 4 +- .../serving/_wallx_infer/infer_config.py | 14 ++-- .../serving/_wallx_infer/model_wrapper.py | 3 - .../harrix/serving/_wallx_infer/robot.py | 21 +++--- .../serving/_wallx_infer/socket_controller.py | 8 ++- .../harrix/serving/_wallx_infer/utils.py | 4 +- wall_x/_vendor/harrix/serving/client.py | 4 +- .../_vendor/harrix/serving/launch_serving.py | 10 +-- .../harrix/serving/policy/wall_x_policy.py | 8 +-- wall_x/_vendor/harrix/serving/scheduler.py | 66 ++++++++++--------- .../harrix/serving/websocket_policy_server.py | 8 ++- wall_x/_vendor/harrix/utils/normalizer.py | 20 ++++-- wall_x/_vendor/x2robot_utils/grounding.py | 4 +- .../_vendor/x2robot_utils/text_templates.py | 51 ++++++++++---- wall_x/data/backends/lerobot/loader.py | 24 ++++--- wall_x/model/qact/tokenizer_mixin.py | 2 - wall_x/trainer/adapters/base_adapter.py | 1 - wall_x/trainer/fsdp_trainer/checkpoint_io.py | 10 ++- 24 files changed, 201 insertions(+), 130 deletions(-) diff --git a/scripts/draw_openloop_plot.py b/scripts/draw_openloop_plot.py index cfbf2816..4ea0fada 100644 --- a/scripts/draw_openloop_plot.py +++ b/scripts/draw_openloop_plot.py @@ -26,7 +26,6 @@ import base64 import json import logging -import os import sys from dataclasses import dataclass, field from pathlib import Path @@ -74,7 +73,9 @@ def get_frame_images(self, frame_idx: int) -> dict[str, np.ndarray]: if self._dataset is None: raise RuntimeError("Episode image loader is not initialized.") if not self._image_cache: - logger.info("Decoding camera frames on demand (first obs_idx=%d)", frame_idx) + logger.info( + "Decoding camera frames on demand (first obs_idx=%d)", frame_idx + ) item = self._dataset[frame_idx] images = { cam_key: tensor_to_rgb_uint8(item[cam_key]) @@ -109,7 +110,11 @@ def _task_block(train_config: dict[str, Any]) -> dict[str, Any]: def _dof_config(train_config: dict[str, Any]) -> dict[str, int]: - return train_config.get("dof_config") or _task_block(train_config).get("dof_config") or {} + return ( + train_config.get("dof_config") + or _task_block(train_config).get("dof_config") + or {} + ) def _agent_pos_config(train_config: dict[str, Any]) -> dict[str, int]: @@ -268,9 +273,7 @@ def build_state_payload( has_right = any(k.startswith("follow_right_") for k in agent_cfg) payload: dict[str, list[float]] = {} if has_left: - left = _arm_follow_pos_from_agent_cfg( - state_vec, agent_cfg, "follow_left_" - ) + left = _arm_follow_pos_from_agent_cfg(state_vec, agent_cfg, "follow_left_") if left is not None: payload["follow1_pos"] = left if has_right: @@ -549,9 +552,7 @@ def load_episode_arrays( start = max(0, int(frame_start)) end = num_steps if frame_end is None else min(num_steps, int(frame_end)) - states = [ - _hf_row_to_numpy(ds.hf_dataset[i], state_key) for i in range(start, end) - ] + states = [_hf_row_to_numpy(ds.hf_dataset[i], state_key) for i in range(start, end)] actions = [ _hf_row_to_numpy(ds.hf_dataset[i], action_key) for i in range(start, end) ] @@ -678,9 +679,9 @@ def plot_openloop( gt_rows = np.asarray(action_gt_list, dtype=np.float32) pred_rows = np.asarray(action_pred_list, dtype=np.float32) - assert gt_rows.shape == pred_rows.shape, ( - f"shape mismatch: pred={pred_rows.shape} gt={gt_rows.shape}" - ) + assert ( + gt_rows.shape == pred_rows.shape + ), f"shape mismatch: pred={pred_rows.shape} gt={gt_rows.shape}" n, dim = gt_rows.shape fig, axes = plt.subplots(dim, 1, figsize=(12, 3.5 * dim), sharex=True) @@ -887,7 +888,9 @@ async def run_openloop_eval( ep_tag = f"ep{ep_idx}" valid_mask = ~np.isnan(pred_full).any(axis=1) - action_l1 = float(np.mean(np.abs(pred_full[valid_mask] - gt_full[valid_mask]))) + action_l1 = float( + np.mean(np.abs(pred_full[valid_mask] - gt_full[valid_mask])) + ) summary = { "episode_index": ep_idx, "instruction": episode.instruction, @@ -908,7 +911,9 @@ async def run_openloop_eval( ml = min(len(aligned_gt), len(aligned_pred)) if ml < 3: - logger.warning("Episode %d: too few aligned rows (%d) to plot.", ep_idx, ml) + logger.warning( + "Episode %d: too few aligned rows (%d) to plot.", ep_idx, ml + ) continue gt_ml = np.asarray(aligned_gt[:ml], dtype=np.float32) diff --git a/scripts/infer_libero.py b/scripts/infer_libero.py index d96a3b61..eb6ad5ec 100644 --- a/scripts/infer_libero.py +++ b/scripts/infer_libero.py @@ -212,7 +212,10 @@ def main() -> int: args = parse_args() _ensure_local_harrix_on_path() - from wall_x._vendor.harrix.eval_config import autofill_from_checkpoint, load_eval_config + from wall_x._vendor.harrix.eval_config import ( + autofill_from_checkpoint, + load_eval_config, + ) raw = _load_or_build_raw_config(args) diff --git a/scripts/vqa_inference.py b/scripts/vqa_inference.py index b8e10a2e..3b901e0f 100644 --- a/scripts/vqa_inference.py +++ b/scripts/vqa_inference.py @@ -89,9 +89,7 @@ def generate(self, image: Image.Image, text: str, **kwargs) -> str: test_question = "To move the red block in the plate with same color, what should you do next? Think step by step." # Local Image - img = Image.open( - "/path/to/cot_example_frame.png" - ).convert("RGB") + img = Image.open("/path/to/cot_example_frame.png").convert("RGB") # Internet Image # import requests # test_image_url = "https://www.ilankelman.org/stopsigns/australia.jpg" diff --git a/wall_x/_vendor/harrix/adapters/qwen_vlact.py b/wall_x/_vendor/harrix/adapters/qwen_vlact.py index 400b71c1..2c7f42ac 100644 --- a/wall_x/_vendor/harrix/adapters/qwen_vlact.py +++ b/wall_x/_vendor/harrix/adapters/qwen_vlact.py @@ -87,12 +87,11 @@ def _normalize_real_prefix(normalizer, tensor, dataset_names): f"normalizer width {width} exceeds tensor dim {tensor.shape[-1]}" ) out = tensor.clone() - out[..., :width] = normalizer.normalize_data( - tensor[..., :width], dataset_names - ) + out[..., :width] = normalizer.normalize_data(tensor[..., :width], dataset_names) out[..., width:] = 0 return out + class QwenVLActInferAdapter(BaseInferAdapter): """Shared constructor and batched flow inference implementation.""" @@ -142,7 +141,10 @@ def __init__(self, cfg: EvalConfig) -> None: self._normalizer_action = normalizer_action self._normalizer_propri = normalizer_propri self._norm_key = resolved_norm_key - if normalizer_propri is not None and resolved_norm_key in normalizer_propri.delta: + if ( + normalizer_propri is not None + and resolved_norm_key in normalizer_propri.delta + ): train_config["_libero_proprio_norm_dim"] = int( normalizer_propri.delta[resolved_norm_key].shape[0] ) diff --git a/wall_x/_vendor/harrix/serving-old/_wallx_infer/robot.py b/wall_x/_vendor/harrix/serving-old/_wallx_infer/robot.py index 19dc8f15..f365ddc0 100644 --- a/wall_x/_vendor/harrix/serving-old/_wallx_infer/robot.py +++ b/wall_x/_vendor/harrix/serving-old/_wallx_infer/robot.py @@ -3,8 +3,13 @@ import re from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig -from wall_x._vendor.harrix.serving._wallx_infer.utils import VehiclePoseHandler, UnifiedTrajectoryProcessor -from wall_x._vendor.harrix.serving._wallx_infer.base_dataclass import RobotStateActionData +from wall_x._vendor.harrix.serving._wallx_infer.utils import ( + VehiclePoseHandler, + UnifiedTrajectoryProcessor, +) +from wall_x._vendor.harrix.serving._wallx_infer.base_dataclass import ( + RobotStateActionData, +) from wall_x._vendor.harrix.serving._wallx_infer.socket_controller import RobotController from wall_x._vendor.harrix.serving._wallx_infer.logger import InferLogger diff --git a/wall_x/_vendor/harrix/serving-old/_wallx_infer/socket_controller.py b/wall_x/_vendor/harrix/serving-old/_wallx_infer/socket_controller.py index b40eace8..6aabbbb7 100644 --- a/wall_x/_vendor/harrix/serving-old/_wallx_infer/socket_controller.py +++ b/wall_x/_vendor/harrix/serving-old/_wallx_infer/socket_controller.py @@ -25,7 +25,9 @@ class RobotRegistry: def __init__(self): self.registry = _DEFAULT_ROBOT_CONFIG self.logger = InferLogger.get_controller_logger("RobotRegistry") - self.logger.debug(f"Initialized robot registry with {len(self.registry)} robots") + self.logger.debug( + f"Initialized robot registry with {len(self.registry)} robots" + ) def register_robot(self, robot_id, host, action_port, keyboard_port): if robot_id in self.registry: @@ -90,7 +92,9 @@ def connect(self): port = self.robot_info["action_port"] self.action_sock.bind((host, port)) self.action_sock.listen(1) - self.logger.info(f"Listening on action port: {host}:{self.action_sock.getsockname()[1]}") + self.logger.info( + f"Listening on action port: {host}:{self.action_sock.getsockname()[1]}" + ) action_thread = threading.Thread(target=self.handle_action_client) diff --git a/wall_x/_vendor/harrix/serving-old/_wallx_infer/utils.py b/wall_x/_vendor/harrix/serving-old/_wallx_infer/utils.py index 9b49aebe..2d0c1242 100644 --- a/wall_x/_vendor/harrix/serving-old/_wallx_infer/utils.py +++ b/wall_x/_vendor/harrix/serving-old/_wallx_infer/utils.py @@ -41,7 +41,9 @@ def run(self): # Numeric input: switch instruction index index = int(user_input) self.new_instruction_index = index - self.logger.info(f"[Keyboard] Switched to instruction index: {index}") + self.logger.info( + f"[Keyboard] Switched to instruction index: {index}" + ) else: self.logger.info(f"[Keyboard] Input {user_input}; no action taken.") diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py b/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py index 7ed332cd..8b15b4a2 100644 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/infer_config.py @@ -141,9 +141,7 @@ def __init__( self._load_all_configs(train_config_path) self._apply_cam_names_from_train_config(cam_names) - def _apply_cam_names_from_train_config( - self, cam_names: list[str] | None - ) -> None: + def _apply_cam_names_from_train_config(self, cam_names: list[str] | None) -> None: """Use train YAML camera mapping when CLI did not override cam_names.""" if cam_names is not None: return @@ -217,9 +215,9 @@ def _load_all_configs(self, train_config_path=None): self._action_dim = sum(dof_config.values()) if self._ar_action_dim is None: - ar_dof_config = self.train_config.get("ar_dof_config") or self.train_config.get( - "task", {} - ).get("ar_dof_config", {}) + ar_dof_config = self.train_config.get( + "ar_dof_config" + ) or self.train_config.get("task", {}).get("ar_dof_config", {}) self._ar_action_dim = sum(ar_dof_config.values()) def _load_train_config(self, train_config_path): @@ -325,7 +323,9 @@ def _load_model_config_for_wallx(self): self.model_config.update_model_config(self.train_config) self.model_config._attn_implementation = "sdpa" - vision_attn = os.environ.get("WALLX_VISION_ATTN_IMPLEMENTATION", "flash_attention_2") + vision_attn = os.environ.get( + "WALLX_VISION_ATTN_IMPLEMENTATION", "flash_attention_2" + ) self.model_config.vision_config._attn_implementation = vision_attn logger.info( "[LoadModelConfig] vision _attn_implementation=%s (override via WALLX_VISION_ATTN_IMPLEMENTATION)", diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py b/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py index 85ae6f0f..5bceedd8 100644 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/model_wrapper.py @@ -1,6 +1,5 @@ import os import torch -import json import copy from safetensors.torch import load_file import numpy as np @@ -27,12 +26,10 @@ resolve_state_bins, resolve_use_state_string_representation, ) -from wall_x.model.core.action.normalizer import Normalizer 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" diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py b/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py index d86e4e7a..c144302a 100644 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/robot.py @@ -3,7 +3,10 @@ import re from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig -from wall_x._vendor.harrix.serving._wallx_infer.utils import VehiclePoseHandler, UnifiedTrajectoryProcessor +from wall_x._vendor.harrix.serving._wallx_infer.utils import ( + VehiclePoseHandler, + UnifiedTrajectoryProcessor, +) from wall_x._vendor.harrix.serving._wallx_infer.base_dataclass import ( RobotStateActionData, dof_dims, @@ -93,9 +96,7 @@ def ingest_websocket_follow_pos( robot_state_action_data.save_state_data_with_key(grip, key) for key, dim in agent_cfg.items(): - if key == "action_padding" or key.startswith( - ("follow_left_", "follow_right_") - ): + if key == "action_padding" or key.startswith(("follow_left_", "follow_right_")): continue norm_key = key.replace("follow_", "").replace("master_", "") if robot_state_action_data.data.get(f"state_{norm_key}") is None: @@ -109,9 +110,7 @@ def export_websocket_follow_pos( ) -> dict[str, list]: """Convert model ``robot_state_action_data`` back to ``follow1_pos``/``follow2_pos``.""" left_arm_action = _stack_state_action_series(robot_state_action_data, side="left") - right_arm_action = _stack_state_action_series( - robot_state_action_data, side="right" - ) + right_arm_action = _stack_state_action_series(robot_state_action_data, side="right") return { "follow1_pos": left_arm_action.tolist(), "follow2_pos": right_arm_action.tolist(), @@ -245,7 +244,9 @@ def _stack_state_action_series( return np.concatenate([pos, rot, grip], axis=1) -def _views_to_camera_observation(config: InferConfig, views: dict) -> dict[str, np.ndarray]: +def _views_to_camera_observation( + config: InferConfig, views: dict +) -> dict[str, np.ndarray]: """Map websocket camera keys to model observation keys. Only cameras listed in ``config.cam_names`` are required (e.g. LIBERO uses @@ -280,9 +281,7 @@ def _views_to_camera_observation(config: InferConfig, views: dict) -> dict[str, observation[obs_key] = view_data break else: - raise KeyError( - f"Missing view for {obs_key!r}, tried keys: {possible_keys}" - ) + raise KeyError(f"Missing view for {obs_key!r}, tried keys: {possible_keys}") return observation diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py b/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py index b40eace8..6aabbbb7 100644 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/socket_controller.py @@ -25,7 +25,9 @@ class RobotRegistry: def __init__(self): self.registry = _DEFAULT_ROBOT_CONFIG self.logger = InferLogger.get_controller_logger("RobotRegistry") - self.logger.debug(f"Initialized robot registry with {len(self.registry)} robots") + self.logger.debug( + f"Initialized robot registry with {len(self.registry)} robots" + ) def register_robot(self, robot_id, host, action_port, keyboard_port): if robot_id in self.registry: @@ -90,7 +92,9 @@ def connect(self): port = self.robot_info["action_port"] self.action_sock.bind((host, port)) self.action_sock.listen(1) - self.logger.info(f"Listening on action port: {host}:{self.action_sock.getsockname()[1]}") + self.logger.info( + f"Listening on action port: {host}:{self.action_sock.getsockname()[1]}" + ) action_thread = threading.Thread(target=self.handle_action_client) diff --git a/wall_x/_vendor/harrix/serving/_wallx_infer/utils.py b/wall_x/_vendor/harrix/serving/_wallx_infer/utils.py index 9b49aebe..2d0c1242 100644 --- a/wall_x/_vendor/harrix/serving/_wallx_infer/utils.py +++ b/wall_x/_vendor/harrix/serving/_wallx_infer/utils.py @@ -41,7 +41,9 @@ def run(self): # Numeric input: switch instruction index index = int(user_input) self.new_instruction_index = index - self.logger.info(f"[Keyboard] Switched to instruction index: {index}") + self.logger.info( + f"[Keyboard] Switched to instruction index: {index}" + ) else: self.logger.info(f"[Keyboard] Input {user_input}; no action taken.") diff --git a/wall_x/_vendor/harrix/serving/client.py b/wall_x/_vendor/harrix/serving/client.py index 1dfdbaf4..c4639f7f 100644 --- a/wall_x/_vendor/harrix/serving/client.py +++ b/wall_x/_vendor/harrix/serving/client.py @@ -215,9 +215,7 @@ def main_sync(args): print(response["action"]) end_time = time.perf_counter() print(f"Time taken: {(end_time - start_time) * 1000} ms") - action_label = torch.load( - "./example_data" - ).numpy() + action_label = torch.load("./example_data").numpy() np.testing.assert_array_equal(response["action"], action_label) print("compare success") client.close_sync() diff --git a/wall_x/_vendor/harrix/serving/launch_serving.py b/wall_x/_vendor/harrix/serving/launch_serving.py index 47a3420a..2821364d 100644 --- a/wall_x/_vendor/harrix/serving/launch_serving.py +++ b/wall_x/_vendor/harrix/serving/launch_serving.py @@ -15,7 +15,6 @@ import os import socket import sys -import yaml import traceback import tyro @@ -38,9 +37,11 @@ def get_wallx_policy(model_config, image_passing_mode, serialize_actions=True): 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) - - + return WallXPolicy( + config=config, + image_passing_mode=image_passing_mode, + serialize_actions=serialize_actions, + ) logger = logging.getLogger(__name__) @@ -97,6 +98,7 @@ class ServerModelConfig: qwen25_prompt_template: str | None = None prompt_priority_order: str | None = None + @dataclasses.dataclass class Args: """Arguments for the serve_wall_x script.""" diff --git a/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py b/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py index 73984592..6ceafed0 100644 --- a/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py +++ b/wall_x/_vendor/harrix/serving/policy/wall_x_policy.py @@ -102,15 +102,15 @@ def _is_single_arm_right_only(self) -> bool: or {} ) skip = {"action_padding"} - has_left = any( - k.startswith("follow_left_") for k in agent_pos if k not in skip - ) + 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]: + 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): diff --git a/wall_x/_vendor/harrix/serving/scheduler.py b/wall_x/_vendor/harrix/serving/scheduler.py index d114f022..8d36fdb3 100644 --- a/wall_x/_vendor/harrix/serving/scheduler.py +++ b/wall_x/_vendor/harrix/serving/scheduler.py @@ -31,26 +31,30 @@ def __init__( self.max_wait_time = max_wait_time_ms / 1000.0 self.max_queue_size = max_queue_size self.timeout = timeout_ms / 1000.0 - + self.queue = deque() self.queue_lock = asyncio.Lock() self.queue_not_empty = asyncio.Condition(self.queue_lock) - + self.running = False self.batch_task = None - + async def start(self): self.running = True self.batch_task = asyncio.create_task(self._batch_loop()) - logger.info(f"RequestScheduler started: max_batch={self.max_batch_size}, max_wait={self.max_wait_time*1000}ms") - + logger.info( + f"RequestScheduler started: max_batch={self.max_batch_size}, max_wait={self.max_wait_time*1000}ms" + ) + async def stop(self): self.running = False # Process all remaining requests before shutting down async with self.queue_lock: remaining = len(self.queue) if remaining: - logger.info(f"Graceful shutdown: processing {remaining} remaining request(s)") + logger.info( + f"Graceful shutdown: processing {remaining} remaining request(s)" + ) while True: async with self.queue_lock: if len(self.queue) == 0: @@ -66,30 +70,29 @@ async def stop(self): if self.batch_task: await self.batch_task logger.info("RequestScheduler stopped") - + async def add_request(self, obs: Dict[str, Any]) -> Dict[str, Any]: request_id = str(uuid.uuid4()) future = asyncio.Future() request = Request( - request_id=request_id, - obs=obs, - future=future, - timestamp=time.monotonic() + request_id=request_id, obs=obs, future=future, timestamp=time.monotonic() ) - + async with self.queue_lock: if len(self.queue) >= self.max_queue_size: - raise RuntimeError(f"Queue full: {len(self.queue)}/{self.max_queue_size}") + raise RuntimeError( + f"Queue full: {len(self.queue)}/{self.max_queue_size}" + ) self.queue.append(request) self.queue_not_empty.notify() - + try: result = await asyncio.wait_for(future, timeout=self.timeout) return result except asyncio.TimeoutError: logger.error(f"Request {request_id} timeout after {self.timeout}s") raise - + async def _batch_loop(self): while self.running: try: @@ -98,57 +101,58 @@ async def _batch_loop(self): await self._process_batch(batch) except Exception as e: logger.error(f"Unexpected error in batch loop: {e}", exc_info=True) - + async def _collect_batch(self) -> List[Request]: async with self.queue_not_empty: while self.running and len(self.queue) == 0: await self.queue_not_empty.wait() - + if not self.running: return [] - + batch = [] deadline = time.monotonic() + self.max_wait_time - + while len(batch) < self.max_batch_size: if len(self.queue) > 0: batch.append(self.queue.popleft()) - + if len(batch) >= self.max_batch_size: break - + if len(self.queue) == 0: remaining = deadline - time.monotonic() if remaining <= 0: break try: await asyncio.wait_for( - self.queue_not_empty.wait(), - timeout=remaining + self.queue_not_empty.wait(), timeout=remaining ) except asyncio.TimeoutError: break - + return batch - + async def _process_batch(self, batch: List[Request]): if not batch: return - + start_time = time.monotonic() logger.info(f"Processing batch of {len(batch)} requests") - + try: obs_list = [req.obs for req in batch] results = await asyncio.to_thread(self.policy.infer_batch, obs_list) - + for req, result in zip(batch, results): if not req.future.done(): req.future.set_result(result) - + infer_time = time.monotonic() - start_time - logger.info(f"Batch processed in {infer_time*1000:.1f}ms, throughput: {len(batch)/infer_time:.1f} req/s") - + logger.info( + f"Batch processed in {infer_time*1000:.1f}ms, throughput: {len(batch)/infer_time:.1f} req/s" + ) + except Exception as e: logger.error(f"Batch processing failed: {e}", exc_info=True) for req in batch: diff --git a/wall_x/_vendor/harrix/serving/websocket_policy_server.py b/wall_x/_vendor/harrix/serving/websocket_policy_server.py index 89aadc93..54ae810e 100644 --- a/wall_x/_vendor/harrix/serving/websocket_policy_server.py +++ b/wall_x/_vendor/harrix/serving/websocket_policy_server.py @@ -81,7 +81,9 @@ def __init__( self._scheduler = RequestScheduler( policy=policy, max_batch_size=max_batch_size, - max_wait_time_ms=max_wait_time_ms if max_wait_time_ms is not None else 0, + max_wait_time_ms=( + max_wait_time_ms if max_wait_time_ms is not None else 0 + ), max_queue_size=max_queue_size, timeout_ms=timeout_ms, ) @@ -112,7 +114,9 @@ async def run(self): process_request=_health_check, ) as server: mode_str = "batched" if self.batching_enabled else "single" - logger.info(f"Server started on {self._host}:{self._port} (mode={mode_str})") + logger.info( + f"Server started on {self._host}:{self._port} (mode={mode_str})" + ) await server.serve_forever() finally: if self._scheduler is not None: diff --git a/wall_x/_vendor/harrix/utils/normalizer.py b/wall_x/_vendor/harrix/utils/normalizer.py index 2494b83e..9df44675 100644 --- a/wall_x/_vendor/harrix/utils/normalizer.py +++ b/wall_x/_vendor/harrix/utils/normalizer.py @@ -40,7 +40,9 @@ def _load_custom_action_stats(train_config: dict) -> dict | None: return json.load(f) -def _normalizer_from_stats(action_stats: dict, train_config: dict, key: str) -> Normalizer: +def _normalizer_from_stats( + action_stats: dict, train_config: dict, key: str +) -> Normalizer: return Normalizer( action_stats, train_config[key], @@ -49,7 +51,9 @@ def _normalizer_from_stats(action_stats: dict, train_config: dict, key: str) -> ) -def _missing_normalizer_error(checkpoint_path: str, train_config: dict) -> FileNotFoundError: +def _missing_normalizer_error( + checkpoint_path: str, train_config: dict +) -> FileNotFoundError: custom = train_config.get("customized_action_statistic_dof", None) return FileNotFoundError( "Public inference requires normalization data. Expected one of: " @@ -75,18 +79,24 @@ def build_normalizers( action_pth = os.path.join(checkpoint_path, "normalizer_action.pth") propri_pth = os.path.join(checkpoint_path, "normalizer_propri.pth") 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)): + 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) if os.path.exists(action_pth): normalizer_action = Normalizer.from_ckpt(action_pth) else: - normalizer_action = _normalizer_from_stats(custom_stats, train_config, "dof_config") + normalizer_action = _normalizer_from_stats( + custom_stats, train_config, "dof_config" + ) if os.path.exists(propri_pth): normalizer_propri = Normalizer.from_ckpt(propri_pth) else: - normalizer_propri = _normalizer_from_stats(custom_stats, train_config, "agent_pos_config") + normalizer_propri = _normalizer_from_stats( + custom_stats, train_config, "agent_pos_config" + ) action_dim = sum(resolve_dof_config(train_config).values()) propri_dim = sum(resolve_agent_pos_config(train_config).values()) diff --git a/wall_x/_vendor/x2robot_utils/grounding.py b/wall_x/_vendor/x2robot_utils/grounding.py index 0852cd01..60576f44 100644 --- a/wall_x/_vendor/x2robot_utils/grounding.py +++ b/wall_x/_vendor/x2robot_utils/grounding.py @@ -96,9 +96,7 @@ def reverse_match(match): orig_x = max(0, min(orig_width - 1, round(x / scale_w))) orig_y = max(0, min(orig_height - 1, round(y / scale_h))) elif model_type in ["qwen2", "qwen3", "qwen3_5"]: - orig_x = max( - 0, min(orig_width - 1, round((x / 1000) * orig_width)) - ) + orig_x = max(0, min(orig_width - 1, round((x / 1000) * orig_width))) orig_y = max( 0, min(orig_height - 1, round((y / 1000) * orig_height)) ) diff --git a/wall_x/_vendor/x2robot_utils/text_templates.py b/wall_x/_vendor/x2robot_utils/text_templates.py index ae1ca20e..ee788df9 100644 --- a/wall_x/_vendor/x2robot_utils/text_templates.py +++ b/wall_x/_vendor/x2robot_utils/text_templates.py @@ -32,22 +32,30 @@ def pad_text_input_to_target_length( device=text_inputs.attention_mask.device, ) if padding_side == "right": - text_inputs["input_ids"] = torch.cat([text_inputs.input_ids, padding], dim=1) + text_inputs["input_ids"] = torch.cat( + [text_inputs.input_ids, padding], dim=1 + ) text_inputs["attention_mask"] = torch.cat( [text_inputs.attention_mask, attention_padding], dim=1 ) else: - text_inputs["input_ids"] = torch.cat([padding, text_inputs.input_ids], dim=1) + text_inputs["input_ids"] = torch.cat( + [padding, text_inputs.input_ids], dim=1 + ) text_inputs["attention_mask"] = torch.cat( [attention_padding, text_inputs.attention_mask], dim=1 ) elif current_length > target_length: if padding_side == "right": text_inputs["input_ids"] = text_inputs.input_ids[:, :target_length] - text_inputs["attention_mask"] = text_inputs.attention_mask[:, :target_length] + text_inputs["attention_mask"] = text_inputs.attention_mask[ + :, :target_length + ] else: text_inputs["input_ids"] = text_inputs.input_ids[:, -target_length:] - text_inputs["attention_mask"] = text_inputs.attention_mask[:, -target_length:] + text_inputs["attention_mask"] = text_inputs.attention_mask[ + :, -target_length: + ] return text_inputs @@ -117,7 +125,9 @@ def preprocesser_call( batch_size = len(prefix_text) if images is not None and len(images) > 0: - image_inputs = processor.image_processor(images=images, return_tensors=return_tensors) + image_inputs = processor.image_processor( + images=images, return_tensors=return_tensors + ) image_grid_thw = image_inputs["image_grid_thw"] else: image_inputs = {} @@ -125,7 +135,9 @@ def preprocesser_call( if videos is not None: if hasattr(processor, "video_processor"): - videos_inputs = processor.video_processor(videos=videos, return_tensors=return_tensors) + videos_inputs = processor.video_processor( + videos=videos, return_tensors=return_tensors + ) else: videos_inputs = processor.image_processor( images=None, videos=videos, return_tensors=return_tensors @@ -144,13 +156,19 @@ def preprocesser_call( ) if norm_state is not None: - norm_state = norm_state.cpu().numpy() if isinstance(norm_state, torch.Tensor) else norm_state + norm_state = ( + norm_state.cpu().numpy() + if isinstance(norm_state, torch.Tensor) + else norm_state + ) agent_pos_mask = ( agent_pos_mask[:, 0, :].cpu().numpy().astype(bool) if isinstance(agent_pos_mask, torch.Tensor) else agent_pos_mask[:, 0, :].astype(bool) ) - discretized = np.digitize(norm_state, bins=np.linspace(-1, 1, state_bins + 1)[:-1]) - 1 + discretized = ( + np.digitize(norm_state, bins=np.linspace(-1, 1, state_bins + 1)[:-1]) - 1 + ) discretized = discretized[:, 0, :] for i in range(batch_size): if "<|propri|>" not in prefix_text[i]: @@ -188,7 +206,9 @@ def preprocesser_call( ) text_inputs = BatchEncoding( data={ - "input_ids": torch.cat([prefix_inputs.input_ids, postfix_inputs.input_ids], dim=1), + "input_ids": torch.cat( + [prefix_inputs.input_ids, postfix_inputs.input_ids], dim=1 + ), "attention_mask": torch.cat( [prefix_inputs.attention_mask, postfix_inputs.attention_mask], dim=1 ), @@ -200,16 +220,23 @@ def preprocesser_call( if pad_token_id is None: pad_token_id = processor.tokenizer.eos_token_id if pad_to_128_multiple: - target_length = 128 * ((max(len(t) for t in text_inputs.input_ids) + 127) // 128) + target_length = 128 * ( + (max(len(t) for t in text_inputs.input_ids) + 127) // 128 + ) text_inputs = pad_text_input_to_target_length( - text_inputs, target_length, pad_token_id=pad_token_id, padding_side=padding_side + text_inputs, + target_length, + pad_token_id=pad_token_id, + padding_side=padding_side, ) text_inputs["labels"] = None if inference_mode else None return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}) -def get_prologue_with_embodied_information(dataset_name, cam_mapping, robot_id, uid, config): +def get_prologue_with_embodied_information( + dataset_name, cam_mapping, robot_id, uid, config +): """Return a generic VLA system prompt without private robot maps.""" role_start = "<|im_start|>" role_end = "<|im_end|>" diff --git a/wall_x/data/backends/lerobot/loader.py b/wall_x/data/backends/lerobot/loader.py index d7698534..e5d5a960 100644 --- a/wall_x/data/backends/lerobot/loader.py +++ b/wall_x/data/backends/lerobot/loader.py @@ -154,16 +154,20 @@ def __init__( "(action=%s, state=%s; raw action dim=%s -> %s)", self._convert_action_euler_to_6d, self._convert_state_euler_to_6d, - _euler_layout_dim(self._dof_config) - if self._convert_action_euler_to_6d - else "-", - sum( - d - for k, d in self._dof_config.items() - if k not in RELATIVE_SKIP_KEYS - ) - if self._convert_action_euler_to_6d - else "-", + ( + _euler_layout_dim(self._dof_config) + if self._convert_action_euler_to_6d + else "-" + ), + ( + sum( + d + for k, d in self._dof_config.items() + if k not in RELATIVE_SKIP_KEYS + ) + if self._convert_action_euler_to_6d + else "-" + ), ) def _maybe_convert_euler_to_6d(self, vec, layout_config: dict, enabled: bool): diff --git a/wall_x/model/qact/tokenizer_mixin.py b/wall_x/model/qact/tokenizer_mixin.py index a5bd3a65..42f70eaf 100644 --- a/wall_x/model/qact/tokenizer_mixin.py +++ b/wall_x/model/qact/tokenizer_mixin.py @@ -15,7 +15,6 @@ SpatialActionTokenizer = None - class ActionTokenizerMixin(ABC): """Base class for action tokenizers""" @@ -931,7 +930,6 @@ def update_placeholder_mask(self, processor, prefix_length, input_ids): return {"ar_action_mask": ar_action_mask} - # ============================================================ # Factory function # ============================================================ diff --git a/wall_x/trainer/adapters/base_adapter.py b/wall_x/trainer/adapters/base_adapter.py index 269da803..af6f61e2 100644 --- a/wall_x/trainer/adapters/base_adapter.py +++ b/wall_x/trainer/adapters/base_adapter.py @@ -1,4 +1,3 @@ -import dataclasses import logging import os from abc import ABC, abstractmethod diff --git a/wall_x/trainer/fsdp_trainer/checkpoint_io.py b/wall_x/trainer/fsdp_trainer/checkpoint_io.py index bb21b6a3..35b2af05 100644 --- a/wall_x/trainer/fsdp_trainer/checkpoint_io.py +++ b/wall_x/trainer/fsdp_trainer/checkpoint_io.py @@ -437,7 +437,11 @@ def _maybe_resize_token_embeddings_for_load( if key not in state_dict: continue ckpt_vocab = state_dict[key].shape[0] - embed = model.get_input_embeddings() if hasattr(model, "get_input_embeddings") else None + embed = ( + model.get_input_embeddings() + if hasattr(model, "get_input_embeddings") + else None + ) if embed is None: return cur_vocab = embed.weight.shape[0] @@ -449,7 +453,9 @@ def _maybe_resize_token_embeddings_for_load( ) if hasattr(model, "resize_token_embeddings"): model.resize_token_embeddings(ckpt_vocab) - elif hasattr(model, "model") and hasattr(model.model, "resize_token_embeddings"): + elif hasattr(model, "model") and hasattr( + model.model, "resize_token_embeddings" + ): model.model.resize_token_embeddings(ckpt_vocab) return