perf: optimize g1 motion tracking state updates

[caozx1110]

Summary

• Optimize G1 motion tracking hot-path state updates by reusing output buffers and avoiding repeated temporary allocations.
• Add backend body-state copy helpers and coverage to keep the optimized path contract-driven across MuJoCo and Motrix.
• Keep benchmarked behavior equivalent with reference tests for relative transforms and reward computation.

Benchmark

Command shape:

uv run --active --no-sync benchmark/benchmark_env_step.py task=g1_motion_tracking/<backend> num_envs=2048 num_steps=50 warmup_steps=10 --skip-plots

Baseline: upstream/main 367d9c4
Current: optimize/g1-mt-update-state 455ecea

Backend  Metric        upstream/main        current        Improvement
-------  ------------  -------------------  -------------  -----------
mujoco   median_step     39.411 ms         34.463 ms      12.55% faster
mujoco   throughput       50881 eps         58340 eps     14.66% higher
mujoco   update_state     9.498 ms          4.540 ms      52.20% faster
motrix   median_step     48.766 ms         39.539 ms      18.92% faster
motrix   throughput       42118 eps         51168 eps     21.49% higher
motrix   update_state    11.457 ms          6.218 ms      45.72% faster

Raw benchmark outputs:

/tmp/unilab_pr_bench_20260525/current_g1_mt_mujoco.json
/tmp/unilab_pr_bench_20260525/current_g1_mt_motrix.json
/tmp/unilab_pr_bench_20260525/upstream_main_g1_mt_mujoco.json
/tmp/unilab_pr_bench_20260525/upstream_main_g1_mt_motrix.json

Validation

make test-all
967 passed, 12 skipped, 254 deselected, 66 warnings

Additional run:

CUDA PPO g1_flip_tracking/mujoco, algo.max_iterations=1000
Completed through model_999.pt, final_mean_reward=60.28, best_mean_reward=63.54

[Copilot]

Pull request overview

This PR optimizes the G1 motion tracking “hot path” by reducing per-step allocations and repeated temporary array creation, mainly by reusing preallocated buffers and adding backend copy helpers to support efficient state reads across MuJoCo and Motrix.

Changes:

• Reworked G1MotionTrackingEnv.update_state() / obs / reward paths to reuse buffers and avoid tile/concatenate allocations.
• Added reusable motion gather buffers (MotionData out-params) and backend copy_body_state_w() helpers, plus row-wise sensor/body getters.
• Added/expanded unit and smoke tests to enforce equivalence of the optimized math paths (relative transforms, anchor-frame writers, reward).

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/utils/test_math_utils.py	Adds tests validating new batched quaternion/transform helpers match reference implementations.
tests/envs/test_env_configs.py	Extends env config/unit tests to cover new optimized fast paths and required buffers/flags.
tests/base/test_sim_backend_smoke.py	Adds backend smoke tests for copy_body_state_w() (MuJoCo & Motrix) and row-wise sensor equivalence.
src/unilab/envs/motion_tracking/g1/tracking.py	Main optimization: preallocated buffers, faster relative-transform math, reward computation reuse, reduced temporary allocations.
src/unilab/envs/motion_tracking/g1/motion_loader.py	Adds reusable MotionData buffer allocation and out parameter to avoid repeated allocations.
src/unilab/envs/motion_tracking/g1/motion_box_loader.py	Extends box motion loader to accept out buffers and optionally fill object state.
src/unilab/envs/common/rotation.py	Adds batched quaternion ops and anchor-frame transform helpers to avoid flatten/tile allocations.
src/unilab/base/backend/mujoco/backend.py	Adds copy_body_state_w(), get_body_pose_w_rows(), and get_sensor_data_rows() for efficient state/sensor reads.
src/unilab/base/backend/motrix/backend.py	Adds copy_body_state_w(), get_body_pose_w_rows(), and get_sensor_data_rows() for efficient state/sensor reads.
src/unilab/base/backend/base.py	Defines default get_body_state_w(), copy_body_state_w(), row-wise body pose getters, and row-wise sensor getter in base backend API.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.