Add heterogeneous object placement support

[zhx06]

Summary

Add heterogeneous object placement support

Detailed description

• ObjectPlacer/RelationSolver: accept per-env bounding boxes for heterogeneous objects
• PooledObjectPlacer: per-env pools with single batched solver call
• NoCollisionLossStrategy: xy_only 2D overlap mode to avoid Z-gradient conflicts
• GR1 example: --mode heterogeneous with Robolab object sets
• Address the issue of bouncing objects by moving the office table to Background class

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Summary

This PR adds heterogeneous object placement support, allowing different environments to have different object variants (e.g., via RigidObjectSet). The implementation threads per-env bounding boxes through the solver, placer, and pooled placer, with a new xy_only mode for collision detection. The changes are substantial but well-structured, with comprehensive test coverage for the new functionality.

Architecture Impact

The change touches the core placement pipeline:

• RelationSolver.solve() now accepts env_bboxes parameter, propagated to _compute_total_loss() and _compute_no_overlap_loss()
• ObjectPlacer.place() gains env_bboxes parameter and new _place_heterogeneous() path
• PooledObjectPlacer splits into homogeneous (flat pool) and heterogeneous (per-env pools) modes
• NoCollisionLossStrategy adds xy_only parameter for 2D overlap detection
• _force_convex_hull() globally modifies physics collision approximations at env creation time

Callers affected: ArenaEnvBuilder._solve_relations() now passes num_envs to pool constructor. The solve_and_place_objects event now passes env_ids to the pool. Any external code creating PooledObjectPlacer with heterogeneous objects must now pass num_envs.

Implementation Verdict

Minor fixes needed — The implementation is solid overall, but there are a few correctness concerns around the _force_convex_hull timing and some edge cases in the heterogeneous pool logic.

Test Coverage

Excellent test coverage in test_heterogeneous_placement.py (488 lines). Tests cover:

• Per-env bbox expansion for homogeneous/heterogeneous objects
• Solver with env_bboxes parameter
• Placer heterogeneous detection and z-height correctness
• Mixed heterogeneous/homogeneous scenes
• PooledObjectPlacer heterogeneous mode: detection, sampling, refill, isolation
• Multi-set scenarios with different variant counts

Missing: No regression tests for _force_convex_hull() behavior. No tests verifying the xy_only collision mode produces correct gradients.

CI Status

No CI checks available yet — manual verification recommended before merge.

Findings

🔴 Critical: isaaclab_arena/environments/arena_env_builder.py:437 — _force_convex_hull may execute too late

The function modifies USD stage collision approximations after gym.make() returns. However, by this point physics initialization may have already cached the collision representations. The USD attribute change might not take effect for already-initialized physics bodies.

def make_registered_and_return_cfg(...):
    ...
    env = gym.make(name, cfg=cfg, render_mode=render_mode)
    _force_convex_hull(env)  # Physics may already be initialized

This needs verification that the stage modification propagates to the physics simulation. If not, this must be called before gym.make() or the scene must be explicitly reloaded.

---

🟡 Warning: isaaclab_arena/relations/pooled_object_placer.py:196-197 — Fallback accepts invalid layouts silently

When no env has valid layouts, the code accepts best-loss fallbacks with only a warning. In heterogeneous mode, this could lead to objects with incorrect geometry being placed:

if best is not None:
    print(f"Warning: env {env_id} had no valid layouts; ...")
    self._layout_pools[env_id].append(best)

Consider adding a configuration option to fail hard instead of accepting invalid layouts, especially since heterogeneous mode implies geometry-sensitive placement.

---

🟡 Warning: isaaclab_arena/relations/object_placer.py:236-240 — Potential off-by-one in candidate slicing

The slicing logic for env bboxes uses env_idx : env_idx + 1 which creates a tensor of shape (1, 3). This is correct, but the comment says "shape (1, 3)" while the docstring says "shape (num_envs, 3)" which is confusing:

env_child_bboxes = {
    obj: AxisAlignedBoundingBox(
        min_point=env_bboxes[obj].min_point[env_idx : env_idx + 1],  # (1, 3)
        max_point=env_bboxes[obj].max_point[env_idx : env_idx + 1],
    )
    for obj in objects
}

The docstring at line 209-211 should clarify that env_bboxes has shape (num_envs, 3) but individual slices are (1, 3).

---

🟡 Warning: isaaclab_arena/assets/object_set.py:75-79 — Assertion prevents valid use case

The assertion prevents random_choice=True with heterogeneous placement, but the error message references "placement pool assumes round-robin" which is implementation detail leaking into API:

assert not (self.heterogeneous_bbox and self.random_choice), (
    f"RigidObjectSet '{name}': random_choice=True is not supported with heterogeneous "
    "placement (len(objects) > 1). The placement pool assumes round-robin variant "
    "assignment (env_idx % num_variants) which conflicts with random spawning order."
)

Consider using raise ValueError(...) instead of assert for API validation (asserts can be disabled with -O).

---

🔵 Improvement: isaaclab_arena/relations/relation_solver.py:174-175 — Incomplete comment removal

The docstring was updated to explain anchor exclusion, but the code still iterates only over non_anchor_objects. The previous code had comments explaining anchor handling that are now removed. Add a brief inline comment for future maintainers:

non_anchor_objects = state.optimizable_objects
# Anchors excluded: On relation handles Z placement, 3D clearance would fight it

---

🔵 Improvement: isaaclab_arena/relations/pooled_object_placer.py:271-280 — Refill calculation may over-solve

The refill logic computes max_demand * self._num_envs which may generate many more layouts than needed:

if needs_refill:
    max_demand = max(demand_per_env.values())
    self._solve_and_store_heterogeneous(max(self._pool_size, max_demand * self._num_envs))

If only env 0 needs 1 more layout and there are 100 envs, this solves 100 new layouts. Consider a more targeted refill strategy that only refills depleted env pools.

---

🔵 Improvement: isaaclab_arena/tests/test_heterogeneous_placement.py:325-328 — Test uses bare except pattern

try:
    pool.sample_without_replacement(2, env_ids=None)
    assert False, "Should have raised AssertionError"
except AssertionError:
    pass

Use pytest.raises(AssertionError) instead for cleaner test code and better failure messages.

---

🔵 Improvement: isaaclab_arena/relations/relation_loss_strategies.py:345-347 — Method signature inconsistency

NoCollisionLossStrategy.compute_loss() now has xy_only parameter, but it's not a RelationLossStrategy subclass so it doesn't need to match the interface. However, the class docstring at line 326-334 doesn't mention xy_only. Update the docstring to document this parameter.

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Follow-up Review

The new commits address several previous concerns:

1. Previous concern about assert vs raise ValueError — Addressed: Line 75-79 in object_set.py now uses raise ValueError(...) instead of assert.

2. Previous concern about pytest.raises — Addressed: Line 325 in test_heterogeneous_placement.py now correctly uses pytest.raises(AssertionError).

Remaining Concerns

The following previous findings were not addressed and remain valid:

• 🔴 _force_convex_hull timing (arena_env_builder.py:437) — Still executes after gym.make(). Needs verification that physics re-reads the USD attribute, or must be moved earlier.

• 🟡 Fallback accepts invalid layouts silently (pooled_object_placer.py:196-197) — Still only prints a warning when accepting invalid layouts.

• 🔵 Refill over-solving (pooled_object_placer.py:271-280) — Still computes max_demand * self._num_envs which may generate excessive layouts.

• 🔵 NoCollisionLossStrategy docstring — Still doesn't document the xy_only parameter.

No new issues introduced by the latest commits. The core implementation remains solid.

[greptile-apps]

Greptile Summary

This PR extends the Isaac Lab Arena object placement system with heterogeneous support: RigidObjectSet now assigns a fixed per-env variant and exposes get_bounding_box_per_env, ObjectPlacer gains a dedicated _place_heterogeneous path that tiles per-env bboxes for one batched solver call, and PooledObjectPlacer is rearchitected around per-env layout pools with independent cursors. A GR1 example environment with --mode heterogeneous and four Robolab variant sets is added, along with a new test module covering the full heterogeneous placement stack.

• RigidObjectSet gains get_variant_indices / get_bounding_box_per_env: variant assignments are generated once and fixed for the object's lifetime.
• PooledObjectPlacer replaces the single shared list with per-env pools and cursors; layouts distributed via greedy shortest-first or env-matched batched solve.
• NoCollisionLossStrategy gains xy_only; _compute_no_overlap_loss now uses it unconditionally and removes the anchor-vs-non-anchor loss loop entirely.

Confidence Score: 3/5

Multiple solver and pool correctness issues remain open from previous review rounds and are still present in this revision; merging as-is risks silent bad placements and behavior regressions in homogeneous scenes.

The anchor-vs-non-anchor collision loss loop was deleted unconditionally affecting all scenes; xy_only=True is hard-coded globally; _store_env_matched_results can append invalid fallback entries to already-valid env pools on multi-batch refill; the env-indexed reset path drains a full round on every partial reset; force_convex_hull is never activated in the GR1 environment for heterogeneous mode.

isaaclab_arena/relations/relation_solver.py, isaaclab_arena/relations/pooled_object_placer.py, isaaclab_arena/relations/placement_events.py, isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py

Important Files Changed

Filename	Overview
isaaclab_arena/relations/relation_solver.py	Removes anchor-vs-non-anchor collision loss entirely and hard-codes xy_only=True for all non-anchor pairwise loss.
isaaclab_arena/relations/pooled_object_placer.py	Rearchitected around per-env pools; _store_env_matched_results fallback check is batch-scoped only and can append invalid layouts to already-valid pools.
isaaclab_arena/relations/placement_events.py	Reusable path correctly consumes only reset-env count; env-indexed path still consumes a full round on every partial reset.
isaaclab_arena/assets/object_set.py	Adds variant-index assignment with lazy initialization and correct spawn_cfg mutation; random_choice default preserved as False.
isaaclab_arena/relations/object_placer.py	Clean split into homogeneous and heterogeneous solve paths with correct bbox tiling and per-env validation.
isaaclab_arena_environments/gr1_table_multi_object_no_collision_environment.py	Adds --mode heterogeneous; force_convex_hull is never set so robolab assets never get convex-hull treatment.
isaaclab_arena/assets/background_library.py	Adds OfficeTableBackground as a kinematic Background asset to fix bouncing physics.
isaaclab_arena/relations/relation_loss_strategies.py	Adds xy_only parameter to NoCollisionLossStrategy.compute_loss; Z-overlap correctly gated.
isaaclab_arena/tests/test_heterogeneous_placement.py	New test module with comprehensive coverage of the heterogeneous stack.

_{Reviews (12): Last reviewed commit: "change comments and naming style" | Re-trigger Greptile}

[greptile-apps]

_force_convex_hull applied to every environment build, not only heterogeneous/robolab ones

_force_convex_hull(env) is called inside make_registered_and_return_cfg, which is invoked for every make_registered call regardless of whether the scene contains any robolab assets or uses heterogeneous mode at all. Any object across the entire USD stage (robot, table, YCB assets, …) whose MeshCollisionAPI approximation was set to convexDecomposition will be silently flipped to convexHull. This can break carefully authored physics for non-robolab objects. The fix should guard the call—either only invoke it when heterogeneous mode is active, or let the caller opt-in explicitly rather than running it unconditionally.

[greptile-apps]

Help text describes objects that are not the actual heterogeneous defaults

The --objects help string says the heterogeneous default is "cracker_box/sugar_box, tomato_soup_can/mustard_bottle", but those names do not appear anywhere in this module. The actual defaults are the entries of HETERO_VARIANT_SETS (bottles, cans, tools, boxes) plus HETERO_FIXED_OBJECTS (banana, lime). A user reading the help text would configure the wrong defaults or be confused when objects don't match.

[greptile-apps]

--mode heterogeneous --objects … silently falls back to homogeneous behaviour

When the user supplies --objects, each name is wrapped in a single-element RigidObjectSet. Because heterogeneous_bbox = len(objects) > 1, a single-element set always has heterogeneous_bbox=False. PooledObjectPlacer.__init__ then detects no heterogeneous objects and runs the flat homogeneous pool path, even though --mode heterogeneous was explicitly requested. There is no warning or error. Consider raising an error or emitting a warning when mode == "heterogeneous" but no object has heterogeneous_bbox=True after construction.

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Follow-up Review

No new commits since the previous review. The remaining concerns from the prior review are still valid:

1. _force_convex_hull timing — Still needs verification that post-gym.make() USD attribute changes are picked up by the physics engine.
2. Fallback accepts invalid layouts silently — Warning-only behavior unchanged.
3. Refill over-solving — Unchanged.
4. NoCollisionLossStrategy docstring — xy_only parameter still undocumented.

These are minor issues that don't block merge. The core heterogeneous placement implementation is correct and well-tested. Ship it with optional follow-up for the docstring and over-solving efficiency.

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Follow-up Review

The new commit (87ea7e6) adds the xy_only parameter documentation to NoCollisionLossStrategy.compute_loss() docstring — this addresses the previous docstring concern.

The other minor items from the previous review remain unchanged (fallback warning-only behavior, refill over-solving efficiency, _force_convex_hull timing verification) but these are non-blocking observations, not bugs.

Ship it — the implementation is correct, well-tested, and the documentation gap has been addressed.

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot

Summary

This PR adds comprehensive heterogeneous object placement support, enabling different environments to have different object variants (e.g., via RigidObjectSet). The implementation cleanly threads per-env bounding boxes through the placement pipeline with a new _place_heterogeneous() path and per-env pool management in PooledObjectPlacer.

Key Changes

• Assets Layer: New get_bounding_box_per_env() method in ObjectBase with specialized override in RigidObjectSet
• Solver: Extended RelationSolver.solve() to accept env_bboxes for per-env geometry
• Loss Strategies: Added xy_only mode to NoCollisionLossStrategy for 2D overlap (avoids Z-gradient conflicts with On constraint)
• ObjectPlacer: Split into _place_homogeneous() and _place_heterogeneous() paths with explicit strategy selection
• PooledObjectPlacer: Heterogeneous mode maintains per-env layout pools with automatic refill
• Environment: New force_convex_hull option and --mode heterogeneous CLI support

Test Coverage

Excellent test coverage in test_heterogeneous_placement.py (487 lines added). Tests comprehensively cover:

• Per-env bbox expansion (homogeneous/heterogeneous)
• Solver with env_bboxes parameter
• Placer heterogeneous detection and z-height correctness
• Mixed heterogeneous/homogeneous scene placement
• PooledObjectPlacer: detection, sampling with env_ids, refill logic, pool isolation
• Multi-set scenarios with different variant counts

CI Status

Pre-commit check is pending — recommend waiting for CI to pass before merge.

---

Findings

🟡 Medium: _force_convex_hull timing may be too late (arena_env_builder.py:410-425)

The function modifies USD stage collision approximations after gym.make() returns. By this point, PhysX may have already cached collision representations during scene initialization. Consider verifying that:

1. PhysX re-reads the approximationAttr on the next simulation step, or
2. Move this call earlier in the initialization sequence

env = gym.make(name, cfg=cfg, render_mode=render_mode)
if self.arena_env.force_convex_hull:
    _force_convex_hull(env)  # Is PhysX already initialized?

---

🟡 Medium: Silent fallback to invalid layouts (pooled_object_placer.py:193-201)

When all candidates fail validation for an env, the code accepts the best-loss layout with only a warning. In production scenarios with geometry-sensitive placement, this could lead to object interpenetration:

if best is not None:
    print(f"Warning: env {env_id} had no valid layouts; accepting best-loss fallback (loss={best.final_loss:.6f}).")
    self._layout_pools[env_id].append(best)

Consider adding a strict_validation option that raises instead of accepting invalid layouts.

---

🔵 Suggestion: Refill strategy may over-generate layouts (pooled_object_placer.py:271-280)

The refill logic computes max_demand * self._num_envs, which could generate many more layouts than needed. If only env 0 needs 1 more layout across 64 envs, this would solve 64 new layouts. A more targeted approach could refill only depleted env pools:

if needs_refill:
    max_demand = max(demand_per_env.values())
    self._solve_and_store_heterogeneous(max(self._pool_size, max_demand * self._num_envs))

---

🔵 Suggestion: Missing xy_only documentation (relation_loss_strategies.py:348-360)

The class-level docstring for NoCollisionLossStrategy doesn't mention the new xy_only parameter. The parameter is documented in compute_loss(), but the class docstring should summarize when to use 2D vs 3D overlap detection:

class NoCollisionLossStrategy:
    """Computes loss to prevent object overlap.
    ...
    When ``xy_only=True``, only XY overlap is used — suitable for objects on the
    same surface where Z overlap is expected and Z gradients would fight the On constraint.

---

🔵 Suggestion: Consider type annotation for _cached_variant_indices (object_set.py:134)

The _cached_variant_indices attribute is dynamically added without type annotation. Consider adding it as an optional instance attribute:

def __init__(self, ...):
    ...
    self._cached_variant_indices: list[int] | None = None

---

Overall Assessment

Ready for merge with minor suggestions. The implementation is well-structured with clean separation between homogeneous and heterogeneous paths. The test coverage is comprehensive. The _force_convex_hull timing concern is worth investigating but may work correctly in practice depending on IsaacSim's USD→PhysX synchronization behavior.

[greptile-apps]

Pool over-drained on every partial reset

sample_without_replacement(env.scene.env_origins.shape[0]) is called on every reset event, consuming one layout from every env's pool — even envs that are not resetting. In typical RL setups where individual environments time out at different steps, a single-env reset event now drains N pool entries instead of 1, making the effective drain rate N × (1/k) where k is the number of resetting envs per call. Because _solve_and_store is the only refill path and it runs the gradient solver, this causes solver invocations to become N/k times more frequent, which can be a severe performance regression at scale.

For the homogeneous path the layouts are fully interchangeable across envs, so there is no correctness requirement to advance all cursors in lockstep; only the resetting envs' positions are written anyway. Consider branching on _uses_env_specific_bboxes and calling sample_without_replacement(len(env_ids)) with sequential index lookup for the homogeneous case, while keeping the full-round advance only for the heterogeneous case where env-index correctness is required.

[greptile-apps]

Silent behavior change: random_choice default flipped from False to True

The old default was False (sequential, round-robin object spawning via MultiUsdFileCfg). Any existing RigidObjectSet constructed without an explicit random_choice argument will now randomly pick variants instead of using the deterministic sequential order. Additionally, existing callers that explicitly passed random_choice=False to get sequential ordering will now hit a ValueError in _generate_variant_indices if get_bounding_box_per_env is later called, because the new code treats random_choice=False + no variant_indices_by_env as an error. Consider keeping the default as False and updating existing callers that need random behaviour to pass random_choice=True explicitly, or document this as a breaking change.

Suggested change

	random_choice: bool = True,
	random_choice: bool = False,

[isaaclab-review-bot]

🤖 Isaac Lab Review Bot — Update Review

New Commits Since Last Review

5 new commits since the previous review (May 11):

• 5ec80e15cb10 — change objectset behavior
• 6cae73f7a93f — pre-commit fix
• 01b97f7dce94 — increase readability
• edbebd103a26 — address comments
• cf5599520cc0 — change comments and naming style

Changes Summary

The new commits primarily address code review feedback with:

1. Documentation improvements: Clarified random_choice vs deterministic member-order assignment behavior in get_variant_indices() docstring
2. Naming conventions: Renamed test functions for consistency (e.g., test_solver_accepts_env_bboxes → test_relation_solver_uses_env_bboxes)
3. Code cleanup: Removed unused self.objects = bboxes assignment in HeterogeneousDummyObject test helper
4. Comment refinements: Simplified pool storage documentation in PooledObjectPlacer

Assessment

✅ These changes look good. The updates improve code clarity without changing functionality. The test naming now follows a consistent pattern (test_<class>_<behavior>_<scenario>), making it easier to identify what each test covers.

Previous Findings Status

The suggestions from the May 11 review remain relevant:

• 🟡 _force_convex_hull timing concern — still worth verifying PhysX re-reads collision attrs
• 🟡 Silent fallback to invalid layouts — consider strict_validation option for production
• 🔵 Documentation suggestions — partially addressed by this commit's comment improvements

CI Status

Pre-commit passed ✅. Ready for final review and merge.

---

Reviewed at commit cf5599520cc013198f358e7aeba7b7a5b51dce56