Triển khai C++: TensorRT, ZMQ, ONNX

Trong bài 1, chúng ta đã nhìn SONIC như một kiến trúc nhiều encoder, một latent token space và một decoder điều khiển toàn thân. Trong bài 2, ta đã chạy sim2sim và eval để biết checkpoint có hoạt động không. Trong bài 3, ta đã đi qua dữ liệu, motion library và lệnh export ONNX. Bài 4 này đi xuống lớp cuối cùng trước robot thật: deployment C++ trong gear_sonic_deploy.

Mục tiêu của bài không phải là giải thích mọi dòng C++. Mục tiêu là giúp bạn đọc được luồng chương trình: deploy.sh dựng command như thế nào, g1_deploy_onnx_ref.cpp tạo các thread thời gian thực ra sao, TRTInference/InferenceEngine.cpp chuyển ONNX sang TensorRT engine như thế nào, observation_config.hpp quyết định observation layout ra sao, encoder.hpp biến observation cao chiều thành token_state, zmq_output_handler.hpp xuất debug stream thế nào, và visualize_motion.py nhận stream đó để vẽ target/measured robot trong MuJoCo.

Nếu bạn từng deploy YOLO hoặc một policy PyTorch đơn giản, stack này sẽ có vẻ hơi dài. Nhưng với humanoid, độ dài đó có lý do: policy chạy ở 50 Hz, command writer ở 500 Hz, input và planner chạy ở thread riêng, ZMQ không được phép làm nghẽn control loop, và mọi tensor phải khớp đúng dimension giữa ONNX, TensorRT và YAML config.

Nguồn kỹ thuật nên mở song song:

Nếu bạn muốn nối bài này với các workflow deployment rộng hơn trên vnrobo, hãy đọc thêm deploy GR00T N1 trên Unitree G1 và ASAP cho Unitree G1 sim-to-real. Hai bài đó không thay thế stack SONIC, nhưng giúp đặt TensorRT, WBC và sim-to-real vào bối cảnh vận hành robot thật.

1. Deployment stack giải quyết bài toán gì?

Trong training, SONIC có thể chạy trong Isaac Lab với Python, Hydra config, vectorized environments và checkpoint .pt. Trong deployment, mục tiêu khác hẳn: nhận trạng thái robot thật hoặc MuJoCo, tạo observation đúng thứ tự, chạy inference đủ nhanh, phát motor command đều và cung cấp debug stream để operator biết chuyện gì đang xảy ra.

Vì vậy gear_sonic_deploy chia việc thành nhiều lớp:

Đây là pattern production robotics quen thuộc: launcher đơn giản, control app chặt chẽ, inference backend cache engine, debug stream tách khỏi command path. Nếu bạn debug, đừng bắt đầu từ neural network ngay. Hãy đi theo thứ tự launcher -> file path -> observation dimension -> TensorRT tensor name -> control loop -> output stream.

2. Export và đặt đúng bốn artifact

Với checkpoint tự train hoặc checkpoint release dạng .pt, bước bridge sang C++ là export ONNX. Lệnh thường dùng từ repo root:

python gear_sonic/eval_agent_trl.py \
  +checkpoint=/path/to/sonic_release/last.pt \
  +headless=True \
  ++num_envs=1 \
  +export_onnx_only=true

Sau bước này, C++ deployment cần bốn artifact:

Tài liệu download model của NVlabs cũng liệt kê đúng các file này khi tải từ Hugging Face: model_encoder.onnx, model_decoder.onnx, observation_config.yaml, và planner_sonic.onnx. Điểm cần để ý là deploy.sh không nhận riêng decoder path mặc định; nó nhận prefix checkpoint qua --cp, rồi tự ghép:

CHECKPOINT_DECODER="${CHECKPOINT}_decoder.onnx"
CHECKPOINT_ENCODER="${CHECKPOINT}_encoder.onnx"

Vì default CHECKPOINT là policy/release/model, hai file mặc định sẽ là:

policy/release/model_decoder.onnx
policy/release/model_encoder.onnx

Obs config mặc định là:

policy/release/observation_config.yaml

Planner mặc định hiện nằm ở:

planner/target_vel/V2/planner_sonic.onnx

Nếu bạn tải model vào layout khác, đừng sửa code vội. Hãy truyền đường dẫn qua deploy.sh:

cd gear_sonic_deploy

bash deploy.sh sim \
  --cp policy/release/model \
  --obs-config policy/release/observation_config.yaml \
  --planner planner/target_vel/V2/planner_sonic.onnx \
  --input-type keyboard \
  --output-type zmq

Với robot thật, thay sim bằng real hoặc tên interface có IP mạng robot:

bash deploy.sh real \
  --cp policy/release/model \
  --obs-config policy/release/observation_config.yaml \
  --planner planner/target_vel/V2/planner_sonic.onnx \
  --input-type manager \
  --output-type all

Beginner hay gặp lỗi “file tồn tại nhưng TensorRT không load được”. Nguyên nhân thường là Git LFS pointer, tải thiếu từ Hugging Face, hoặc dùng observation_config.yaml không cùng phiên bản với ONNX. Với SONIC, YAML config là một phần của model contract, không phải file cấu hình trang trí.

3. deploy.sh: từ mode sim đến lệnh C++

deploy.sh là launcher thực dụng. Nó làm bốn việc chính.

Thứ nhất, nó resolve interface. Nếu bạn chạy:

bash deploy.sh sim

script chọn loopback interface, thường là lo trên Linux hoặc lo0 trên macOS, và set ENV_TYPE=sim. Nếu chạy:

bash deploy.sh real

script tìm interface có IP dạng 192.168.123.x, tức mạng hay dùng với Unitree. Nếu không tìm thấy, nó fallback sang interface non-loopback đầu tiên hoặc tên mặc định.

Thứ hai, nó kiểm tra file. Các file bắt buộc là decoder ONNX, encoder ONNX, obs config, planner ONNX và motion data directory. Đây là check rất đáng giá vì tránh để binary C++ chết muộn sau khi đã build xong.

Thứ ba, nó source môi trường và build bằng just build. Script cũng kiểm tra TensorRT_ROOT, cmake, clang, git, và just. Với sim mode, nó thêm --disable-crc-check vì MuJoCo không cần CRC theo cách robot thật cần.

Thứ tư, nó dựng command cuối:

just run g1_deploy_onnx_ref "$TARGET" "$CHECKPOINT_DECODER" "$MOTION_DATA" \
  --obs-config "$OBS_CONFIG" \
  --encoder-file "$CHECKPOINT_ENCODER" \
  --planner-file "$PLANNER" \
  --input-type "$INPUT_TYPE" \
  --output-type "$OUTPUT_TYPE" \
  --zmq-host "$ZMQ_HOST"

Điểm cần nhớ: argument thứ hai của binary là decoder model, còn encoder được truyền bằng --encoder-file. Điều này hợp lý vì decoder/policy là thành phần bắt buộc để tạo action; encoder chỉ được dùng khi token_state bật trong observation config. Nếu policy không dùng token, encoder có thể bị ignore. Nếu policy dùng token nhưng không có encoder file, hệ thống có thể nhận token từ bên ngoài, nhưng bạn phải hiểu rủi ro zero-token hoặc token timeout.

4. g1_deploy_onnx_ref.cpp: xương sống runtime

File g1_deploy_onnx_ref.cpp có comment đầu file rất rõ: app tạo bốn thread bằng CreateRecurrentThreadEx.

State machine runtime là:

INIT -> WAIT_FOR_CONTROL -> CONTROL

Ở INIT, robot đợi LowState hợp lệ và ramp về default standing pose. Ở WAIT_FOR_CONTROL, hệ thống publish config cho debug subscriber và chờ operator start. Ở CONTROL, mỗi tick 50 Hz chạy pipeline chính:

1. GatherRobotStateToLogger: lấy IMU, joint state, hand state và log vào StateLogger.
2. GatherInputInterfaceData: snapshot input từ keyboard, gamepad, ZMQ hoặc manager.
3. UpdateHeadingState: đồng bộ heading giữa robot và reference motion.
4. GatherObservations: fill policy observation buffer theo observation_config.yaml.
5. Nếu encoder đang dùng, tạo encoder observation, chạy EncoderEngine, ghi token_state.
6. CreatePolicyCommand: chạy decoder/policy TensorRT để lấy action 29 motor.
7. Update Dex3 hands nếu có hand target.
8. Publish ZMQ/ROS2 output.
9. Ghi motion log nếu bật.
10. Advance frame hoặc blend planner motion.

Nhìn như nhiều bước, nhưng mental model rất đơn giản:

robot state + reference/input
        |
        v
observation_config.yaml quyết định vector obs
        |
        +--> encoder ONNX/TensorRT -> token_state
        |
        v
decoder ONNX/TensorRT -> action 29 motor
        |
        v
low-level command writer 500 Hz

Chính vì control thread chạy 50 Hz, mọi thứ trong tick phải tránh block. ZMQ output dùng non-blocking send. TensorRT engine được convert/cache từ trước. CUDA graph được capture cho policy và encoder để giảm overhead launch kernel. Những chi tiết này nghe “performance”, nhưng trong humanoid chúng cũng là safety: loop không đều có thể làm command lỗi nhịp.

5. TensorRT inference: ONNX không chạy trực tiếp

TRTInference/InferenceEngine.cpp là lớp chung cho encoder, decoder và planner. Luồng của nó gồm hai pha.

Pha build:

1. Kiểm tra file ONNX có tồn tại.
2. Hash nội dung ONNX, tên GPU và precision.
3. Tạo tên .trt cùng thư mục với ONNX, có prefix như policy_, encoder_, planner_.
4. Nếu .trt đã có và hash khớp, dùng lại.
5. Nếu không, dùng TensorRT builder, network explicit batch và ONNX parser để parse model.
6. Tạo optimization profile cho dynamic axes nếu có.
7. Set FP16 nếu option yêu cầu và GPU hỗ trợ.
8. Build serialized network, ghi hash 64 byte rồi ghi engine xuống disk.

Pha runtime:

1. Deserialize .trt, bỏ qua 64 byte hash đầu file.
2. Tạo execution context.
3. Duyệt tất cả I/O tensors, tính shape và data type.
4. cudaMalloc input/output buffer.
5. Gán tensor address vào context.
6. SetInputData copy host -> device.
7. Enqueue gọi enqueueV3.
8. GetOutputData copy device -> host.

Điểm beginner cần hiểu: file .onnx là format trao đổi model; file .trt là engine đã optimize cho GPU, TensorRT version và precision cụ thể. Vì hash có cả tên GPU và precision, đổi GPU hoặc đổi FP16/FP32 có thể làm engine cũ bị xem là stale và build lại. Tài liệu deployment của NVlabs cũng nhấn mạnh cần đúng version TensorRT; dùng sai version có thể làm planner xuất motion sai, đây không phải lỗi cosmetic.

6. observation_config.yaml: contract quan trọng nhất

Trong release config, policy observation tổng là 436 chiều:

observations:
  - name: "token_state"
    enabled: true
  - name: "his_base_angular_velocity_10frame_step1"
    enabled: true
  - name: "his_body_joint_positions_10frame_step1"
    enabled: true
  - name: "his_body_joint_velocities_10frame_step1"
    enabled: true
  - name: "his_last_actions_10frame_step1"
    enabled: true
  - name: "his_gravity_dir_10frame_step1"
    enabled: true

Encoder section khai báo:

encoder:
  dimension: 64
  use_fp16: false
  encoder_modes:
    - name: "g1"
      mode_id: 0
    - name: "teleop"
      mode_id: 1
    - name: "smpl"
      mode_id: 2

observation_config.hpp parse file YAML đơn giản này và validate logic: nếu observation token_state được bật, encoder section phải có dimension > 0; nếu token_state tắt, encoder section bị ignore. Với mỗi observation, C++ có registry tên -> dimension -> function gather. Khi init, app cộng dimension và offset, rồi so với input dimension thật của model. Nếu lệch, app nên fail sớm.

Đây là lỗi phổ biến khi tự export ONNX: bạn thay model_encoder.onnx nhưng dùng nhầm observation_config.yaml của checkpoint khác. GitHub issue của repo từng có ví dụ encoder model input 1751 chiều nhưng config cộng ra 1762 chiều. Sai 11 chiều cũng đủ làm deployment dừng, vì TensorRT không thể đoán observation nào nên bỏ.

Rule thực dụng:

model_encoder.onnx
model_decoder.onnx
observation_config.yaml
planner_sonic.onnx

nên đi theo cùng release/export. Đừng trộn “model mới” với “config cũ” trừ khi bạn tự kiểm tra tensor shape và observation offsets.

7. Chọn input type: keyboard, gamepad, ZMQ, manager

g1_deploy_onnx_ref.cpp hỗ trợ nhiều input interface. Với bài này, năm mode quan trọng là:

ROS2 input cũng có trong code nếu build có HAS_ROS2, nhưng với yêu cầu bài này, hãy xem ROS2 là option phụ thuộc môi trường build. Với người mới, thứ tự học hợp lý là keyboard -> gamepad -> zmq -> manager. Đừng bắt đầu bằng manager nếu bạn chưa biết mode đơn lẻ nào đang fail.

Ví dụ sim2sim keyboard:

bash deploy.sh sim \
  --input-type keyboard \
  --output-type zmq

Ví dụ ZMQ manager và debug output:

bash deploy.sh sim \
  --input-type zmq_manager \
  --output-type all \
  --zmq-host localhost

Với output, code tạo ZMQ khi --output-type zmq hoặc all. ROS2 output được tạo khi --output-type ros2 hoặc all, nhưng chỉ nếu binary được build với ROS2. Nếu bạn truyền ros2 trong build không có ROS2, validation có thể không cho qua hoặc không tạo interface tùy compile flag.

8. ZMQ output và visualize_motion.py

zmq_output_handler.hpp publish qua một ZMQ PUB socket. Mặc định output port là 5557 và topic là g1_debug. Message không phải JSON; nó là msgpack payload có prefix topic:

[topic_prefix][msgpack payload]

Payload gồm state và visualization fields:

Socket options được chọn cho control loop: send HWM 10, send buffer 32 KB, linger 0, và send non-blocking. Ý nghĩa là nếu subscriber chậm, hệ thống drop message cũ thay vì làm nghẽn robot. Đây là quyết định đúng cho debug telemetry: robot command quan trọng hơn việc GUI nhận đủ mọi frame.

visualize_motion.py là subscriber/debug viewer. Nó có ba mode input:

python visualize_motion.py --motion_dir reference/example/
python visualize_motion.py --csv_path some_motion.csv
python visualize_motion.py \
  --realtime_debug_url tcp://localhost:5557 \
  --realtime_debug_topic g1_debug

Ở realtime mode, script subscribe topic, unpack msgpack, rồi cập nhật một dictionary có target pose, measured pose, VR 3-point và nhiệt độ motor nếu có. Sau đó nó dùng MuJoCo viewer để hiển thị nhiều robot overlay: target, measured, reference/ghost và robot nhiệt độ. Khi bạn thấy target đi một đường còn measured trễ hoặc lệch nhiều, đó là tín hiệu để xem lại policy, planner, observation, hoặc trạng thái robot.

9. Checklist debug cho beginner

Khi deployment không chạy, hãy đi theo checklist tuyến tính:

1. deploy.sh có tìm đúng interface không? Với sim phải là loopback; với real phải là mạng robot.
2. TensorRT_ROOT có đúng version không? Tài liệu NVlabs yêu cầu version cụ thể cho desktop và Jetson.
3. model_decoder.onnx, model_encoder.onnx, observation_config.yaml, planner_sonic.onnx có phải cùng release/export không?
4. --cp có là prefix không? Ví dụ policy/release/model, không phải policy/release/model_decoder.onnx.
5. Policy input tensor có tên obs_dict và output có tên action không?
6. Encoder input tensor có tên obs_dict và output có tên encoded_tokens không?
7. Observation dimension log có khớp model input dimension không?
8. Planner path có chứa version V0, V1 hoặc V2 không? Code dùng path để suy ra planner version.
9. Nếu dùng ZMQ, publisher/subscriber có cùng host, port và topic không?
10. Nếu GUI không hiện gì, kiểm tra --output-type zmq|all và visualize_motion.py --realtime_debug_url tcp://localhost:5557.

Một lưu ý an toàn: đừng dùng robot thật để debug những lỗi có thể debug trong MuJoCo. Tài liệu Quick Start của NVlabs cũng khuyến nghị làm quen sim2sim trước khi deploy hardware. Với humanoid, “chạy được binary” chưa đồng nghĩa “an toàn để start control”. Hãy để sim, ZMQ output và visualize viewer nói cho bạn biết pipeline đã đúng nhịp chưa.

10. Kết luận

Deployment SONIC không chỉ là “load ONNX rồi infer”. Nó là một hệ thống realtime nhỏ: shell launcher chọn interface và file, C++ app chia thread, TensorRT build/cache engine, observation config khóa contract tensor, encoder tạo token, decoder tạo action, planner tạo reference motion, ZMQ xuất debug, và MuJoCo viewer giúp operator nhìn được target so với measured.

Nếu bạn nắm một câu sau bài này, hãy nắm câu này: đừng debug deployment bằng cảm giác; debug bằng artifact contract. Contract đó gồm bốn file model_encoder.onnx, model_decoder.onnx, planner_sonic.onnx, observation_config.yaml, cộng với input/output mode bạn chọn. Khi bốn file khớp nhau và ZMQ viewer phản ánh đúng target/measured state, bạn mới có nền tảng để đi tiếp sang teleop và VLA trong bài 5.

Bài viết liên quan

• Kiến trúc SONIC cho WBC humanoid
• Dữ liệu BONES-SEED và huấn luyện SONIC
• Teleop và VLA với SONIC