Add Rotation Initialization for Objects#738
Conversation
zhx06
requested review from
alexmillane,
cvolkcvolk,
peterd-NV,
qianl-nv,
viiik-inside and
xyao-nv
as code owners
![isaaclab-review-bot[bot]](https://avatars.githubusercontent.com/u/157846462?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤖 Isaac Lab Review Bot
Summary
This PR adds random yaw initialization for placed objects via a new --random_yaw_init CLI flag. The implementation uses conservative bounding boxes (axis-aligned boxes enclosing the rotated object) for collision detection, ensuring layout validity while adding scene variety. The code is well-structured, follows Arena patterns, and includes comprehensive test coverage.
Findings
🔵 Suggestion - isaaclab_arena/utils/pose.py:60-67 - Consider caching or inlining the quaternion composition for hot paths. The current implementation creates two temporary Pose objects per call to rotate_quat_by_yaw. While correct, if this is called in tight loops (multi-env × objects), a direct quaternion multiply formula would be more efficient:
# Direct Z-rotation composition (avoids Pose allocation):
# yaw_q = (0, 0, sin(θ/2), cos(θ/2)); base_q = (x, y, z, w)
# result_w = w*cos - z*sin; result_z = w*sin + z*cos; result_x/y unchanged for pure-Z🔵 Suggestion - isaaclab_arena/relations/object_placer.py:232 - The yaw is sampled uniformly in [-π, π) but the comment says "uniform in [-pi, pi)". Since torch.rand produces [0, 1), the formula (2u - 1) * π actually gives [-π, π), which is correct. However, for symmetry, [0, 2π) might be more conventional (same distribution, different range). This is purely stylistic.
🔵 Suggestion - isaaclab_arena/relations/placement_result.py:31-32 - The orientations field uses field(default_factory=dict) which is the right pattern, but consider documenting that this dict is keyed by ObjectBase instances (same as positions) in the docstring for API clarity.
🟡 Warning - isaaclab_arena/relations/relation_solver_state.py:112-119 - When initial_orientations is provided but an object has no entry in some env's dict, d.get(obj, 0.0) correctly defaults to 0. However, the conditional if any(yaw != 0.0 for yaw in yaws) means a single non-zero yaw causes bbox rotation for all envs. This is intentional (and documented), but worth noting: if only env[0] has yaw=π/2 while others have 0, the per-env rotated bbox tensor will have different values per row, which is correct. The logic is sound.
🔵 Suggestion - isaaclab_arena/utils/bounding_box.py:231-239 - The Z-extent broadcasting logic handles the N=1, M>1 case explicitly. Consider adding a brief comment explaining this edge case for future maintainers:
# When N=1 box is rotated by M angles, Z extents must broadcast to match M rows.Test Coverage
Excellent. The PR adds 8 new/modified test files with thorough coverage:
test_bounding_box.py: Testsrotated_around_zfor single and batched angles, including 45°/90° edge casestest_no_collision_loss.py: End-to-end test proving the solver uses rotated bboxes (regression guard)test_object_placer_reproducibility.py: Tests yaw composition with markers, multi-env distinct yaws, state mismatch errorstest_placement_events.py: Runtime path verification forsolve_and_place_objectswith yawed posestest_validate_placement.py: Rotation-aware overlap detection
The tests are well-designed—test_solver_uses_rotated_bbox_for_collision would fail if the solver reverted to unrotated bboxes, providing effective regression protection.
CI Status
⏳ Pre-commit check is pending
Verdict
No issues found. Clean implementation with proper device handling, correct quaternion math, and defensive coding (e.g., yaw == 0.0 fast paths). The conservative AABB approach is the right trade-off for a non-optimized parameter. Ready to merge once CI passes.
Update (5b21206): The new commit addresses reviewer comments and adds significant improvements:
Changes reviewed:
- Metrics system refactor -
MetricBase.compute_metric_from_recording()replaced withget_metric_term_cfg()returning aMetricTermCfgcontaining a callable. The env now has aMetricsManagerand exposescompute_metrics()directly. This is a cleaner separation of config vs computation. - Module split -
isaaclab_arena_manager_based_env.pysplit into_env.py(runtime) and_env_cfg.py(config). Good for maintainability. - Bug fix -
compile_env_notebook.py: Fixedrotation_wxyz→rotation_xyzwtypo. - Parameter tuning -
PickAndPlaceTask.force_thresholdchanged from 1.0 → 0.1 (more sensitive success detection). - Docs - Added CloudXR network requirement notices to teleoperation workflows.
- Tests updated - All tests now use
env.unwrapped.compute_metrics()instead of the removedcompute_metrics(env)function.
Assessment: All changes look good. The metrics refactor follows Isaac Lab patterns (configclass + term cfg + manager). No new issues introduced.
Greptile SummaryThis PR introduces per-object random yaw initialization for non-anchor placeable objects, adding a
Confidence Score: 5/5The changes are self-contained: existing code paths are unaffected when random_yaw_init is off (empty dicts default all yaws to 0.0), and the rotated-bbox math is well-tested at multiple granularities. The quaternion composition, conservative AABB refit, per-candidate seeding, and batched bbox shape handling are all correct. Loss strategies implicitly broadcast (batch_size, 3) bboxes correctly with existing (N,3) position tensors. The PlacementResult.orientations field propagates the yaw consistently through pooling and reset events. No files require special attention. The most complex new code — rotated_around_z in bounding_box.py and the _local_bboxes precomputation in relation_solver_state.py — is directly covered by the new tests. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[ObjectPlacer.place] --> B[_generate_initial_positions per candidate]
A --> C[_generate_initial_orientations per candidate]
C --> D{random_yaw_init?}
D -- Yes --> E[Sample yaw uniform in neg-pi to pi per non-anchor]
D -- No --> F[Return empty dict]
E --> G[RelationSolver.solve with initial_orientations]
F --> G
B --> G
G --> H[RelationSolverState: Precompute _local_bboxes rotated_around_z per-env]
H --> I[Loss computation uses rotated bboxes]
I --> J[_validate_placement uses same rotated bboxes]
J --> K[PlacementCandidate with .orientations stored]
K --> L{apply_positions_to_objects?}
L -- Yes --> M[_apply_positions: rotate_quat_by_yaw per env]
L -- No --> N[PlacementResult .orientations in pool]
N --> O[PooledObjectPlacer samples on reset]
O --> P1[solve_and_place_objects: rotate_quat_by_yaw runtime]
O --> P2[_apply_static_initial_poses: rotate_quat_by_yaw at init]
M --> Q[Pose written to sim]
P1 --> Q
P2 --> Q
Reviews (2): Last reviewed commit: "address comments" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def rotate_quat_by_yaw( | ||
| base_xyzw: tuple[float, float, float, float], yaw_rad: float | ||
| ) -> tuple[float, float, float, float]: | ||
| """Rotate ``base_xyzw`` (xyzw) by an extra yaw about Z. Returns base unchanged when yaw is 0.""" | ||
| if yaw_rad == 0.0: | ||
| return base_xyzw | ||
| yaw_xyzw = (0.0, 0.0, math.sin(yaw_rad / 2.0), math.cos(yaw_rad / 2.0)) | ||
| return Pose(rotation_xyzw=base_xyzw).multiply(Pose(rotation_xyzw=yaw_xyzw)).rotation_xyzw |
There was a problem hiding this comment.
compose_poses → matrix_from_quat → matrix multiply → quat_from_matrix to compose two quaternions that are always pure Z-rotations. Direct quaternion multiplication is simpler, removes the isaaclab import dependency from the fast path, and avoids the round-trip precision loss through matrix form. Since base_xyzw can only ever be a pure-Z quaternion in this codebase (identity or from RotateAroundSolution), a direct formula works correctly.
| def rotate_quat_by_yaw( | |
| base_xyzw: tuple[float, float, float, float], yaw_rad: float | |
| ) -> tuple[float, float, float, float]: | |
| """Rotate ``base_xyzw`` (xyzw) by an extra yaw about Z. Returns base unchanged when yaw is 0.""" | |
| if yaw_rad == 0.0: | |
| return base_xyzw | |
| yaw_xyzw = (0.0, 0.0, math.sin(yaw_rad / 2.0), math.cos(yaw_rad / 2.0)) | |
| return Pose(rotation_xyzw=base_xyzw).multiply(Pose(rotation_xyzw=yaw_xyzw)).rotation_xyzw | |
| def rotate_quat_by_yaw( | |
| base_xyzw: tuple[float, float, float, float], yaw_rad: float | |
| ) -> tuple[float, float, float, float]: | |
| """Rotate ``base_xyzw`` (xyzw) by an extra yaw about Z. Returns base unchanged when yaw is 0.""" | |
| if yaw_rad == 0.0: | |
| return base_xyzw | |
| bx, by, bz, bw = base_xyzw | |
| sz = math.sin(yaw_rad / 2.0) | |
| cz = math.cos(yaw_rad / 2.0) | |
| # Hamilton product: base * yaw_quat (both are Z-axis rotations, so they commute) | |
| return ( | |
| bx * cz + by * sz, | |
| -bx * sz + by * cz, | |
| bz * cz + bw * sz, | |
| -bz * sz + bw * cz, | |
| ) |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| poses = [] | ||
| for env_idx in range(num_envs): | ||
| pos = layouts[env_idx].positions.get(obj) | ||
| if pos is None: | ||
| break | ||
| rotation_xyzw = rotate_quat_by_yaw(get_rotation_xyzw(obj), layouts[env_idx].orientations.get(obj, 0.0)) | ||
| poses.append(Pose(position_xyz=pos, rotation_xyzw=rotation_xyzw)) |
There was a problem hiding this comment.
get_rotation_xyzw(obj) is called once per env_idx inside the inner loop, but it returns the same base rotation for every environment (it reads from the object's static marker, not from the layout). Hoisting it out avoids num_envs redundant calls per object.
| poses = [] | |
| for env_idx in range(num_envs): | |
| pos = layouts[env_idx].positions.get(obj) | |
| if pos is None: | |
| break | |
| rotation_xyzw = rotate_quat_by_yaw(get_rotation_xyzw(obj), layouts[env_idx].orientations.get(obj, 0.0)) | |
| poses.append(Pose(position_xyz=pos, rotation_xyzw=rotation_xyzw)) | |
| poses = [] | |
| base_rotation_xyzw = get_rotation_xyzw(obj) | |
| for env_idx in range(num_envs): | |
| pos = layouts[env_idx].positions.get(obj) | |
| if pos is None: | |
| break | |
| rotation_xyzw = rotate_quat_by_yaw(base_rotation_xyzw, layouts[env_idx].orientations.get(obj, 0.0)) | |
| poses.append(Pose(position_xyz=pos, rotation_xyzw=rotation_xyzw)) |
| device = self._min_point.device | ||
| angles = torch.as_tensor(angle_rad, dtype=torch.float32, device=device).reshape(-1) # (M,) | ||
|
|
||
| cos = torch.cos(angles).unsqueeze(1) # (M, 1) |
There was a problem hiding this comment.
(N, 4) corners against (M, 1) cos/sin will raise a cryptic size-mismatch error rather than a clear message. A short assertion would make the constraint visible at call time.
| device = self._min_point.device | |
| angles = torch.as_tensor(angle_rad, dtype=torch.float32, device=device).reshape(-1) # (M,) | |
| cos = torch.cos(angles).unsqueeze(1) # (M, 1) | |
| device = self._min_point.device | |
| angles = torch.as_tensor(angle_rad, dtype=torch.float32, device=device).reshape(-1) # (M,) | |
| n, m = self._min_point.shape[0], angles.shape[0] | |
| assert n == 1 or m == 1 or n == m, ( | |
| f"rotated_around_z: box count (N={n}) and angle count (M={m}) must satisfy N==1, M==1, or N==M." | |
| ) | |
| cos = torch.cos(angles).unsqueeze(1) # (M, 1) |
Signed-off-by: zhx06 <zihaox@nvidia.com>
zhx06
force-pushed
the
zxiao/random_object_rotation_init
branch
from
a3ea787 to
5b21206
Compare
| resolve_on_reset: Optional override for whether to draw fresh layouts from | ||
| the placement pool on reset. When ``False``, fixed per-environment | ||
| initial poses are applied immediately. | ||
| random_yaw_init: If True, give each non-anchor object a random fixed yaw about Z |
There was a problem hiding this comment.
Docstring is lil jorgan-y, to make it clean by
If True, randomly rotates non-anchor objects around the vertical (Z) axis at startup to add visual variety to the scene.
| return compose_poses(self, other) | ||
|
|
||
|
|
||
| def rotate_quat_by_yaw( |
There was a problem hiding this comment.
Can you add a unit test for this math func?
| base_xyzw: tuple[float, float, float, float], yaw_rad: float | ||
| ) -> tuple[float, float, float, float]: | ||
| """Rotate base_xyzw (xyzw) by an extra yaw about Z. Returns base unchanged when yaw is 0.""" | ||
| if yaw_rad == 0.0: |
There was a problem hiding this comment.
How about wrapping yaw_rad to [-pi, pi] to avoid unnecessary compute and potential errors caused by precision?
There was a problem hiding this comment.
I see, in _generate_initial_orientations() it outs [-pi, pi].
But for this helper func alone, it could be reused with user defined yaw_rad.
There was a problem hiding this comment.
Got it, thanks for clarifying! Since _generate_initial_orientations() outputs [-π, π] but this helper expects [0, 2π], the ranges are intentionally different. Makes sense to keep them as-is. 👍
| """Number of attempts made.""" | ||
|
|
||
| orientations: dict[ObjectBase, float] = field(default_factory=dict) | ||
| """Per-object fixed yaw (radians about Z) for this layout. Empty when unrotated.""" |
There was a problem hiding this comment.
It will be more informative if you clarify w.r.t to what frame.
fixed yaw
| bbox = obj.get_bounding_box().to(self._device) | ||
| if initial_orientations is not None: | ||
| yaws = [d.get(obj, 0.0) for d in initial_orientations] | ||
| if any(yaw != 0.0 for yaw in yaws): |
There was a problem hiding this comment.
again, I strongly encourage wrapping yaw before branching and applying the transformation
There was a problem hiding this comment.
Either you can ensure yaw is only provided thru _generate_initial_orientations(), so [-pi, pi] is enforced.
Or you can add a 'if not wrapper, wrap it' to be safe.
There was a problem hiding this comment.
Both approaches work. The second option (defensive wrap) is more robust if yaw could come from other sources in the future. Thanks for clarifying!
| env_id_tensor = torch.tensor([cur_env], device=env.device) | ||
| positions = layouts_per_env[local_idx].positions | ||
| result = layouts_per_env[local_idx] | ||
| positions = result.positions |
There was a problem hiding this comment.
Is it needed? You are only reusing it once next line.
results.orientations stays its own tho.
| via the same conservative enclosing AABB the solver uses, so validation matches solving.""" | ||
| bbox = obj.get_bounding_box() | ||
| yaw = (orientations or {}).get(obj, 0.0) | ||
| if yaw != 0.0: |
| else: | ||
| poses = [ | ||
| Pose(position_xyz=positions_per_env[env_idx][obj], rotation_xyzw=rotation_xyzw) | ||
| Pose( |
There was a problem hiding this comment.
Shall this func be updated to _apply_pose()
| self._opt_idx_tensor = None | ||
| self._optimizable_positions = None | ||
|
|
||
| # Precompute each object's local bbox on the solver device. Rotate to the conservative |
There was a problem hiding this comment.
# Precompute original bounding boxes on the GPU/device.
# If an object is rotated (yaw != 0) in any environment, calculate a larger,
# unrotated "enclosing" box to fit it, creating a separate box per environment.
# Otherwise, keep a single base box and let the loss function broadcast it to save memory.
^ Is my understanding correct?
| action="store_true", | ||
| default=False, | ||
| help=( | ||
| "Initialize non-anchor placeable objects with a random yaw about the up (Z) axis for scene variety. " |
There was a problem hiding this comment.
help=(
"Randomly rotate objects (except anchors) around the Z-axis for scene variety. "
"Collisions use a larger enclosing box; the solver won't optimize this rotation."
),
2 participants
Summary
Add random yaw initialization for placed objects
Detailed description
--random_yaw_initflag for per-object random z-rotation.