perf: accelerate G1 locomotion hot paths with numba#695
Merged
Conversation
Move the shared numeric helpers out of unilab.envs.common into unilab.utils
and extract per-env private geometry/coordinate helpers (tracking, stewart,
allegro, sharpa, go2_arm) into a new unilab.utils.geometry module. Each
extraction preserves the caller-visible dtype cast (e.g. get_global_dtype
for allegro rotation axis) so refactored code is bit-identical to the
pre-refactor implementation across random inputs.
Notable moves:
- envs/common/{rotation,math}.py -> utils/{rotation,math}.py, then math.py's
lone np_sample_uniform folded into utils/geometry.py so utils/ no longer
hosts a single-function module.
- utils/geometry.py picks up the frame-transform, roll/pitch, gravity-in-body,
quat-normalize, orientation-error, angvel-from-pair, uniform-quat-sample,
and sphere/cart conversions that were duplicated across envs.
- Dead helpers _normalize_quat and _orientation_error_local removed from
go2_arm/base.py; call sites now hit the shared util directly.
- All 30+ import sites (src/tests/scripts/docs/AGENTS.md) updated; utils
package whitelist and sphinx toctree adjusted.
Extend SimBackend.set_state contract to return an optional
{"timing": {...}} dict so the DR manager can surface per-substep timings
next to the outer dr_reset_set_state_ms wall-clock. Both Motrix and
MuJoCo backends emit the same 16-key schema for column stability.
Instrument Motrix set_state into: qpos_convert, mask, data_slice,
data_reset, clear_forces, geom_overrides, reset_rand, set_dof_vel,
set_dof_pos, actuator_ctrl, forward_kinematic, refresh_pose_cache,
invalidate_velocity, internal_gap. MuJoCo populates pool_reset +
state_scatter; the motrix-only keys stay at 0.0 for schema parity.
Apply owner-layer Motrix set_state optimizations, all driven by the new
sub-timing report:
- Reusable scratch buffers for the bool mask and the qpos motrix
conversion, sized once at backend init; no per-reset np.zeros /
np.array(copy=True) allocation.
- Hoist env_ids_intp = np.asarray(env_indices, dtype=np.intp) to the
top of set_state and pass it into _clear_applied_body_forces,
_apply_init_geom_size_overrides, _apply_reset_randomization, and
_refresh_link_pose_cache so the helpers stop re-converting.
- In-place _mujoco_qpos_to_motrix_into(dst, qpos) for the hot path;
original _mujoco_qpos_to_motrix stays for cold paths.
- Skip np.ascontiguousarray on actuator ctrl when the slice is already
c-contiguous.
Benchmark harness (benchmark_offpolicy_collector_active.py) exposes the
new keys via NP_ENV_STEP_TIMING_KEYS + CSV fields, plus two new
formatters _format_set_state_detail_table (motrix keyset) and
_format_set_state_mujoco_table.
New tests: schema + parity contract for both backends
(tests/base/test_sim_backend_set_state_timing.py), DR manager merge +
malformed-timing tolerance (tests/dr/test_manager.py), benchmark CSV
schema + table formatters (tests/benchmark/test_offpolicy_collector_active_benchmark.py).
New microbenchmark: benchmark/benchmark_motrix_set_state_ab.py drives
set_state under baseline vs optimized variants to isolate the Python-
side wins from Rust-boundary noise.
Local A/B (Apple M5 Max, 8192 env, motrix, 100 measured steps):
- dr_reset_set_state_ms: 3.274 ms (walk_flat numba), fully broken down
by sub-key (Data reset 37% + FK 20% + Data slice 16% dominant, all in
Rust boundary).
- reset_done_ms: g1_walk_flat -0.102 ms, g1_motion_tracking -0.157 ms
vs b96224a baseline (issue #663 comment 4904777445).
- set_state microbenchmark: mask 1.6-2.3x, qpos_convert 1.05-1.35x,
actuator_ctrl 1.07-1.17x; outer set_state 1.08-1.16x at low reset
ratios (which is the collector's normal regime).
Refs: #679, #665
fix: 限制 off-policy Torch CPU 线程预算
feat: show collector env step breakdown
…mpling benchmark: add SAC replay sampling placement test
# Conflicts: # src/unilab/algos/torch/offpolicy/double_buffer_runner.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Validation