Skip to content

refactor: consolidate numpy geometry helpers under unilab.utils#677

Closed
TATP-233 wants to merge 2 commits into
mainfrom
refactor/utils-geometry-consolidation
Closed

refactor: consolidate numpy geometry helpers under unilab.utils#677
TATP-233 wants to merge 2 commits into
mainfrom
refactor/utils-geometry-consolidation

Conversation

@TATP-233

@TATP-233 TATP-233 commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Move unilab.envs.common.{rotation,math}unilab.utils.{rotation,math} so the shared numeric helpers live in utils/ alongside device, tensor, nan_guard. Delete the now-empty unilab.envs.common package.
  • Fold math.py's lone np_sample_uniform into unilab.utils.geometry so utils/ no longer hosts a single-function module.
  • Extract per-env private geometry helpers (previously duplicated across motion_tracking/g1/tracking.py, manipulation/stewart/balance.py, manipulation/allegro_inhand/rotation.py, manipulation/sharpa_inhand/rotation.py, locomotion/go2_arm/base.py, locomotion/go2_arm/manip_loco.py) into a new unilab.utils.geometry module: quaternion normalize, roll/pitch, gravity-z-in-body, orientation-error-local, angular-velocity-from-pair, uniform quaternion sampling, spherical↔cartesian, and a fused frame-transform + 6D-flatten writer.
  • Delete dead code: _normalize_quat and _orientation_error_local in go2_arm/base.py no longer have callers after the util replaces them.
  • Update all 30+ import sites across src/, tests/, scripts/, sphinx docs, AGENTS.md pointer, and the utils allowlist test.

Numerical equivalence

Each replaced helper was probed against its pre-refactor implementation on random inputs (batched float32 quaternions/positions, various rotation axes, sphere↔cart samples). All 14 checked functions produce bit-identical output:

Function Bit-equal
np_gravity_z_in_body_from_quat
np_roll_pitch_from_quat (+ stewart float32 wrapper)
normalize_rotation_axis (allegro)
compute_ball_angvel (allegro)
sample_random_quaternion (sharpa)
np_quat_orientation_error_local (go2_arm)
_sphere2cart / _cart2sphere (go2_arm aliases)
np_write_relative_anchor_transform_pos_rot6d (tracking, pos + rot6d)
np_sample_uniform (moved)

normalize_rotation_axis initially drifted by 1 ULP (~6e-08) for non-unit input tuples because np_normalize_axis computes the norm as a Python float (float64). Fixed by pre-casting the input tuple to get_global_dtype() before delegating, so the norm and division happen at the training dtype exactly as before.

Validation

  • make test-all1410 passed, 27 skipped, 258 deselected (43.15s, coverage 73%).
  • Standalone bit-equivalence probes for all 14 refactored functions on random float32 inputs — all bit-identical.
  • All env module imports smoke-tested (G1MotionTrackingEnv, StewartBalanceEnv, AllegroRotationPPO, SharpaInhandRotationEnv, Go2ArmBaseEnv, Go2ArmManipLocoEnv, G1BoxTrackingEnv).

Test plan

  • make test-all passes
  • Bit-equivalence probes on random inputs for every refactored helper
  • sample_random_quaternion monkeypatch test in tests/test_sharpa.py still works (import at module scope preserved)
  • _cart2sphere direct import in tests/envs/locomotion/go2_arm/test_manip_loco_contract.py still works (aliased at import)
  • unilab.utils allowlist test updated with geometry, rotation (and no longer requires math)

TATP-233 added 2 commits July 7, 2026 20:40
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.

@TATP-233 TATP-233 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

check

@TATP-233

TATP-233 commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator Author

关闭本 PR:本次改动已直接合并进 perf/g1-walk-numba-accel 分支(merge commit e7bedcc0),未走单独入 main 的流程。

合并过程中解决了两处冲突:

  • tests/utils/test_utils_package_policy.py:白名单取并集,同时保留 geometrynumba_geometryrotation
  • tests/utils/test_numba_geometry.py:perf 分支在切分之后新增的 unilab.envs.common.rotation 导入改写为 unilab.utils.rotation

合并后 tests/utils + tests/envs + tests/test_sharpa.py 共 335 用例通过;剩余 4 个 test_g1_motion_tracking_clip_end_* 失败为 perf 分支合并前既存的 _numba_accelerator AttributeError(git stash 已核对),与本次重构无关。

@TATP-233 TATP-233 closed this Jul 7, 2026
@TATP-233 TATP-233 deleted the refactor/utils-geometry-consolidation branch July 7, 2026 14:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant